This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new fdd60d0e [REFACTOR] Export public APIs through FFI-compatible
boundaries (#7)
fdd60d0e is described below
commit fdd60d0e46eb84f9db924a108a4867edb0941e2c
Author: Tianqi Chen <[email protected]>
AuthorDate: Tue Oct 20 21:00:42 2026 -0400
[REFACTOR] Export public APIs through FFI-compatible boundaries (#7)
Public C++ callers could not link several out-of-line compiler APIs
under hidden visibility. Export those APIs and inline the axis
predicates and sequential-pass pointer accessor. Layout operations use
exported nonvirtual reference-wrapper methods with explicit FFI type
dispatch; layout and document nodes remain data-only and need no C++
vtables.
Use FFI strings, owning callbacks, tuples and optionals across exported
signatures. Native symbolic-analysis maps and recovery callbacks retain
their existing interfaces, and registration-only hooks remain in their
owning implementation files. Header-defined native pass adapters
preserve mutable callables and move ownership. C++ consumers must
rebuild for the changed DLL signatures and use the layout
reference-wrapper methods.
---
include/tvm/codegen/codegen.h | 10 ++++----
include/tvm/ir/transform.h | 40 ++++++++++++++++++++++--------
include/tvm/runtime/logging.h | 10 +++++++-
include/tvm/script/ir_builder/frame.h | 23 +++++++++--------
include/tvm/script/ir_builder/ir.h | 39 +++++++++++++++--------------
include/tvm/script/printer/config.h | 4 +--
include/tvm/script/printer/printer.h | 2 +-
include/tvm/sym/analyzer.h | 15 +++++------
include/tvm/target/target.h | 10 ++++----
include/tvm/tirx/index_map.h | 27 ++++++++++++--------
include/tvm/tirx/layout.h | 21 +++++++++++-----
include/tvm/tirx/script/ir_builder/frame.h | 3 +--
include/tvm/tirx/script/ir_builder/ir.h | 16 ++++++------
include/tvm/tirx/stmt_functor.h | 4 +--
include/tvm/tirx/tile/context.h | 12 ++++-----
include/tvm/tirx/transform.h | 29 ++++++++++++++++++----
include/tvm/tirx/type.h | 2 +-
src/codegen/module_export.cc | 8 +++---
src/ir/transform.cc | 34 +++++++++++--------------
src/runtime/logging.cc | 10 --------
src/script/printer/printer.cc | 8 +++---
src/sym/const_int_bound.cc | 6 ++++-
src/target/target.cc | 6 ++---
src/tirx/ir/index_map.cc | 26 +++++++++----------
src/tirx/ir/layout/axis_registry.cc | 12 ---------
src/tirx/ir/stmt_functor.cc | 6 ++---
src/tirx/ir/transform.cc | 28 ++++++++++-----------
src/tirx/ir/type.cc | 2 +-
28 files changed, 226 insertions(+), 187 deletions(-)
diff --git a/include/tvm/codegen/codegen.h b/include/tvm/codegen/codegen.h
index a7f465c4..02f61ccf 100644
--- a/include/tvm/codegen/codegen.h
+++ b/include/tvm/codegen/codegen.h
@@ -54,8 +54,8 @@ TVM_TIRX_DLL ffi::Module Build(IRModule mod, Target target);
* \param c_symbol_prefix Optional symbol prefix of the blob symbol.
* \return cstr The C string representation of the file.
*/
-TVM_TIRX_DLL std::string PackImportsToC(const ffi::Module& m, bool system_lib,
- const std::string& c_symbol_prefix =
"");
+TVM_TIRX_DLL ffi::String PackImportsToC(const ffi::Module& m, bool system_lib,
+ const ffi::String& c_symbol_prefix =
"");
/*!
* \brief Pack imported device libraries into an object file using LLVM.
@@ -70,9 +70,9 @@ TVM_TIRX_DLL std::string PackImportsToC(const ffi::Module& m,
bool system_lib,
* \param c_symbol_prefix Optional blob symbol prefix, valid only for system
libraries.
*/
TVM_TIRX_DLL void PackImportsToObjectFile(const ffi::Module& mod, bool
system_lib,
- const std::string&
llvm_target_string,
- const std::string& file_name,
- const std::string& c_symbol_prefix =
"");
+ const ffi::String&
llvm_target_string,
+ const ffi::String& file_name,
+ const ffi::String& c_symbol_prefix =
"");
namespace attr {
inline constexpr ffi::reflection::AttrKey<bool> kParallelStridePattern{
diff --git a/include/tvm/ir/transform.h b/include/tvm/ir/transform.h
index 0d7763e6..c515e95a 100644
--- a/include/tvm/ir/transform.h
+++ b/include/tvm/ir/transform.h
@@ -32,6 +32,7 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/creator.h>
#include <tvm/ffi/reflection/registry.h>
+#include <tvm/ffi/rvalue_ref.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/module.h>
#include <tvm/ir/with_context.h>
@@ -196,8 +197,9 @@ class PassContext : public ffi::ObjectRef {
// The exit of a pass context scope.
TVM_TIRX_DLL void ExitWithScope();
// Register configuration key value type.
- TVM_TIRX_DLL static void RegisterConfigOption(const char* key, ffi::String
value_type_str,
-
std::function<ffi::Any(ffi::Any)> legalization);
+ TVM_TIRX_DLL static void RegisterConfigOption(
+ const char* key, ffi::String value_type_str,
+ ffi::TypedFunction<ffi::Any(ffi::Any)> legalization);
// Classes to get the Python `with` like syntax.
friend class Internal;
@@ -324,7 +326,7 @@ class Pass : public ffi::ObjectRef {
*
* \return The transformed module.
*/
- IRModule operator()(IRModule mod, const PassContext& pass_ctx) const;
+ TVM_TIRX_DLL IRModule operator()(IRModule mod, const PassContext& pass_ctx)
const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Pass, ffi::ObjectRef, PassObj);
@@ -374,7 +376,7 @@ class SequentialObj : public PassObj {
*
* \return Return the updated module.
*/
- IRModule operator()(IRModule mod, const PassContext& pass_ctx) const final;
+ TVM_TIRX_DLL IRModule operator()(IRModule mod, const PassContext& pass_ctx)
const final;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("transform.Sequential", SequentialObj,
PassObj);
};
@@ -400,21 +402,39 @@ class Sequential : public Pass {
explicit Sequential(ffi::UnsafeInit tag) : Pass(tag) {}
- const SequentialObj* operator->() const;
+ const SequentialObj* operator->() const { return static_cast<const
SequentialObj*>(get()); }
using ContainerType = SequentialObj;
};
-/*
- * \brief Create a module pass.
+/*!
+ * \brief Create a module pass from an owning FFI callback.
*
- * \param pass_func The packed function that contains the optimization.
+ * \param pass_func The callback that transforms the module.
* \param opt_level The optimization level of the module pass.
* \param name The name of the module pass.
+ * \return The created module pass.
+ */
+TVM_TIRX_DLL Pass
+CreateModulePass(ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>,
PassContext)> pass_func,
+ int opt_level, ffi::String name);
+
+/*!
+ * \brief Create a module pass from a native C++ callback.
*
+ * \param pass_func The callback that transforms the module.
+ * \param opt_level The optimization level of the module pass.
+ * \param name The name of the module pass.
* \return The created module pass.
*/
-TVM_TIRX_DLL Pass CreateModulePass(std::function<IRModule(IRModule,
PassContext)> pass_func,
- int opt_level, ffi::String name);
+inline Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)>
pass_func,
+ int opt_level, ffi::String name) {
+ ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>, PassContext)> callback
=
+ [pass_func = std::move(pass_func)](ffi::RValueRef<IRModule> mod,
+ PassContext ctx) -> IRModule {
+ return pass_func(*std::move(mod), std::move(ctx));
+ };
+ return CreateModulePass(std::move(callback), opt_level, std::move(name));
+}
/*!
* \brief A special trace pass that prints the header and IR to LOG(INFO).
diff --git a/include/tvm/runtime/logging.h b/include/tvm/runtime/logging.h
index 68718acc..ad61d761 100644
--- a/include/tvm/runtime/logging.h
+++ b/include/tvm/runtime/logging.h
@@ -304,7 +304,15 @@ class VLogContext {
}
}
- std::string str() const;
+ std::string str() const {
+ std::stringstream result;
+ for (const auto* entry : context_stack_) {
+ TVM_FFI_ICHECK_NOTNULL(entry);
+ result << entry->str();
+ result << ": ";
+ }
+ return result.str();
+ }
private:
std::vector<std::stringstream*> context_stack_;
diff --git a/include/tvm/script/ir_builder/frame.h
b/include/tvm/script/ir_builder/frame.h
index 7e7b2158..cb078bf1 100644
--- a/include/tvm/script/ir_builder/frame.h
+++ b/include/tvm/script/ir_builder/frame.h
@@ -68,7 +68,7 @@ class IRModuleFrameObj : public IRBuilderFrameObj {
IRBuilderFrameObj);
public:
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -160,9 +160,10 @@ class ForFrameObj : public StmtFrameObj {
public:
/*! \brief Apply source target names before entry, preserving variable
identity. */
- void SetNames(ffi::Optional<ffi::Variant<ffi::String,
ffi::Array<ffi::String>>> names);
+ TVM_TIRX_DLL void SetNames(
+ ffi::Optional<ffi::Variant<ffi::String, ffi::Array<ffi::String>>> names);
/*! \brief Construct the loop nest with this frame's stored source location.
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -213,7 +214,7 @@ class AssertFrameObj : public StmtFrameObj {
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -259,7 +260,7 @@ class RegionFrameObj : public StmtFrameObj {
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.RegionFrame",
RegionFrameObj, StmtFrameObj);
/*! \brief Construct the region and add it to the enclosing builder frame. */
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*! \brief Managed reference to RegionFrameObj. */
@@ -297,7 +298,7 @@ class WhileFrameObj : public StmtFrameObj {
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -346,7 +347,7 @@ class IfFrameObj : public StmtFrameObj {
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -380,12 +381,12 @@ class ThenFrameObj : public StmtFrameObj {
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
- void EnterWithScope() final;
+ TVM_TIRX_DLL void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
@@ -419,12 +420,12 @@ class ElseFrameObj : public StmtFrameObj {
* \brief The method called when entering RAII scope.
* \sa tvm::support::With
*/
- void EnterWithScope() final;
+ TVM_TIRX_DLL void EnterWithScope() final;
/*!
* \brief The method called when exiting RAII scope.
* \sa tvm::support::With
*/
- void ExitWithScope() final;
+ TVM_TIRX_DLL void ExitWithScope() final;
};
/*!
diff --git a/include/tvm/script/ir_builder/ir.h
b/include/tvm/script/ir_builder/ir.h
index dec738f7..924c82d8 100644
--- a/include/tvm/script/ir_builder/ir.h
+++ b/include/tvm/script/ir_builder/ir.h
@@ -70,10 +70,10 @@ TVM_TIRX_DLL void DefFunction(const ffi::String& func_name,
const BaseFunc& func
* \param annotations The optional annotations of the For statement.
* \return The ForFrame.
*/
-ForFrame For(ForKind kind, PrimExpr start, PrimExpr stop,
- ffi::Optional<PrimExpr> step = std::nullopt,
- ffi::Optional<PrimType> dtype = std::nullopt,
- ffi::Optional<ffi::Map<ffi::String, Any>> annotations =
std::nullopt);
+TVM_TIRX_DLL ForFrame For(ForKind kind, PrimExpr start, PrimExpr stop,
+ ffi::Optional<PrimExpr> step = std::nullopt,
+ ffi::Optional<PrimType> dtype = std::nullopt,
+ ffi::Optional<ffi::Map<ffi::String, Any>>
annotations = std::nullopt);
/*!
* \brief The grid For statement.
@@ -82,8 +82,9 @@ ForFrame For(ForKind kind, PrimExpr start, PrimExpr stop,
* each loop var takes the dtype of its own extent.
* \return The ForFrame.
*/
-ForFrame Grid(ffi::Array<ffi::Variant<PrimExpr, ffi::Tuple<PrimExpr,
PrimExpr>>> extents,
- ffi::Optional<PrimType> dtype = std::nullopt);
+TVM_TIRX_DLL ForFrame
+Grid(ffi::Array<ffi::Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>>
extents,
+ ffi::Optional<PrimType> dtype = std::nullopt);
/*!
* \brief The assertion statement.
@@ -92,53 +93,53 @@ ForFrame Grid(ffi::Array<ffi::Variant<PrimExpr,
ffi::Tuple<PrimExpr, PrimExpr>>>
* \param message_parts The error message parts (stored as separate fragments
in the IR).
* \return The AssertFrame.
*/
-AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
- ffi::Array<ffi::String> message_parts);
+TVM_TIRX_DLL AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
+ ffi::Array<ffi::String> message_parts);
/*!
* \brief Create a while loop.
* \param condition The termination condition of the loop.
* \return The result WhileFrame.
*/
-WhileFrame While(PrimExpr condition);
+TVM_TIRX_DLL WhileFrame While(PrimExpr condition);
/*!
* \brief Create a return statement.
* \param value The value to return.
* \return The same statement that was added to the parent frame.
*/
-tvm::Stmt Return(Expr value);
+TVM_TIRX_DLL tvm::Stmt Return(Expr value);
/*!
* \brief Create a break statement.
* \return The same statement that was added to the parent frame.
*/
-tvm::Stmt Break();
+TVM_TIRX_DLL tvm::Stmt Break();
/*!
* \brief Create a continue statement.
* \return The same statement that was added to the parent frame.
*/
-tvm::Stmt Continue();
+TVM_TIRX_DLL tvm::Stmt Continue();
/*!
* \brief Create an if statement.
* \param condition The condition of if statement.
* \return The result IfFrame.
*/
-IfFrame If(PrimExpr condition);
+TVM_TIRX_DLL IfFrame If(PrimExpr condition);
/*!
* \brief Create a then.
* \return The result ThenFrame.
*/
-ThenFrame Then();
+TVM_TIRX_DLL ThenFrame Then();
/*!
* \brief Create an else.
* \return The result ElseFrame.
*/
-ElseFrame Else();
+TVM_TIRX_DLL ElseFrame Else();
/*!
* \brief Construct a result-free region with lexical body parameters.
@@ -147,16 +148,16 @@ ElseFrame Else();
* \param body_params Existing body parameters, or omitted to generate them
from the operation.
* \param attrs Additional operation attributes.
*/
-RegionFrame Region(Op op, ffi::Array<Expr> args,
- ffi::Optional<ffi::Array<Var>> body_params = std::nullopt,
- DictAttrs attrs = DictAttrs());
+TVM_TIRX_DLL RegionFrame Region(Op op, ffi::Array<Expr> args,
+ ffi::Optional<ffi::Array<Var>> body_params =
std::nullopt,
+ DictAttrs attrs = DictAttrs());
/*!
* \brief Evaluate the input expression.
* \param value The input expression to evaluate.
* \return The same statement that was added to the parent frame.
*/
-tvm::Stmt Evaluate(Expr value);
+TVM_TIRX_DLL tvm::Stmt Evaluate(Expr value);
} // namespace ir
} // namespace ir_builder
diff --git a/include/tvm/script/printer/config.h
b/include/tvm/script/printer/config.h
index 440ca25b..7e80f1cc 100644
--- a/include/tvm/script/printer/config.h
+++ b/include/tvm/script/printer/config.h
@@ -108,7 +108,7 @@ class PrinterConfigObj : public ffi::Object {
.def_ro("extra_config", &PrinterConfigObj::extra_config);
}
- ffi::Array<ffi::String> GetBuiltinKeywords();
+ TVM_TIRX_DLL ffi::Array<ffi::String> GetBuiltinKeywords();
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.PrinterConfig", PrinterConfigObj,
ffi::Object);
@@ -129,7 +129,7 @@ class TVM_TIRX_DLL PrinterConfig : public ffi::ObjectRef {
* address string. Defined in src/script/printer/script_printer.cc so that
* <tvm/runtime/logging.h> is not pulled into this public header.
*/
-TVM_TIRX_DLL std::string RedirectedReprPrinterMethod(const ffi::ObjectRef&
obj);
+TVM_TIRX_DLL ffi::String RedirectedReprPrinterMethod(const ffi::ObjectRef&
obj);
} // namespace tvm
#endif // TVM_SCRIPT_PRINTER_CONFIG_H_
diff --git a/include/tvm/script/printer/printer.h
b/include/tvm/script/printer/printer.h
index b9f76d16..6aef9a47 100644
--- a/include/tvm/script/printer/printer.h
+++ b/include/tvm/script/printer/printer.h
@@ -38,7 +38,7 @@ namespace tvm {
* \param config Optional translation and rendering configuration.
* \return The rendered script or fallback representation.
*/
-TVM_TIRX_DLL std::string Script(const ffi::ObjectRef& node,
+TVM_TIRX_DLL ffi::String Script(const ffi::ObjectRef& node,
const ffi::Optional<PrinterConfig>& config =
std::nullopt);
namespace script {
diff --git a/include/tvm/sym/analyzer.h b/include/tvm/sym/analyzer.h
index ce90fcbb..b889061c 100644
--- a/include/tvm/sym/analyzer.h
+++ b/include/tvm/sym/analyzer.h
@@ -494,7 +494,7 @@ class IntSetAnalyzer {
*/
TVM_TIRX_DLL void Bind(const Var& var, const Range& new_range, bool
allow_override = false);
- std::function<void()> EnterConstraint(const PrimExpr& constraint);
+ TVM_TIRX_DLL std::function<void()> EnterConstraint(const PrimExpr&
constraint);
private:
friend class AnalyzerObj;
@@ -548,7 +548,8 @@ class Z3Prover {
* \param is_assume Whether the constraint comes from an assumption.
* \return an exit function that must be called to cleanup the constraint
can be nullptr.
*/
- std::function<void()> EnterConstraint(const PrimExpr& constraint, bool
is_assume = false);
+ TVM_TIRX_DLL std::function<void()> EnterConstraint(const PrimExpr&
constraint,
+ bool is_assume = false);
/*!
* \brief Get the SMTLIB2 representation of the current context.
@@ -556,28 +557,28 @@ class Z3Prover {
* \param expr The optional expression to check.
* \return The SMTLIB2 string.
*/
- ffi::String GetSMTLIB2(const ffi::Optional<PrimExpr> expr);
+ TVM_TIRX_DLL ffi::String GetSMTLIB2(const ffi::Optional<PrimExpr> expr);
/*!
* \brief Get statistics about Z3 prover.
*
* \return The statistics string.
*/
- ffi::String GetStats();
+ TVM_TIRX_DLL ffi::String GetStats();
/*!
* \brief Set timeout in milliseconds for Z3 prover.
*
* \param timeout_ms The timeout in milliseconds.
*/
- void SetTimeoutMs(unsigned timeout_ms);
+ TVM_TIRX_DLL void SetTimeoutMs(unsigned timeout_ms);
/*!
* \brief Set resource limitation for Z3 prover.
*
* \param rlimit the resource limitation.
*/
- void SetRLimit(unsigned rlimit);
+ TVM_TIRX_DLL void SetRLimit(unsigned rlimit);
/*!
* \brief Get the Z3 model for the given expression if satisfiable.
@@ -585,7 +586,7 @@ class Z3Prover {
* \param expr The expression to get the model for.
* \return The model as a string.
*/
- ffi::String GetModel(const PrimExpr& expr);
+ TVM_TIRX_DLL ffi::String GetModel(const PrimExpr& expr);
/*!
* \brief Count the number of integer values that satisfy the current
constraints.
diff --git a/include/tvm/target/target.h b/include/tvm/target/target.h
index d1a2fe99..5e9300f0 100644
--- a/include/tvm/target/target.h
+++ b/include/tvm/target/target.h
@@ -61,7 +61,7 @@ class TargetObj : public ffi::Object {
* \brief The JSON string representation of the target
* \return JSON string of the target configuration (e.g. {"kind": "llvm",
"mcpu": "cortex-a53"})
*/
- TVM_TIRX_DLL const std::string& str() const;
+ TVM_TIRX_DLL ffi::String str() const;
/*! \return Export target to JSON-like configuration */
TVM_TIRX_DLL ffi::Map<ffi::String, ffi::Any> ToConfig() const;
/*! \return The ffi::Optional<Target> typed target host of the TargetObj */
@@ -77,7 +77,7 @@ class TargetObj : public ffi::Object {
* \return True if the target's `TargetObj::keys` contains the
* specified key, False otherwise.
*/
- TVM_TIRX_DLL bool HasKey(const std::string& query_key) const;
+ TVM_TIRX_DLL bool HasKey(const ffi::String& query_key) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -139,7 +139,7 @@ class TargetObj : public ffi::Object {
private:
/*! \brief Internal string repr. */
- mutable std::string str_repr_;
+ mutable ffi::String str_repr_;
friend class TargetInternal;
};
@@ -181,10 +181,10 @@ class Target : public ffi::ObjectRef {
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Target, ffi::ObjectRef,
TargetObj);
- static Target WithHost(const Target& target, const ffi::Optional<Target>&
host);
+ TVM_TIRX_DLL static Target WithHost(const Target& target, const
ffi::Optional<Target>& host);
/*! \return The target with the host stripped out */
- Target WithoutHost() const;
+ TVM_TIRX_DLL Target WithoutHost() const;
private:
// enable with syntax.
diff --git a/include/tvm/tirx/index_map.h b/include/tvm/tirx/index_map.h
index b4e73824..115458f2 100644
--- a/include/tvm/tirx/index_map.h
+++ b/include/tvm/tirx/index_map.h
@@ -27,6 +27,7 @@
#define TVM_TIRX_INDEX_MAP_H_
#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
#include <tvm/ir/expr.h>
#include <tvm/runtime/tensor.h>
@@ -97,7 +98,7 @@ class IndexMapNode : public ffi::Object {
* \returns The indices in the output space. Contains one value for
* each expression in `final_indices`.
*/
- ffi::Array<PrimExpr> MapIndices(const ffi::Array<PrimExpr>& indices) const;
+ TVM_TIRX_DLL ffi::Array<PrimExpr> MapIndices(const ffi::Array<PrimExpr>&
indices) const;
/*!
* \brief Map indices to the output space using an existing analyzer.
*
@@ -122,7 +123,7 @@ class IndexMapNode : public ffi::Object {
* \returns The ranges in the output space. Contains one value for
* each expression in `final_indices`.
*/
- ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges) const;
+ TVM_TIRX_DLL ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges)
const;
/*! \brief Map a memory range to the output space using an existing analyzer.
*
* If contiguous memory locations in the input space are not
@@ -136,7 +137,8 @@ class IndexMapNode : public ffi::Object {
* \returns The ranges in the output space. Contains one value for
* each expression in `final_indices`.
*/
- ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges, const
sym::Analyzer& analyzer) const;
+ TVM_TIRX_DLL ffi::Array<Range> MapRanges(const ffi::Array<Range>& ranges,
+ const sym::Analyzer& analyzer)
const;
/*! \brief Map a tensor shape to the output space using a fresh analyzer.
*
@@ -145,7 +147,7 @@ class IndexMapNode : public ffi::Object {
* \returns The tensor shape in the output space. Contains one
* value for each expression in `final_indices`.
*/
- ffi::Array<PrimExpr> MapShape(const ffi::Array<PrimExpr>& shape) const;
+ TVM_TIRX_DLL ffi::Array<PrimExpr> MapShape(const ffi::Array<PrimExpr>&
shape) const;
/*! \brief Map a tensor shape to the output space using an existing analyzer.
*
* \param shape The tensor shape in the input space. Should contain
@@ -163,7 +165,7 @@ class IndexMapNode : public ffi::Object {
*
* \returns The transformed Tensor.
*/
- ffi::Tensor MapTensor(ffi::Tensor arr_src) const;
+ TVM_TIRX_DLL ffi::Tensor MapTensor(ffi::Tensor arr_src) const;
/*!
* \brief Convert to string representation in Python.
@@ -171,7 +173,8 @@ class IndexMapNode : public ffi::Object {
* \return The stringified lambda expression in Python.
*/
TVM_TIRX_DLL ffi::String ToPythonString(
- const std::function<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map = nullptr) const;
+ const ffi::TypedFunction<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map =
+ nullptr) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -218,7 +221,7 @@ class IndexMap : public ffi::ObjectRef {
* assumed to be correct and bijective, and is returned.
* \param initial_ranges The ranges of the input indices.
*/
- IndexMap Inverse(ffi::Array<Range> initial_ranges) const;
+ TVM_TIRX_DLL IndexMap Inverse(ffi::Array<Range> initial_ranges) const;
/*! \brief Generate the inverse mapping using an existing analyzer.
*
* The range of the input indices is required in order to ensure
@@ -241,8 +244,9 @@ class IndexMap : public ffi::ObjectRef {
* \param f_name_map The optional name map to rename the variables.
* \return The renamed index map.
*/
- TVM_TIRX_DLL IndexMap RenameVariables(
- const std::function<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map = nullptr) const;
+ TVM_TIRX_DLL IndexMap
+ RenameVariables(const ffi::TypedFunction<ffi::Optional<ffi::String>(const
Var& var)>& f_name_map =
+ nullptr) const;
/*! \brief Generate the inverse mapping using a fresh analyzer.
*
@@ -254,7 +258,8 @@ class IndexMap : public ffi::ObjectRef {
* \return The inverted index map, along with the predicate for
* which the inverse maps to a valid range.
*/
- std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(ffi::Array<Range>
initial_ranges) const;
+ TVM_TIRX_DLL ffi::Tuple<IndexMap, PrimExpr> NonSurjectiveInverse(
+ ffi::Array<Range> initial_ranges) const;
/*! \brief Generate the inverse mapping using an existing analyzer.
*
* Determine the inverse, where the output range may contain
@@ -267,7 +272,7 @@ class IndexMap : public ffi::ObjectRef {
* \return The inverted index map, along with the predicate for
* which the inverse maps to a valid range.
*/
- TVM_TIRX_DLL std::pair<IndexMap, PrimExpr> NonSurjectiveInverse(
+ TVM_TIRX_DLL ffi::Tuple<IndexMap, PrimExpr> NonSurjectiveInverse(
ffi::Array<Range> initial_ranges, const sym::Analyzer& analyzer) const;
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(IndexMap, ffi::ObjectRef,
IndexMapNode);
diff --git a/include/tvm/tirx/layout.h b/include/tvm/tirx/layout.h
index 7c2527ce..1085dd90 100644
--- a/include/tvm/tirx/layout.h
+++ b/include/tvm/tirx/layout.h
@@ -85,22 +85,31 @@ using FAxisFuser =
ffi::TypedFunction<ffi::Optional<Iter>(Target, ThreadScope, T
class AxisNode : public ffi::EnumObj {
public:
/*! \brief Check if the axis is a thread axis. */
- bool IsThreadAxis() const;
+ bool IsThreadAxis() const {
+ static ffi::reflection::TypeAttrColumn
state_column(ffi::type_attr::kEnumState);
+ static ffi::EnumState state =
+
ffi::Any(state_column[_GetOrAllocRuntimeTypeIndex()]).as_or_throw<ffi::EnumState>();
+ auto column = state->attrs.Get("thread");
+ auto thread = column ? column->Get(ffi::GetRef<ffi::ObjectRef>(this)) :
std::nullopt;
+ TVM_FFI_ICHECK(thread.has_value())
+ << "Axis '" << _str_index << "' has no thread classification";
+ return thread->cast<bool>();
+ }
/*! \brief Check if the axis is a memory axis. */
- bool IsMemoryAxis() const;
+ bool IsMemoryAxis() const { return !IsThreadAxis(); }
/*! \brief The backend owning this axis, if any. */
- ffi::Optional<ffi::String> GetBackend() const;
+ TVM_TIRX_DLL ffi::Optional<ffi::String> GetBackend() const;
/*! \brief Get the scope of the (thread) axis. */
- ffi::Optional<ThreadScope> GetScope() const;
+ TVM_TIRX_DLL ffi::Optional<ThreadScope> GetScope() const;
/*! \brief Get the subscope of the (thread) axis. */
- ffi::Optional<ThreadScope> GetSubscope() const;
+ TVM_TIRX_DLL ffi::Optional<ThreadScope> GetSubscope() const;
/*! \brief Get the fuser of the (thread) axis. */
- ffi::Optional<FAxisFuser> GetFuser() const;
+ TVM_TIRX_DLL ffi::Optional<FAxisFuser> GetFuser() const;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Axis", AxisNode, ffi::EnumObj);
};
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index bbef2c86..01427e4e 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -29,7 +29,6 @@
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/tile/thread_scope.h>
-#include <functional>
#include <utility>
namespace tvm {
@@ -86,7 +85,7 @@ class TVM_TIRX_DLL FunctionFrameObj : public ir::StmtFrameObj
{
void ExitWithScope() final;
/*! \brief Register validation for an extension-owned function attribute. */
- using AttrValidator = std::function<void(const FunctionFrameObj*, const
ffi::Any&)>;
+ using AttrValidator = ffi::TypedFunction<void(const FunctionFrameObj*, const
ffi::Any&)>;
static void RegisterAttrValidator(ffi::String key, AttrValidator validator);
void ValidateAttrs() const;
diff --git a/include/tvm/tirx/script/ir_builder/ir.h
b/include/tvm/tirx/script/ir_builder/ir.h
index d7000095..c16c189f 100644
--- a/include/tvm/tirx/script/ir_builder/ir.h
+++ b/include/tvm/tirx/script/ir_builder/ir.h
@@ -49,8 +49,8 @@ using tvm::tirx::ThreadScope;
* \brief The primitive function statement.
* \return The FunctionFrame.
*/
-FunctionFrame Function(bool is_private = false, bool persistent = false,
- bool is_declaration = false);
+TVM_TIRX_DLL FunctionFrame Function(bool is_private = false, bool persistent =
false,
+ bool is_declaration = false);
/*!
* \brief The Function variable arguments adding function.
@@ -58,7 +58,7 @@ FunctionFrame Function(bool is_private = false, bool
persistent = false,
* \param var The variable argument.
* \return The variable.
*/
-Var Arg(ffi::String name, Var var);
+TVM_TIRX_DLL Var Arg(ffi::String name, Var var);
/*!
* \brief The Function tensor arguments adding function.
@@ -66,26 +66,26 @@ Var Arg(ffi::String name, Var var);
* \param tensor The tensor argument.
* \return The tensor.
*/
-TensorVar Arg(ffi::String name, TensorVar tensor);
+TVM_TIRX_DLL TensorVar Arg(ffi::String name, TensorVar tensor);
/*!
* \brief The Function naming statement.
* \param name The name of the Function.
*/
-void FuncName(ffi::String name);
+TVM_TIRX_DLL void FuncName(ffi::String name);
/*!
* \brief The Function annotation statement.
* \param attrs The annotations of the Function.
*/
-void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
+TVM_TIRX_DLL void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
/*!
* \brief The Function return type statement.
* \param ret_type The return type of the Function.
* \return The return type.
*/
-Type FuncRet(Type ret_type);
+TVM_TIRX_DLL Type FuncRet(Type ret_type);
/*!
* \brief Store data in a tensor destination.
@@ -94,7 +94,7 @@ Type FuncRet(Type ret_type);
* \param value The value to be stored.
* \return The same statement that was added to the parent frame.
*/
-tvm::Stmt TensorStore(Expr dest, ffi::Array<PrimExpr> indices, PrimExpr value);
+TVM_TIRX_DLL tvm::Stmt TensorStore(Expr dest, ffi::Array<PrimExpr> indices,
PrimExpr value);
/*!
* \brief Create a TIR var that represents a pointer
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index bf7d59c2..d532e8fc 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -294,7 +294,7 @@ class TVM_TIRX_DLL StmtExprMutator : public
tvm::ExprMutator {
* \return The result.
*/
TVM_TIRX_DLL Stmt SubstituteWithDataTypeLegalization(
- Stmt stmt, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
+ Stmt stmt, ffi::TypedFunction<ffi::Optional<PrimExpr>(const Var&)> vmap);
/*!
* \brief Substitute the var specified by vmap and legalize data types after
substitution.
@@ -306,7 +306,7 @@ TVM_TIRX_DLL Stmt SubstituteWithDataTypeLegalization(
* \return The result.
*/
TVM_TIRX_DLL PrimExpr SubstituteWithDataTypeLegalization(
- PrimExpr expr, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap);
+ PrimExpr expr, ffi::TypedFunction<ffi::Optional<PrimExpr>(const Var&)>
vmap);
/*!
* \brief Check if the statement contains the specified node type.
diff --git a/include/tvm/tirx/tile/context.h b/include/tvm/tirx/tile/context.h
index dc44797f..97bff8b4 100644
--- a/include/tvm/tirx/tile/context.h
+++ b/include/tvm/tirx/tile/context.h
@@ -104,20 +104,20 @@ class DispatchContextObj : public ffi::Object {
}
/*! \brief Add a tensor to be allocated in the kernel. */
- void AddAllocTensor(TensorVar tensor, ffi::Array<PrimExpr> allocated_addr =
{},
- ffi::Map<ffi::String, ffi::Any> annotations = {});
+ TVM_TIRX_DLL void AddAllocTensor(TensorVar tensor, ffi::Array<PrimExpr>
allocated_addr = {},
+ ffi::Map<ffi::String, ffi::Any> annotations
= {});
/*! \brief Add an initialization statement to be inserted. */
- void AddInitStmt(Stmt stmt, bool host = false);
+ TVM_TIRX_DLL void AddInitStmt(Stmt stmt, bool host = false);
/*! \brief Add a statement to be inserted after a tensor's definition. */
- void AddPostTensorDefStmt(TensorVar tensor, Stmt stmt);
+ TVM_TIRX_DLL void AddPostTensorDefStmt(TensorVar tensor, Stmt stmt);
/*! \brief Set a value in the shared state cache. */
- void SharedStateSet(ffi::String key, ffi::ObjectRef value);
+ TVM_TIRX_DLL void SharedStateSet(ffi::String key, ffi::ObjectRef value);
/*! \brief Get a value from the shared state cache. */
- ffi::Optional<ffi::ObjectRef> SharedStateGet(ffi::String key);
+ TVM_TIRX_DLL ffi::Optional<ffi::ObjectRef> SharedStateGet(ffi::String key);
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.DispatchContext",
DispatchContextObj, ffi::Object);
};
diff --git a/include/tvm/tirx/transform.h b/include/tvm/tirx/transform.h
index 1fe81fee..cca75c53 100644
--- a/include/tvm/tirx/transform.h
+++ b/include/tvm/tirx/transform.h
@@ -49,19 +49,38 @@ using tvm::transform::Sequential;
/*! \brief Wrap a function body in a device scope region. */
TVM_TIRX_DLL Pass DecorateDeviceScope();
-/*
- * \brief Create a function pass that optimizes Functions.
+/*!
+ * \brief Create a function pass from an owning FFI callback.
*
- * \param pass_func The packed function that contains the optimization.
+ * \param pass_func The callback that transforms each function, returning
std::nullopt to remove it.
* \param opt_level The optimization level of the function pass.
* \param name The name of the function pass.
- *
* \return The created function pass.
*/
TVM_TIRX_DLL Pass CreateFunctionPass(
- std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func,
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ pass_func,
int opt_level, ffi::String name);
+/*!
+ * \brief Create a function pass from a native C++ callback.
+ *
+ * \param pass_func The callback that transforms each function, returning
std::nullopt to remove it.
+ * \param opt_level The optimization level of the function pass.
+ * \param name The name of the function pass.
+ * \return The created function pass.
+ */
+inline Pass CreateFunctionPass(
+ std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func,
+ int opt_level, ffi::String name) {
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ callback = [pass_func = std::move(pass_func)](ffi::RValueRef<Function>
func, IRModule mod,
+ PassContext ctx) ->
ffi::Optional<Function> {
+ return pass_func(*std::move(func), std::move(mod), std::move(ctx));
+ };
+ return CreateFunctionPass(std::move(callback), opt_level, std::move(name));
+}
+
/*!
* \brief Lower vectorization loops.
*
diff --git a/include/tvm/tirx/type.h b/include/tvm/tirx/type.h
index 1cd7a1e3..ddcc10f9 100644
--- a/include/tvm/tirx/type.h
+++ b/include/tvm/tirx/type.h
@@ -111,7 +111,7 @@ class TensorTypeNode : public TypeNode {
TVM_TIRX_DLL bool IsScalar(bool alloc_or_decl = true) const;
/*! \brief Return the constant element count, or nullopt for symbolic
extents or overflow. */
- TVM_TIRX_DLL std::optional<int64_t> ConstantAllocationSize() const;
+ TVM_TIRX_DLL ffi::Optional<int64_t> ConstantAllocationSize() const;
/*! \brief Determine the offset in the tensor of the given index.
*
diff --git a/src/codegen/module_export.cc b/src/codegen/module_export.cc
index 098afbf4..a500cd22 100644
--- a/src/codegen/module_export.cc
+++ b/src/codegen/module_export.cc
@@ -269,8 +269,8 @@ std::string PackImportsToBytes(const ffi::Module& mod) {
} // namespace
-std::string PackImportsToC(const ffi::Module& mod, bool system_lib,
- const std::string& c_symbol_prefix) {
+ffi::String PackImportsToC(const ffi::Module& mod, bool system_lib,
+ const ffi::String& c_symbol_prefix) {
if (c_symbol_prefix.length() != 0) {
TVM_FFI_ICHECK(system_lib)
<< "c_symbol_prefix advanced option should be used in conjuction with
system-lib";
@@ -314,8 +314,8 @@ std::string PackImportsToC(const ffi::Module& mod, bool
system_lib,
}
void PackImportsToObjectFile(const ffi::Module& mod, bool system_lib,
- const std::string& llvm_target_string, const
std::string& file_name,
- const std::string& c_symbol_prefix) {
+ const ffi::String& llvm_target_string, const
ffi::String& file_name,
+ const ffi::String& c_symbol_prefix) {
if (c_symbol_prefix.length() != 0) {
TVM_FFI_ICHECK(system_lib)
<< "c_symbol_prefix advanced option should be used in conjuction with
system-lib";
diff --git a/src/ir/transform.cc b/src/ir/transform.cc
index 96773b29..80a7cda7 100644
--- a/src/ir/transform.cc
+++ b/src/ir/transform.cc
@@ -107,7 +107,7 @@ bool PassContext::PassEnabled(const PassInfo& info) const {
class PassConfigManager {
public:
void Register(std::string key, ffi::String value_type_str,
- std::function<ffi::Any(ffi::Any)> legalization) {
+ ffi::TypedFunction<ffi::Any(ffi::Any)> legalization) {
TVM_FFI_ICHECK_EQ(key2vtype_.count(key), 0U);
ValueTypeInfo info;
info.type_str = value_type_str;
@@ -135,7 +135,7 @@ class PassConfigManager {
TVM_FFI_CHECK(value != nullptr, AttributeError) << key << " is None";
- TVM_FFI_CHECK(info.legalization, AttributeError)
+ TVM_FFI_CHECK(info.legalization != nullptr, AttributeError)
<< "Config option \'" << key << "\' was defined without a
legalization function.";
auto legalized = info.legalization(value);
if (!legalized.same_as(value)) {
@@ -165,14 +165,14 @@ class PassConfigManager {
private:
struct ValueTypeInfo {
std::string type_str;
- std::function<ffi::Any(ffi::Any)> legalization;
+ ffi::TypedFunction<ffi::Any(ffi::Any)> legalization;
};
std::unordered_map<std::string, ValueTypeInfo> key2vtype_;
};
void PassContext::RegisterConfigOption(const char* key, ffi::String
value_type_str,
- std::function<ffi::Any(ffi::Any)>
legalization) {
+ ffi::TypedFunction<ffi::Any(ffi::Any)>
legalization) {
PassConfigManager::Global()->Register(key, value_type_str, legalization);
}
@@ -430,7 +430,7 @@ class ModulePassObj : public PassObj {
* implement the algorithm in the `pass_func` and let it run on a module. It
* will then remove the dead code including the unused functions in the
module.
*/
- std::function<IRModule(IRModule, PassContext)> pass_func;
+ ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>, PassContext)>
pass_func;
explicit ModulePassObj(PassInfo pass_info) : pass_info(std::move(pass_info))
{}
explicit ModulePassObj(ffi::UnsafeInit tag) : pass_info(tag) {}
@@ -459,7 +459,8 @@ class ModulePassObj : public PassObj {
class ModulePass : public Pass {
public:
- ModulePass(std::function<IRModule(IRModule, PassContext)> pass_func,
PassInfo pass_info);
+ ModulePass(ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>,
PassContext)> pass_func,
+ PassInfo pass_info);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ModulePass, Pass,
ModulePassObj);
};
@@ -471,7 +472,9 @@ PassInfo::PassInfo(int opt_level, ffi::String name) {
data_ = std::move(pass_info);
}
-ModulePass::ModulePass(std::function<IRModule(IRModule, PassContext)>
pass_func, PassInfo pass_info)
+ModulePass::ModulePass(
+ ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>, PassContext)>
pass_func,
+ PassInfo pass_info)
: Pass(ffi::UnsafeInit{}) {
auto n = ffi::make_object<ModulePassObj>(std::move(pass_info));
n->pass_func = std::move(pass_func);
@@ -491,7 +494,7 @@ IRModule ModulePassObj::operator()(IRModule mod, const
PassContext& pass_ctx) co
// TVMScript-rendered location.
IRModule input_mod = mod;
try {
- mod = pass_func(std::move(mod), pass_ctx);
+ mod = pass_func(ffi::RValueRef<IRModule>(std::move(mod)), pass_ctx);
} catch (ffi::Error& err) {
throw EnrichPassErrorWithContext(err, input_mod, pass_info->name);
}
@@ -513,10 +516,6 @@ Sequential::Sequential(tvm::ffi::Array<Pass> passes,
ffi::String name) : Pass(ff
data_ = std::move(n);
}
-const SequentialObj* Sequential::operator->() const {
- return static_cast<const SequentialObj*>(get());
-}
-
IRModule SequentialObj::operator()(IRModule mod, const PassContext& pass_ctx)
const {
for (const Pass& pass : passes) {
VLOG(0) << "Running pass " << pass->Info()->name;
@@ -532,8 +531,8 @@ IRModule SequentialObj::operator()(IRModule mod, const
PassContext& pass_ctx) co
return mod;
}
-Pass CreateModulePass(std::function<IRModule(IRModule, PassContext)>
pass_func, int opt_level,
- ffi::String name) {
+Pass CreateModulePass(ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>,
PassContext)> pass_func,
+ int opt_level, ffi::String name) {
PassInfo pass_info = PassInfo(opt_level, name);
return ModulePass(std::move(pass_func), pass_info);
}
@@ -563,12 +562,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("transform.MakeModulePass",
[](ffi::TypedFunction<IRModule(ffi::RValueRef<IRModule>,
PassContext)> pass_func,
- PassInfo pass_info) {
- auto wrapped_pass_func = [pass_func](IRModule mod, PassContext
ctx) {
- return pass_func(ffi::RValueRef<IRModule>(std::move(mod)), ctx);
- };
- return ModulePass(wrapped_pass_func, pass_info);
- })
+ PassInfo pass_info) { return ModulePass(std::move(pass_func),
pass_info); })
.def("transform.RunPass",
[](Pass pass, ffi::RValueRef<IRModule> mod) { return
pass(*std::move(mod)); });
}
diff --git a/src/runtime/logging.cc b/src/runtime/logging.cc
index 388902ee..847d9b12 100644
--- a/src/runtime/logging.cc
+++ b/src/runtime/logging.cc
@@ -136,16 +136,6 @@ LogFatal::Entry& LogFatal::GetEntry() {
return result;
}
-std::string VLogContext::str() const {
- std::stringstream result;
- for (const auto* entry : context_stack_) {
- TVM_FFI_ICHECK_NOTNULL(entry);
- result << entry->str();
- result << ": ";
- }
- return result.str();
-}
-
} // namespace detail
} // namespace runtime
} // namespace tvm
diff --git a/src/script/printer/printer.cc b/src/script/printer/printer.cc
index 58d8f52a..833b35ea 100644
--- a/src/script/printer/printer.cc
+++ b/src/script/printer/printer.cc
@@ -457,7 +457,7 @@ std::string RenderFallbackWithInvisiblePathInfo(const
ffi::String& script,
} // namespace
-std::string Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig>& cfg) {
+ffi::String Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig>& cfg) {
PrinterConfig config = cfg.value_or(PrinterConfig());
static ffi::reflection::TypeAttrColumn translate(
tvm::script::printer::type_attr::kDocTranslate.name);
@@ -467,15 +467,15 @@ std::string Script(const ffi::ObjectRef& node, const
ffi::Optional<PrinterConfig
!node.as<TypeNode>())) {
return RenderFallbackWithInvisiblePathInfo(ffi::ReprPrint(ffi::Any(node)),
config);
}
- return std::string(script::printer::RenderScript(node, config));
+ return script::printer::RenderScript(node, config);
}
-std::string RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
+ffi::String RedirectedReprPrinterMethod(const ffi::ObjectRef& obj) {
try {
PrinterConfig config;
config->extra_config.Set("ir.comment_imports", true);
// Call translation directly so an unsupported type cannot recurse through
ffi repr.
- return std::string(script::printer::RenderScript(obj, config));
+ return script::printer::RenderScript(obj, config);
} catch (const tvm::ffi::Error& e) {
LOG(WARNING) << "TVMScript printer falls back to the basic address printer
with the error:\n"
<< e.what();
diff --git a/src/sym/const_int_bound.cc b/src/sym/const_int_bound.cc
index ff335127..1ce9aa17 100644
--- a/src/sym/const_int_bound.cc
+++ b/src/sym/const_int_bound.cc
@@ -878,9 +878,13 @@ ConstIntBound ConstIntBoundAnalyzer::operator()(const
PrimExpr& expr) const {
}
ConstIntBound ConstIntBoundAnalyzer::operator()(const PrimExpr& expr,
BoundMapType* bound) {
+ struct RestoreMemo {
+ Impl* impl;
+ BoundMapType* previous;
+ ~RestoreMemo() { impl->bound_ = previous; }
+ } restore{impl_.get(), impl_->bound_};
impl_->bound_ = bound;
Entry ret = impl_->Dispatch(expr);
- impl_->bound_ = nullptr;
return ConstIntBound(ret.min_value, ret.max_value);
}
diff --git a/src/target/target.cc b/src/target/target.cc
index e0fcabfa..56dee045 100644
--- a/src/target/target.cc
+++ b/src/target/target.cc
@@ -93,9 +93,9 @@ static TargetKind GetTargetKind(const ffi::String& name) {
return kind.value();
}
-const std::string& TargetObj::str() const {
+ffi::String TargetObj::str() const {
if (str_repr_.empty()) {
- str_repr_ = std::string(ffi::json::Stringify(ToConfig()));
+ str_repr_ = ffi::json::Stringify(ToConfig());
}
return str_repr_;
}
@@ -176,7 +176,7 @@ int TargetObj::GetTargetDeviceType() const {
return kind->default_device_type;
}
-bool TargetObj::HasKey(const std::string& query_key) const {
+bool TargetObj::HasKey(const ffi::String& query_key) const {
return std::any_of(keys.begin(), keys.end(),
[&query_key](const auto& key) { return key == query_key;
});
}
diff --git a/src/tirx/ir/index_map.cc b/src/tirx/ir/index_map.cc
index fcbeac6e..b28003ff 100644
--- a/src/tirx/ir/index_map.cc
+++ b/src/tirx/ir/index_map.cc
@@ -67,10 +67,10 @@ IndexMap IndexMap::FromFunc(int ndim,
return IndexMap(initial_indices, func(callback_indices),
std::move(inverse_index_map));
}
-std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self,
- const ffi::Array<Range>&
initial_ranges,
- sym::IterMapLevel
check_level,
- sym::AnalyzerObj* analyzer) {
+ffi::Tuple<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self,
+ const ffi::Array<Range>&
initial_ranges,
+ sym::IterMapLevel
check_level,
+ sym::AnalyzerObj* analyzer)
{
TVM_FFI_ICHECK(analyzer != nullptr);
sym::Analyzer analyzer_ref = ffi::GetRef<sym::Analyzer>(analyzer);
if (self->inverse_index_map.has_value()) {
@@ -78,7 +78,7 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const
IndexMap& self,
// case, the user-defined inverse is assumed to be correct and
// bijective.
PrimExpr padding_predicate = IntImm::Bool(false);
- return {self->inverse_index_map.value().as_or_throw<IndexMap>(),
padding_predicate};
+ return ffi::Tuple(self->inverse_index_map.value().as_or_throw<IndexMap>(),
padding_predicate);
}
// Dummy variables to represent the inverse's inputs.
@@ -158,17 +158,17 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const
IndexMap& self,
padding_predicate = output_var_analyzer->Simplify(padding_predicate, 10);
}
- return {IndexMap(output_vars, inverse_exprs), padding_predicate};
+ return ffi::Tuple(IndexMap(output_vars, inverse_exprs), padding_predicate);
}
-std::pair<IndexMap, PrimExpr> IndexMap::NonSurjectiveInverse(
+ffi::Tuple<IndexMap, PrimExpr> IndexMap::NonSurjectiveInverse(
ffi::Array<Range> initial_ranges) const {
sym::Analyzer analyzer;
return NonSurjectiveInverse(initial_ranges, analyzer);
}
-std::pair<IndexMap, PrimExpr> IndexMap::NonSurjectiveInverse(ffi::Array<Range>
initial_ranges,
- const
sym::Analyzer& analyzer) const {
+ffi::Tuple<IndexMap, PrimExpr>
IndexMap::NonSurjectiveInverse(ffi::Array<Range> initial_ranges,
+ const
sym::Analyzer& analyzer) const {
return IndexMapInverseImpl(*this, initial_ranges,
sym::IterMapLevel::NoCheck, analyzer.get());
}
@@ -205,7 +205,7 @@ ffi::Array<PrimExpr> IndexMapNode::MapIndices(const
ffi::Array<PrimExpr>& indice
ffi::Array<PrimExpr> output = final_indices.Map([&](PrimExpr index) {
PrimExpr result = SubstituteWithDataTypeLegalization(
- std::move(index), [&](const Var& var) { return vmap.Get(var); });
+ std::move(index), [&](const Var& var) -> ffi::Optional<PrimExpr> {
return vmap.Get(var); });
return analyzer_ptr->Simplify(result);
});
return output;
@@ -363,7 +363,7 @@ ffi::Tensor IndexMapNode::MapTensor(ffi::Tensor arr_src)
const {
}
IndexMap IndexMap::RenameVariables(
- const std::function<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map) const {
+ const ffi::TypedFunction<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map) const {
std::unordered_set<std::string> used_names;
ffi::Map<Var, Var> var_remap;
UniqueNameSupply name_supply;
@@ -461,7 +461,7 @@ std::string IndexMap2PythonLambdaExpr(const
ffi::Array<PrimVar>& initial_indices
}
ffi::String IndexMapNode::ToPythonString(
- const std::function<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map) const {
+ const ffi::TypedFunction<ffi::Optional<ffi::String>(const Var& var)>&
f_name_map) const {
auto index_map = ffi::GetRef<IndexMap>(this).RenameVariables(f_name_map);
std::string lambda_expr =
IndexMap2PythonLambdaExpr(index_map->initial_indices,
index_map->final_indices);
@@ -516,7 +516,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::Optional<sym::Analyzer> opt_analyzer) {
sym::Analyzer analyzer = opt_analyzer.has_value() ?
opt_analyzer.value() : sym::Analyzer();
auto result = forward.NonSurjectiveInverse(initial_ranges, analyzer);
- return ffi::Array<ffi::ObjectRef>{result.first, result.second};
+ return ffi::Array<ffi::ObjectRef>{result.get<0>(), result.get<1>()};
});
}
diff --git a/src/tirx/ir/layout/axis_registry.cc
b/src/tirx/ir/layout/axis_registry.cc
index df5bae9c..c614b96c 100644
--- a/src/tirx/ir/layout/axis_registry.cc
+++ b/src/tirx/ir/layout/axis_registry.cc
@@ -52,18 +52,6 @@ ffi::Optional<T> AxisAttr(const AxisNode* axis, const char*
key) {
return value.value().cast<T>();
}
-bool AxisNode::IsThreadAxis() const {
- auto thread = AxisAttr<bool>(this, "thread");
- TVM_FFI_ICHECK(thread.has_value()) << "Axis '" << _str_index << "' has no
thread classification";
- return thread.value();
-}
-
-bool AxisNode::IsMemoryAxis() const {
- auto thread = AxisAttr<bool>(this, "thread");
- TVM_FFI_ICHECK(thread.has_value()) << "Axis '" << _str_index << "' has no
thread classification";
- return !thread.value();
-}
-
ffi::Optional<ffi::String> AxisNode::GetBackend() const {
return AxisAttr<ffi::String>(this, "backend");
}
diff --git a/src/tirx/ir/stmt_functor.cc b/src/tirx/ir/stmt_functor.cc
index 8034176c..a10adb8c 100644
--- a/src/tirx/ir/stmt_functor.cc
+++ b/src/tirx/ir/stmt_functor.cc
@@ -443,8 +443,8 @@ class IRSubstituteWithDataTypeLegalization : public
DataTypeLegalizer {
}
};
-Stmt SubstituteWithDataTypeLegalization(Stmt stmt,
-
std::function<ffi::Optional<PrimExpr>(const Var&)> vmap) {
+Stmt SubstituteWithDataTypeLegalization(
+ Stmt stmt, ffi::TypedFunction<ffi::Optional<PrimExpr>(const Var&)> vmap) {
auto general_vmap = [vmap = std::move(vmap)](const Var& var) ->
ffi::Optional<Expr> {
if (auto replacement = vmap(var)) return Expr(replacement.value());
return std::nullopt;
@@ -455,7 +455,7 @@ Stmt SubstituteWithDataTypeLegalization(Stmt stmt,
}
PrimExpr SubstituteWithDataTypeLegalization(
- PrimExpr expr, std::function<ffi::Optional<PrimExpr>(const Var&)> vmap) {
+ PrimExpr expr, ffi::TypedFunction<ffi::Optional<PrimExpr>(const Var&)>
vmap) {
auto general_vmap = [vmap = std::move(vmap)](const Var& var) ->
ffi::Optional<Expr> {
if (auto replacement = vmap(var)) return Expr(replacement.value());
return std::nullopt;
diff --git a/src/tirx/ir/transform.cc b/src/tirx/ir/transform.cc
index 051d1595..3b841cac 100644
--- a/src/tirx/ir/transform.cc
+++ b/src/tirx/ir/transform.cc
@@ -64,7 +64,8 @@ class FunctionPassObj : public PassObj {
PassInfo pass_info;
/*! \brief The pass function called on each. */
- std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func;
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ pass_func;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -96,14 +97,16 @@ class FunctionPass : public Pass {
* \param pass_info The pass info.
*/
TVM_TIRX_DLL FunctionPass(
- std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func,
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ pass_func,
PassInfo pass_info);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(FunctionPass, Pass,
FunctionPassObj);
};
FunctionPass::FunctionPass(
- std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func,
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ pass_func,
PassInfo pass_info)
: Pass(ffi::UnsafeInit{}) {
auto n = ffi::make_object<FunctionPassObj>(std::move(pass_info));
@@ -126,7 +129,7 @@ IRModule FunctionPassObj::operator()(IRModule mod, const
PassContext& pass_ctx)
// use move semantics as follows to avoid only copy.
kv.second.reset();
Function func = *std::move(opt_func);
- auto updated = pass_func(std::move(func), mod, pass_ctx);
+ auto updated = pass_func(ffi::RValueRef<Function>(std::move(func)), mod,
pass_ctx);
kv.second = Any(std::move(updated));
if (kv.second == nullptr) {
deleted_list.push_back(kv.first.as_or_throw<GlobalVar>());
@@ -144,7 +147,8 @@ IRModule FunctionPassObj::operator()(IRModule mod, const
PassContext& pass_ctx)
}
Pass CreateFunctionPass(
- std::function<ffi::Optional<Function>(Function, IRModule, PassContext)>
pass_func,
+ ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule, PassContext)>
+ pass_func,
int opt_level, ffi::String name) {
PassInfo pass_info = PassInfo(opt_level, name);
return FunctionPass(std::move(pass_func), pass_info);
@@ -155,15 +159,11 @@ TVM_FFI_STATIC_INIT_BLOCK() {
FunctionPassObj::RegisterReflection(); }
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def(
- "tirx.transform.CreateFunctionPass",
[](ffi::TypedFunction<ffi::Optional<Function>(
- ffi::RValueRef<Function>,
IRModule, PassContext)>
- pass_func,
- PassInfo pass_info) {
- auto wrapped_pass_func = [pass_func](Function func, IRModule mod,
PassContext ctx) {
- return pass_func(ffi::RValueRef<Function>(std::move(func)), mod,
ctx);
- };
- return FunctionPass(wrapped_pass_func, pass_info);
- });
+ "tirx.transform.CreateFunctionPass",
+ [](ffi::TypedFunction<ffi::Optional<Function>(ffi::RValueRef<Function>,
IRModule,
+ PassContext)>
+ pass_func,
+ PassInfo pass_info) { return FunctionPass(std::move(pass_func),
pass_info); });
}
// Pattern A (RM): auto-default repr from reflection.
diff --git a/src/tirx/ir/type.cc b/src/tirx/ir/type.cc
index e8754ad4..8839eb80 100644
--- a/src/tirx/ir/type.cc
+++ b/src/tirx/ir/type.cc
@@ -68,7 +68,7 @@ bool TensorTypeNode::IsScalar(bool alloc_or_decl) const {
layout.value().Canonicalize().as_or_throw<TileLayout>().IsTrivial()));
}
-std::optional<int64_t> TensorTypeNode::ConstantAllocationSize() const {
+ffi::Optional<int64_t> TensorTypeNode::ConstantAllocationSize() const {
int64_t result = 1;
for (const PrimExpr& extent : shape) {
const auto* size = extent.as<IntImmNode>();
