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 043ace5f1 [REFACTOR] Require sequence bodies and mark function
declarations (#153)
043ace5f1 is described below
commit 043ace5f1045097ff122606df80e022db62ef15d
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Nov 1 13:25:05 2026 -0500
[REFACTOR] Require sequence bodies and mark function declarations (#153)
---
include/tvm/ir/function.h | 9 +++++++-
include/tvm/script/ir_builder/frame.h | 4 ++--
include/tvm/script/ir_builder/ir.h | 2 +-
include/tvm/tirx/function.h | 14 +++++++------
include/tvm/tirx/script/ir_builder/frame.h | 2 +-
include/tvm/tirx/stmt_functor.h | 2 +-
python/tvm_tirx/tirx/function.py | 17 +++++++--------
rust/src/tirx/function.rs | 12 +++++------
rust/tests/tirx/function.rs | 12 +++++++----
rust/tests/tirx/main.rs | 2 +-
src/backend/cuda/codegen/codegen_cuda.cc | 6 +++---
src/backend/cuda/codegen/codegen_cuda_host.cc | 4 ++--
src/backend/cuda/transform/bind_backend_config.cc | 14 ++++++-------
src/backend/cuda/transform/lower_iket.cc | 16 +++++++--------
src/backend/metal/codegen/codegen_metal.cc | 4 ++--
src/backend/opencl/codegen/codegen_opencl.cc | 4 ++--
src/backend/trn/codegen/codegen_trn.cc | 4 ++--
src/backend/trn/transform/lower_trainium_layout.cc | 4 ++--
src/backend/vulkan/codegen/codegen_spirv.cc | 4 ++--
src/backend/webgpu/codegen/codegen_webgpu.cc | 6 +++---
src/codegen/build_common.h | 1 +
src/codegen/c/codegen_c.cc | 7 +++++--
src/codegen/c/codegen_c_host.cc | 4 ++++
src/codegen/llvm/codegen_aarch64.cc | 4 ++--
src/codegen/llvm/codegen_cpu.cc | 4 ++++
src/codegen/llvm/codegen_llvm.cc | 7 +++++--
src/codegen/llvm/llvm_module.cc | 1 +
src/script/ir_builder/frame.cc | 4 +++-
src/script/ir_builder/ir.cc | 8 +++++++-
src/tirx/analysis/check_well_formed.cc | 5 +++++
src/tirx/ir/function.cc | 24 +++++++++-------------
src/tirx/script/ir_builder/frame.cc | 21 ++++++++-----------
src/tirx/script/printer/function.cc | 5 ++---
src/tirx/tile/lower.cc | 6 +++---
src/tirx/transform/common_subexpr_elim.cc | 6 +++---
src/tirx/transform/decorate_device_scope.cc | 4 ++--
src/tirx/transform/flatten_tensor.cc | 8 ++++----
src/tirx/transform/hoist_if.cc | 6 +++---
src/tirx/transform/lower_intrin.cc | 2 +-
src/tirx/transform/lower_tensor_layouts.cc | 4 ++--
src/tirx/transform/lower_tvm_builtin.cc | 4 ++--
src/tirx/transform/make_packed_api.cc | 12 +++++------
src/tirx/transform/narrow_datatype.cc | 4 ++--
src/tirx/transform/remove_no_op.cc | 7 +++----
src/tirx/transform/renormalize_split_pattern.cc | 2 +-
src/tirx/transform/split_host_device.cc | 2 +-
src/tirx/transform/stmt_simplify.cc | 4 ++--
src/tirx/transform/storage_rewrite.cc | 9 ++++----
src/tirx/transform/unroll_loop.cc | 4 ++--
tests/cpp/tirx/thread_scope_test.cc | 4 +---
tests/python/tirx/ir/test_function.py | 4 ++--
tests/python/tirx/test_constructor.py | 6 +++---
52 files changed, 181 insertions(+), 154 deletions(-)
diff --git a/include/tvm/ir/function.h b/include/tvm/ir/function.h
index a6d01fb11..cf4e763bb 100644
--- a/include/tvm/ir/function.h
+++ b/include/tvm/ir/function.h
@@ -209,9 +209,16 @@ class BaseFunc : public Expr {
/*!
* \brief Generic attribute names that can be attached to any function.
*
- * \sa tvm::tirx::attr, tvm::relax::attr
+ * \sa tvm::tirx::attr
*/
namespace attr {
+/*!
+ * \brief Whether this function is a declaration.
+ *
+ * Type: bool
+ */
+inline constexpr ffi::reflection::AttrKey<bool> kDeclaration{"declaration"};
+
/*!
* \brief Indicates the special calling convention.
*
diff --git a/include/tvm/script/ir_builder/frame.h
b/include/tvm/script/ir_builder/frame.h
index cb078bf18..e0d00ad47 100644
--- a/include/tvm/script/ir_builder/frame.h
+++ b/include/tvm/script/ir_builder/frame.h
@@ -47,8 +47,8 @@ class IRModuleFrameObj : public IRBuilderFrameObj {
/*! \brief A map from string names to global variables that ensures global
uniqueness. */
ffi::Map<ffi::String, GlobalVar> global_var_map;
/*!
- * \brief A map from GlobalVar to all global functions.
- * \note Only defined functions are in the map, while declared functions are
not included.
+ * \brief A map from GlobalVar to global functions and pending declarations.
+ * \note Declarations carry attr::kDeclaration until replaced by a
definition.
*/
ffi::Map<GlobalVar, BaseFunc> functions;
/*! \brief IRModule's attributes. */
diff --git a/include/tvm/script/ir_builder/ir.h
b/include/tvm/script/ir_builder/ir.h
index 924c82d89..90d464c3c 100644
--- a/include/tvm/script/ir_builder/ir.h
+++ b/include/tvm/script/ir_builder/ir.h
@@ -46,7 +46,7 @@ TVM_TIRX_DLL IRModuleFrame IRModule();
* \brief Declare a Function without given the specific function
implementation.
* \note It is usually used in cross-function call. And we can specify the
function by `DefFunction`
* \param func_name The function unique name.
- * \param func_signature A Function w/o body, which used to specify the
function signature
+ * \param func_signature A Function marked with attr::kDeclaration, specifying
the signature
* (i.e. func params and func return type/shape).
* \return The corresponding GlobalVar.
*/
diff --git a/include/tvm/tirx/function.h b/include/tvm/tirx/function.h
index 8fd9de61c..224b1798f 100644
--- a/include/tvm/tirx/function.h
+++ b/include/tvm/tirx/function.h
@@ -53,8 +53,10 @@ class FunctionObj : public BaseFuncObj {
ffi::Array<tvm::Var> params;
/*! \brief The return type of the function. */
Type ret_type = Type::Missing();
- /*! \brief The body of the function, absent for a declaration. */
- ffi::Optional<SeqStmt> body;
+ /*! \brief The body of the function, empty for a marked declaration. */
+ SeqStmt body;
+
+ explicit FunctionObj(ffi::UnsafeInit tag) : body(tag) {}
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
@@ -108,7 +110,8 @@ class Function : public BaseFunc {
*
* \param params The parameters of the function.
*
- * \param body The body of the function, or std::nullopt for a declaration.
+ * \param body The body of the function. Declarations use an empty sequence
and
+ * tvm::attr::kDeclaration.
*
* \param ret_type The return type of the function.
*
@@ -116,9 +119,8 @@ class Function : public BaseFunc {
*
* \param loc The location of this object in the source code.
*/
- TVM_TIRX_DLL Function(ffi::Array<tvm::Var> params, ffi::Optional<SeqStmt>
body,
- Type ret_type = VoidType(), DictAttrs attrs =
DictAttrs(),
- Location loc = UnknownLoc());
+ TVM_TIRX_DLL Function(ffi::Array<tvm::Var> params, SeqStmt body, Type
ret_type = VoidType(),
+ DictAttrs attrs = DictAttrs(), Location loc =
UnknownLoc());
explicit Function(ffi::ObjectPtr<FunctionObj> node) :
BaseFunc(std::move(node)) {}
diff --git a/include/tvm/tirx/script/ir_builder/frame.h
b/include/tvm/tirx/script/ir_builder/frame.h
index 01427e4eb..2f7af4994 100644
--- a/include/tvm/tirx/script/ir_builder/frame.h
+++ b/include/tvm/tirx/script/ir_builder/frame.h
@@ -55,7 +55,7 @@ class TVM_TIRX_DLL FunctionFrameObj : public ir::StmtFrameObj
{
ffi::Map<ffi::String, Any> attrs;
/*! \brief Whether it is a persistent kernel. */
bool persistent;
- /*! \brief Whether this frame declares a bodyless signature. */
+ /*! \brief Whether this frame collects a signature for later definition. */
bool is_declaration{false};
/*! \brief Finalized function and its module identity. */
ffi::Optional<tvm::tirx::Function> function;
diff --git a/include/tvm/tirx/stmt_functor.h b/include/tvm/tirx/stmt_functor.h
index 54279a832..60e97354a 100644
--- a/include/tvm/tirx/stmt_functor.h
+++ b/include/tvm/tirx/stmt_functor.h
@@ -258,7 +258,7 @@ class TVM_TIRX_DLL StmtExprMutator : public
tvm::ExprMutator {
return ffi::Optional<SeqStmt>(std::move(result).ValueUnchecked());
}
- /*! \brief Mutate an optional statement, preserving an absent function body.
*/
+ /*! \brief Mutate an optional statement, preserving absence. */
TVM_FFI_INLINE UnchangedOr<ffi::Optional<Stmt>> Mutate(
const ffi::Optional<Stmt>& stmt, InplaceMode inplace_mode =
InplaceMode::kDisallow) {
if (!stmt.has_value()) return ffi::Unchanged();
diff --git a/python/tvm_tirx/tirx/function.py b/python/tvm_tirx/tirx/function.py
index d3c6fae0c..071fc90d3 100644
--- a/python/tvm_tirx/tirx/function.py
+++ b/python/tvm_tirx/tirx/function.py
@@ -34,14 +34,15 @@ class Function(BaseFunc, Scriptable):
"""A TIRx function definition or declaration.
The function's ``ty`` is a FuncType derived from parameter types and
- ``ret_type``. A present body is normalized to a SeqStmt.
+ ``ret_type``. The body is normalized to a SeqStmt.
+ A declaration has an empty body and the ``declaration`` attribute set to
True.
Parameters
----------
params : list[tvm_tirx.ir.Var]
Input parameters, in calling order.
- body : Stmt or Sequence[Stmt] or None
- The function body, or None for a declaration without an implementation.
+ body : Stmt or Sequence[Stmt]
+ The function body. An empty sequence is a valid body.
ret_type : tvm_tirx.ir.Type, optional
The declared return type. Omitting it or supplying MissingType produces
an empty TupleType, the type returned by the C++ VoidType() helper.
@@ -52,12 +53,12 @@ class Function(BaseFunc, Scriptable):
The source location of the function.
"""
- body: SeqStmt | None
+ body: SeqStmt
def __init__(
self,
params,
- body: Stmt | Sequence[Stmt] | None,
+ body: Stmt | Sequence[Stmt],
ret_type=None,
attrs=None,
loc: Location = UNKNOWN_LOC,
@@ -83,13 +84,13 @@ class Function(BaseFunc, Scriptable):
loc,
) # type: ignore
- def with_body(self, new_body: Stmt | Sequence[Stmt] | None, loc: Location
= UNKNOWN_LOC):
+ def with_body(self, new_body: Stmt | Sequence[Stmt], loc: Location =
UNKNOWN_LOC):
"""Create a function with the same parameters, return type, and
attributes but a new body.
Parameters
----------
- new_body : Stmt | Sequence[Stmt] | None
- The new body, or None for a declaration.
+ new_body : Stmt | Sequence[Stmt]
+ The new body.
loc : Location, optional
The source location of the new function. Defaults to UNKNOWN_LOC;
diff --git a/rust/src/tirx/function.rs b/rust/src/tirx/function.rs
index 289e4628b..e74fcf037 100644
--- a/rust/src/tirx/function.rs
+++ b/rust/src/tirx/function.rs
@@ -21,7 +21,7 @@
//! Tensor functions.
//!
//! The objects `python/tvm_tirx/tirx/function.py` defines. A [`Function`] is a
-//! [`BaseFunc`] with parameters, a return type and an optional body, built
with [`Function::new`];
+//! [`BaseFunc`] with parameters, a return type and a required sequence body,
built with [`Function::new`];
//! [`renew_def`] gives its variables fresh definitions.
// `custom-new` reserves `new` for a constructor with the C++ one's semantics
(`ty` derived from the
@@ -67,7 +67,7 @@ pub struct FunctionObj {
base: BaseFuncObj,
pub params: Array<Var>,
pub ret_type: Type,
- pub body: Option<SeqStmt>,
+ pub body: SeqStmt,
}
const _: () = {
@@ -102,7 +102,7 @@ impl FunctionObj {
attrs: DictAttrs,
params: Array<Var>,
ret_type: Type,
- body: Option<SeqStmt>,
+ body: SeqStmt,
) -> Self {
let base = BaseFuncObj::new(loc, ty, attrs);
Self { base, params, ret_type, body }
@@ -117,7 +117,7 @@ impl Function {
attrs: impl Into<DictAttrs>,
params: impl Into<Array<Var>>,
ret_type: impl Into<Type>,
- body: impl Into<Option<SeqStmt>>,
+ body: impl Into<SeqStmt>,
) -> Self {
let obj = FunctionObj::new(
loc.into(),
@@ -140,7 +140,7 @@ impl Function {
/// (C++ `Function(params, body, ret_type, attrs, loc)`).
pub fn new(
params: impl IntoIterator<Item = impl Into<Var>>,
- body: Option<&SeqStmt>,
+ body: &SeqStmt,
ret_type: impl Into<Type>,
attrs: impl Into<DictAttrs>,
loc: impl Into<Location>,
@@ -153,7 +153,7 @@ impl Function {
};
let param_types = params.iter().map(|param|
param.ty.clone()).collect::<Vec<Type>>();
let ty = FuncType::new(param_types, ret_type.clone());
- Self::from_fields(loc.into(), ty, attrs.into(), params, ret_type,
body.cloned())
+ Self::from_fields(loc.into(), ty, attrs.into(), params, ret_type,
body.clone())
}
}
diff --git a/rust/tests/tirx/function.rs b/rust/tests/tirx/function.rs
index 882e353cd..794f3eb2d 100644
--- a/rust/tests/tirx/function.rs
+++ b/rust/tests/tirx/function.rs
@@ -52,9 +52,13 @@ fn functions_agree_with_cpp() {
[(FfiString::from("note"), Any::from(1i64))].into_iter().collect();
let attrs = DictAttrs::new(note);
for (what, body, ret_type) in [
- ("Function", Some(body.clone()), &int32),
- ("Function(void)", Some(body.clone()), &missing),
- ("Function(no body)", None, &int32),
+ ("Function", body.clone(), &int32),
+ ("Function(void)", body.clone(), &missing),
+ (
+ "Function(empty body)",
+ SeqStmt::new(Vec::<tvm_tirx::ir::Stmt>::new(), UNKNOWN_LOC),
+ &int32,
+ ),
] {
let native_function = native(
"tirx.Function",
@@ -68,7 +72,7 @@ fn functions_agree_with_cpp() {
);
agree(
what,
- Ok(Function::new(params.iter(), body.as_ref(), ret_type, &attrs,
UNKNOWN_LOC)),
+ Ok(Function::new(params.iter(), &body, ret_type, &attrs,
UNKNOWN_LOC)),
native_function,
);
}
diff --git a/rust/tests/tirx/main.rs b/rust/tests/tirx/main.rs
index 7f2b95b3e..7a58fe011 100644
--- a/rust/tests/tirx/main.rs
+++ b/rust/tests/tirx/main.rs
@@ -98,7 +98,7 @@ fn void_function(params: Vec<Var>) -> Function {
.try_into()
.unwrap();
let missing: Type =
global("ir.MissingType").call_tuple(()).unwrap().try_into().unwrap();
- Function::new(params, Some(&body), missing, DictAttrs::new(Map::new()),
UNKNOWN_LOC)
+ Function::new(params, &body, missing, DictAttrs::new(Map::new()),
UNKNOWN_LOC)
}
fn module_of(name: &str, func: &Function) -> IRModule {
diff --git a/src/backend/cuda/codegen/codegen_cuda.cc
b/src/backend/cuda/codegen/codegen_cuda.cc
index 6d4f6b228..8cb871814 100644
--- a/src/backend/cuda/codegen/codegen_cuda.cc
+++ b/src/backend/cuda/codegen/codegen_cuda.cc
@@ -307,7 +307,7 @@ void CodeGenCUDA::InitFuncState(const Function& func) {
void CodeGenCUDA::DeclareFunction(const GlobalVar& gvar, const Function& func)
{
if (!RegisterFunctionName(gvar, func)) return;
// Definitions supply their body-derived qualifiers in AddFunction.
- if (!func->body.has_value()) {
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
InitFuncState(func);
PrintFunctionSignature(GetFunctionName(gvar), func, fwd_decl_stream);
fwd_decl_stream << ";\n";
@@ -316,7 +316,7 @@ void CodeGenCUDA::DeclareFunction(const GlobalVar& gvar,
const Function& func) {
void CodeGenCUDA::AddFunction(const GlobalVar& gvar, const Function& func) {
DeclareFunction(gvar, func);
- if (!func->body.has_value()) return;
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return;
InitFuncState(func);
std::ostringstream parameters;
PrintFunctionParameters(func, parameters);
@@ -327,7 +327,7 @@ void CodeGenCUDA::AddFunction(const GlobalVar& gvar, const
Function& func) {
stream << " {\n";
PreFunctionBody(func);
int func_scope = BeginScope();
- PrintStmt(func->body.value());
+ PrintStmt(func->body);
EndScope(func_scope);
PrintIndent();
stream << "}\n\n";
diff --git a/src/backend/cuda/codegen/codegen_cuda_host.cc
b/src/backend/cuda/codegen/codegen_cuda_host.cc
index ba19c9018..32e06edc0 100644
--- a/src/backend/cuda/codegen/codegen_cuda_host.cc
+++ b/src/backend/cuda/codegen/codegen_cuda_host.cc
@@ -138,7 +138,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
}
void AddFunction(const GlobalVar& gvar, const Function& func) override {
- if (!func->body.has_value()) {
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
DeclareFunction(gvar, func);
return;
}
@@ -148,7 +148,7 @@ class CodeGenCUDAHost : public CodeGenCHost {
PrintFunctionSignature(GetFunctionName(gvar), func, stream);
stream << " {\n TVM_FFI_SAFE_CALL_BEGIN();\n try {\n";
int scope = BeginScope();
- PrintStmt(func->body.value());
+ PrintStmt(func->body);
EndScope(scope);
// CUDADeviceGuard reports errors by throwing. Clear the reported CUDA
// last-error state so a subsequent successful launch is not blamed for it.
diff --git a/src/backend/cuda/transform/bind_backend_config.cc
b/src/backend/cuda/transform/bind_backend_config.cc
index 9d87b3ba2..13223a4ce 100644
--- a/src/backend/cuda/transform/bind_backend_config.cc
+++ b/src/backend/cuda/transform/bind_backend_config.cc
@@ -50,9 +50,8 @@ class BackendConfigBinder : public StmtExprMutator {
.cast<ffi::Array<ffi::Any>>();
target_ = resolved[0].cast<Target>();
defaults_ = resolved[1].cast<ffi::String>();
- if (!func->body.has_value()) return func;
- auto body =
- Mutate(func->body.value(),
InplaceMode::kDisallow).ValueOrUnchanged(func->body.value());
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return func;
+ auto body = Mutate(func->body,
InplaceMode::kDisallow).ValueOrUnchanged(func->body);
func.CopyOnWrite()->body = body;
return WithAttrs(func, {{tvm::attr::kTarget.name, target_},
{"backend_config", defaults_}});
}
@@ -102,9 +101,8 @@ class DeviceHelperSpecializer : public StmtExprMutator {
active_(active) {}
Function Rewrite(Function func) {
- if (!func->body.has_value()) return func;
- auto body =
- Mutate(func->body.value(),
InplaceMode::kDisallow).ValueOrUnchanged(func->body.value());
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return func;
+ auto body = Mutate(func->body,
InplaceMode::kDisallow).ValueOrUnchanged(func->body);
func.CopyOnWrite()->body = body;
return func;
}
@@ -219,9 +217,9 @@ Pass SpecializeEntryHelpers() {
pending.pop_back();
if (!reachable.insert(gv.get()).second) continue;
auto func = result->Lookup(gv).as<Function>();
- if (!func || !func.value()->body.has_value()) continue;
+ if (!func ||
func.value()->GetAttr(tvm::attr::kDeclaration).value_or(false)) continue;
auto collector = ffi::make_object<HelperCallCollector>();
- collector->Visit(func.value()->body.value());
+ collector->Visit(func.value()->body);
for (auto callee : collector->calls) {
if (result->functions.count(callee)) pending.push_back(callee);
}
diff --git a/src/backend/cuda/transform/lower_iket.cc
b/src/backend/cuda/transform/lower_iket.cc
index dae1cdadb..ec8792501 100644
--- a/src/backend/cuda/transform/lower_iket.cc
+++ b/src/backend/cuda/transform/lower_iket.cc
@@ -1199,24 +1199,24 @@ IRModule LowerIketImpl(IRModule module) {
if (!IketEnabled(module)) {
for (const auto& [global_var, base_function] : module->functions) {
const auto* function_node = base_function.as<FunctionObj>();
- if (!function_node || !function_node->body.has_value()) continue;
+ if (!function_node ||
function_node->GetAttr(tvm::attr::kDeclaration).value_or(false))
+ continue;
Function function = ffi::GetRef<Function>(function_node);
auto collector = ffi::make_object<AnnotationCollector>();
collector->Visit(function->body);
- TokenTensorSet tokens = CollectTokenTensors(function->body.value());
+ TokenTensorSet tokens = CollectTokenTensors(function->body);
if (!collector->has_annotations && tokens.empty()) continue;
if (collector->has_annotations) {
auto verifier = ffi::make_object<TokenVerifier>(tokens);
verifier->Visit(function->body);
- TokenDeclarationMap token_declarations =
CollectTokenDeclarations(function->body.value());
+ TokenDeclarationMap token_declarations =
CollectTokenDeclarations(function->body);
auto schema_verifier =
ffi::make_object<RangeEndSchemaVerifier>(token_declarations,
&collector->declarations);
schema_verifier->Visit(function->body);
ValidateRangeSchemas(collector->declarations);
}
auto strip = ffi::make_object<StripIket>(std::move(tokens));
- Stmt stripped =
-
strip->Mutate(function->body.value()).ValueOrUnchanged(function->body.value());
+ Stmt stripped =
strip->Mutate(function->body).ValueOrUnchanged(function->body);
Stmt body =
ffi::make_object<RemoveStrippedIketNoOps>()->Mutate(stripped).ValueOrUnchanged(stripped);
if (!function->body.same_as(body)) {
@@ -1230,7 +1230,7 @@ IRModule LowerIketImpl(IRModule module) {
std::vector<KernelIketInfo> kernels;
for (const auto& [global_var, base_function] : module->functions) {
const auto* function_node = base_function.as<FunctionObj>();
- if (!function_node || !function_node->body.has_value()) continue;
+ if (!function_node ||
function_node->GetAttr(tvm::attr::kDeclaration).value_or(false)) continue;
Function function = ffi::GetRef<Function>(function_node);
auto collector = ffi::make_object<AnnotationCollector>();
collector->Visit(function->body);
@@ -1240,10 +1240,10 @@ IRModule LowerIketImpl(IRModule module) {
TVM_FFI_CHECK(IsCudaDeviceFunction(function), ValueError)
<< "IKET annotations are only valid in a split CUDA device kernel";
- TokenTensorSet tokens = CollectTokenTensors(function->body.value());
+ TokenTensorSet tokens = CollectTokenTensors(function->body);
auto verifier = ffi::make_object<TokenVerifier>(tokens);
verifier->Visit(function->body);
- TokenDeclarationMap token_declarations =
CollectTokenDeclarations(function->body.value());
+ TokenDeclarationMap token_declarations =
CollectTokenDeclarations(function->body);
auto schema_verifier =
ffi::make_object<RangeEndSchemaVerifier>(token_declarations,
&collector->declarations);
schema_verifier->Visit(function->body);
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 62e5ae2a9..130de7f7c 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -91,7 +91,7 @@ CodeGenMetal::CodeGenMetal(Target target) : target_(target) {
}
void CodeGenMetal::AddFunction(const GlobalVar& gvar, const Function& func) {
- TVM_FFI_CHECK(func->body.has_value(), ValueError)
+ TVM_FFI_CHECK(!func->GetAttr(tvm::attr::kDeclaration).value_or(false),
ValueError)
<< "Kernel code generation requires a function body";
// NOTE: There is no inter-function calls among Metal kernels.
// For now we keep the metal codegen without inter-function call
@@ -208,7 +208,7 @@ void CodeGenMetal::AddFunction(const GlobalVar& gvar, const
Function& func) {
// the function scope.
stream << ") {\n";
int func_scope = this->BeginScope();
- this->PrintStmt(func->body.value());
+ this->PrintStmt(func->body);
this->EndScope(func_scope);
this->PrintIndent();
this->stream << "}\n\n";
diff --git a/src/backend/opencl/codegen/codegen_opencl.cc
b/src/backend/opencl/codegen/codegen_opencl.cc
index 9b3a577bd..fd249318e 100644
--- a/src/backend/opencl/codegen/codegen_opencl.cc
+++ b/src/backend/opencl/codegen/codegen_opencl.cc
@@ -132,8 +132,8 @@ CodeGenOpenCL::CodeGenOpenCL() {
void CodeGenOpenCL::InitFuncState(const Function& f) {
CodeGenC::InitFuncState(f);
texture_access_qualifiers_.clear();
- if (f->body.has_value()) {
-
this->SetTextureScope(ffi::make_object<InferTextureAccess>()->Infer(f->body.value()));
+ if (!f->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
+
this->SetTextureScope(ffi::make_object<InferTextureAccess>()->Infer(f->body));
}
for (Var arg : f->params) {
auto ptr_type = arg->ty.as<PtrTypeObj>();
diff --git a/src/backend/trn/codegen/codegen_trn.cc
b/src/backend/trn/codegen/codegen_trn.cc
index f0613bcc6..f14ace278 100644
--- a/src/backend/trn/codegen/codegen_trn.cc
+++ b/src/backend/trn/codegen/codegen_trn.cc
@@ -81,7 +81,7 @@ CodeGenTrainium::CodeGenTrainium(Target target) :
target_(target) {
}
void CodeGenTrainium::AddFunction(const GlobalVar& gvar, const Function& func)
{
- TVM_FFI_CHECK(func->body.has_value(), ValueError)
+ TVM_FFI_CHECK(!func->GetAttr(tvm::attr::kDeclaration).value_or(false),
ValueError)
<< "Kernel code generation requires a function body";
// NOTE: There is no inter-function calls among Trainium kernels.
// For now we keep the Trainium codegen without inter-function call
@@ -141,7 +141,7 @@ void CodeGenTrainium::AddFunction(const GlobalVar& gvar,
const Function& func) {
// the function scope.
stream << "):\n";
int func_scope = this->BeginScope();
- this->PrintStmt(func->body.value());
+ this->PrintStmt(func->body);
this->PrintIndent();
stream << "return ";
for (size_t i = 0; i < output_vids.size(); i++) {
diff --git a/src/backend/trn/transform/lower_trainium_layout.cc
b/src/backend/trn/transform/lower_trainium_layout.cc
index 1f1bcd664..44db1f01d 100644
--- a/src/backend/trn/transform/lower_trainium_layout.cc
+++ b/src/backend/trn/transform/lower_trainium_layout.cc
@@ -355,9 +355,9 @@ namespace transform {
Pass LowerTrainiumLayout() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
auto* n = f.CopyOnWrite();
- auto [body, params] = TrainiumLayoutApplier::Lower(n->body.value(),
n->params);
+ auto [body, params] = TrainiumLayoutApplier::Lower(n->body, n->params);
n->body = std::move(body);
n->params = std::move(params);
return f;
diff --git a/src/backend/vulkan/codegen/codegen_spirv.cc
b/src/backend/vulkan/codegen/codegen_spirv.cc
index 89d85c986..8e0fba03b 100644
--- a/src/backend/vulkan/codegen/codegen_spirv.cc
+++ b/src/backend/vulkan/codegen/codegen_spirv.cc
@@ -114,7 +114,7 @@ const VarObj* AsBufferVarObj(const Expr& expr) {
CodeGenSPIRV::CodeGenSPIRV(Target target) : spirv_support_(target) {}
runtime::SPIRVShader CodeGenSPIRV::BuildFunction(const Function& f, const
std::string& name) {
- TVM_FFI_CHECK(f->body.has_value(), ValueError)
+ TVM_FFI_CHECK(!f->GetAttr(tvm::attr::kDeclaration).value_or(false),
ValueError)
<< "Kernel code generation requires a function body";
this->InitFuncState();
TVM_FFI_ICHECK(f->HasNonzeroAttr(tvm::tirx::attr::kNoAlias.name))
@@ -178,7 +178,7 @@ runtime::SPIRVShader CodeGenSPIRV::BuildFunction(const
Function& f, const std::s
}
}
}
- this->Dispatch(f->body.value());
+ this->Dispatch(f->body);
builder_->SetLocalSize(func_ptr, workgroup_size_);
builder_->MakeInst(spv::OpReturn);
builder_->MakeInst(spv::OpFunctionEnd);
diff --git a/src/backend/webgpu/codegen/codegen_webgpu.cc
b/src/backend/webgpu/codegen/codegen_webgpu.cc
index 44069f57e..306e0f483 100644
--- a/src/backend/webgpu/codegen/codegen_webgpu.cc
+++ b/src/backend/webgpu/codegen/codegen_webgpu.cc
@@ -201,7 +201,7 @@ CodeGenWebGPU::CodeGenWebGPU(Target target) :
target_(target) {
}
runtime::FunctionInfo CodeGenWebGPU::AddFunction(const Function& f, bool
skip_readonly_decl) {
- TVM_FFI_CHECK(f->body.has_value(), ValueError)
+ TVM_FFI_CHECK(!f->GetAttr(tvm::attr::kDeclaration).value_or(false),
ValueError)
<< "Kernel code generation requires a function body";
// clear previous generated state.
this->InitFuncState(f);
@@ -241,7 +241,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const
Function& f, bool skip_re
// runtime classifies storage-buffer arguments by the serialized "handle"
dtype.
constexpr DLDataType kRuntimeOpaqueHandleType{kDLOpaqueHandle, 64, 1};
- WebGPUWorkGroupInfo info =
WebGPUWorkgroupInfoCollector::Collect(f->body.value());
+ WebGPUWorkGroupInfo info = WebGPUWorkgroupInfoCollector::Collect(f->body);
std::vector<Var> pod_args;
int num_buffer = 0;
@@ -354,7 +354,7 @@ runtime::FunctionInfo CodeGenWebGPU::AddFunction(const
Function& f, bool skip_re
<< val_pod_args << "." << packGridDimX << ") { return; }\n";
// the function scope.
int func_scope = this->BeginScope();
- this->PrintStmt(f->body.value());
+ this->PrintStmt(f->body);
this->EndScope(func_scope);
this->PrintIndent();
this->stream << "}\n\n";
diff --git a/src/codegen/build_common.h b/src/codegen/build_common.h
index 36d555ac1..d039be492 100644
--- a/src/codegen/build_common.h
+++ b/src/codegen/build_common.h
@@ -50,6 +50,7 @@ inline ffi::Map<ffi::String, runtime::FunctionInfo>
ExtractFuncInfo(const IRModu
TVM_FFI_ICHECK(kv.second->IsInstance<tirx::FunctionObj>())
<< "Can only lower IR Module with Functions";
auto f = kv.second.as_or_throw<tirx::Function>();
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) continue;
ffi::Array<DLDataType> arg_types;
ffi::Array<runtime::ArgExtraTags> arg_extra_tags;
diff --git a/src/codegen/c/codegen_c.cc b/src/codegen/c/codegen_c.cc
index bbe009df6..668a8dbd6 100644
--- a/src/codegen/c/codegen_c.cc
+++ b/src/codegen/c/codegen_c.cc
@@ -153,6 +153,9 @@ void CodeGenC::PrintFunctionParameters(const Function&
func, std::ostream& os) {
}
bool CodeGenC::RegisterFunctionName(const GlobalVar& gvar, const Function&
func) {
+ TVM_FFI_CHECK(!func->GetAttr(tvm::attr::kDeclaration).value_or(false) ||
func->body->seq.empty(),
+ ValueError)
+ << "Function declaration " << gvar->name_hint << " must have an empty
body";
if (internal_functions_.count(gvar)) return false;
auto function_name = [&]() -> ffi::String {
@@ -194,7 +197,7 @@ void CodeGenC::AddFunction(const GlobalVar& gvar, const
Function& f) {
// If the function has already been forward-declared, this is a
// no-op.
DeclareFunction(gvar, f);
- if (!f->body.has_value()) return;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return;
auto function_name = GetFunctionName(gvar);
// clear previous generated state.
@@ -204,7 +207,7 @@ void CodeGenC::AddFunction(const GlobalVar& gvar, const
Function& f) {
stream << " {\n";
this->PreFunctionBody(f);
int func_scope = this->BeginScope();
- this->PrintStmt(f->body.value());
+ this->PrintStmt(f->body);
this->EndScope(func_scope);
this->PrintIndent();
this->stream << "}\n\n";
diff --git a/src/codegen/c/codegen_c_host.cc b/src/codegen/c/codegen_c_host.cc
index 55a917f0a..6a88ef3f8 100644
--- a/src/codegen/c/codegen_c_host.cc
+++ b/src/codegen/c/codegen_c_host.cc
@@ -74,6 +74,10 @@ void CodeGenCHost::AddFunction(const GlobalVar& gvar, const
Function& func) {
void CodeGenCHost::AddFunction(const GlobalVar& gvar, const Function& func,
bool emit_fwd_func_decl) {
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
+ DeclareFunction(gvar, func);
+ return;
+ }
auto global_symbol = func->GetAttr(tvm::attr::kGlobalSymbol);
if (global_symbol) {
function_names_.push_back(global_symbol.value());
diff --git a/src/codegen/llvm/codegen_aarch64.cc
b/src/codegen/llvm/codegen_aarch64.cc
index c3fb8037d..1d9519048 100644
--- a/src/codegen/llvm/codegen_aarch64.cc
+++ b/src/codegen/llvm/codegen_aarch64.cc
@@ -84,8 +84,8 @@ void
CodeGenAArch64::SetComputeScopeAttributes(llvm::Function* func) {
void CodeGenAArch64::SetPStateAttributes(llvm::Function* func, const Function&
f) {
// These string Function attrs are exposed directly through T.func_attr and
with_attr.
if (auto sm =
f->GetAttr<ffi::String>(tvm::codegen::aarch64::attr::kPStateSM)) {
- // A locally streaming body does not change a bodyless declaration's
interface.
- if (sm.value() != "body" || f->body.has_value()) {
+ // A locally streaming body does not change a declaration's interface.
+ if (sm.value() != "body" ||
!f->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
func->addFnAttr(MakeStringRef("aarch64_pstate_sm_" + sm.value()));
}
}
diff --git a/src/codegen/llvm/codegen_cpu.cc b/src/codegen/llvm/codegen_cpu.cc
index 2f00189f1..84e5e12c7 100644
--- a/src/codegen/llvm/codegen_cpu.cc
+++ b/src/codegen/llvm/codegen_cpu.cc
@@ -232,6 +232,10 @@ llvm::DISubprogram* CodeGenCPU::CreateDebugFunction(const
GlobalVar& gvar, const
}
void CodeGenCPU::AddFunction(const GlobalVar& gvar, const Function& func) {
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
+ DeclareFunction(gvar, func);
+ return;
+ }
di_subprogram_ = CreateDebugFunction(gvar, func);
EmitDebugLocation(func->loc);
CodeGenLLVM::AddFunction(gvar, func);
diff --git a/src/codegen/llvm/codegen_llvm.cc b/src/codegen/llvm/codegen_llvm.cc
index 301ca6fa6..1b579bbcf 100644
--- a/src/codegen/llvm/codegen_llvm.cc
+++ b/src/codegen/llvm/codegen_llvm.cc
@@ -294,6 +294,9 @@ std::tuple<std::string, llvm::Function::LinkageTypes>
CodeGenLLVM::GetLinkage(
}
llvm::Function* CodeGenLLVM::DeclareFunctionInternal(const GlobalVar& gvar,
const Function& func) {
+ TVM_FFI_CHECK(!func->GetAttr(tvm::attr::kDeclaration).value_or(false) ||
func->body->seq.empty(),
+ ValueError)
+ << "Function declaration " << gvar->name_hint << " must have an empty
body";
if (auto it = functions_.find(gvar.get()); it != functions_.end()) {
return it->second;
}
@@ -349,7 +352,7 @@ void CodeGenLLVM::AddFunctionInternal(const GlobalVar&
gvar, const Function& f)
this->InitFuncState();
function_ = DeclareFunctionInternal(gvar, f);
- if (!f->body.has_value()) return;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return;
// set var map and align information
auto arg_it = function_->arg_begin();
@@ -368,7 +371,7 @@ void CodeGenLLVM::AddFunctionInternal(const GlobalVar&
gvar, const Function& f)
llvm::LLVMContext* ctx = llvm_target_->GetContext();
llvm::BasicBlock* entry = llvm::BasicBlock::Create(*ctx, "entry", function_);
builder_->SetInsertPoint(entry);
- this->Dispatch(f->body.value());
+ this->Dispatch(f->body);
// Body-local alignment assumptions need not hold on paths that returned
// before reaching them. Keep them as llvm.assume instead of strengthening
diff --git a/src/codegen/llvm/llvm_module.cc b/src/codegen/llvm/llvm_module.cc
index d76fa3763..bbdc613c5 100644
--- a/src/codegen/llvm/llvm_module.cc
+++ b/src/codegen/llvm/llvm_module.cc
@@ -263,6 +263,7 @@ void LLVMModuleObj::Init(const IRModule& mod, const Target&
target) {
continue;
}
auto f = kv.second.as_or_throw<tirx::Function>();
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) continue;
auto global_symbol = f->GetAttr(tvm::attr::kGlobalSymbol);
bool is_entry_func = f->HasNonzeroAttr(tvm::tirx::attr::kIsEntryFunc.name);
diff --git a/src/script/ir_builder/frame.cc b/src/script/ir_builder/frame.cc
index 935be5c66..1cd5c7bd6 100644
--- a/src/script/ir_builder/frame.cc
+++ b/src/script/ir_builder/frame.cc
@@ -103,7 +103,9 @@ void IRModuleFrameObj::ExitWithScope() {
for (const auto& kv : functions) {
const GlobalVar& gv = kv.first;
const BaseFunc& func = kv.second;
- TVM_FFI_CHECK(func.defined(), ValueError) << "function " << gv->name_hint
<< " is not defined";
+ TVM_FFI_CHECK(func.defined() &&
!func->GetAttr(tvm::attr::kDeclaration).value_or(false),
+ ValueError)
+ << "function " << gv->name_hint << " is declared but not defined";
func_map.Set(gv, func);
}
IRBuilder builder = IRBuilder::Current();
diff --git a/src/script/ir_builder/ir.cc b/src/script/ir_builder/ir.cc
index e3c50d320..d332944d7 100644
--- a/src/script/ir_builder/ir.cc
+++ b/src/script/ir_builder/ir.cc
@@ -98,12 +98,18 @@ void DefFunction(const ffi::String& func_name, const
BaseFunc& func) {
TVM_FFI_CHECK(it != frame->global_var_map.end(), ValueError)
<< "function " << func_name << " does not exist, please declare it
first.";
const GlobalVar& gv = (*it).second;
- frame->functions.Set(gv, func);
+ TVM_FFI_CHECK(!func->GetAttr(tvm::attr::kDeclaration).value_or(false),
ValueError)
+ << "function " << func_name << " requires a definition";
+ auto previous = frame->functions.Get(gv);
+ TVM_FFI_CHECK(!previous ||
previous.value()->GetAttr(tvm::attr::kDeclaration).value_or(false),
+ ValueError)
+ << "function " << func_name << " has already been defined";
if (auto ty = GetGlobalVarType(func)) {
gv->ty = ty.value();
} else {
TVM_FFI_THROW(InternalError) << "Unsupported function type: " <<
func->GetTypeKey();
}
+ frame->functions.Set(gv, func);
}
void ModuleAttrs(ffi::Map<ffi::String, Any> attrs, bool allow_overwrite) {
diff --git a/src/tirx/analysis/check_well_formed.cc
b/src/tirx/analysis/check_well_formed.cc
index faf0bbbbf..7261e9b62 100644
--- a/src/tirx/analysis/check_well_formed.cc
+++ b/src/tirx/analysis/check_well_formed.cc
@@ -290,6 +290,11 @@ class WellFormednessVerifier : public StmtExprVisitor {
}
Result Visit_(const FunctionObj* op) {
+ if (op->GetAttr(tvm::attr::kDeclaration).value_or(false)) {
+ Verify(false) << "ValueError: Function at <current location> is a
declaration, "
+ "but a definition is required.";
+ return std::nullopt;
+ }
int enclosing_depth = loop_depth_;
loop_depth_ = 0;
WithScope([&]() -> Result {
diff --git a/src/tirx/ir/function.cc b/src/tirx/ir/function.cc
index eb346a7e8..513481591 100644
--- a/src/tirx/ir/function.cc
+++ b/src/tirx/ir/function.cc
@@ -54,7 +54,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMutate(
}));
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Type>, mapped_ret_type,
mutator->MutateExpected(self->ret_type));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Stmt>>,
mapped_body,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
mutator->MutateExpected(self->body));
if (mapped_params.UnchangedOrSameAs(self->params) &&
mapped_ret_type.UnchangedOrSameAs(self->ret_type) &&
@@ -66,9 +66,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMutate(
copy->ret_type =
std::move(mapped_ret_type).ValueOrUnchanged(std::move(copy->ret_type));
if (!mapped_body.IsUnchanged()) {
auto replacement = std::move(mapped_body).ValueUnchecked();
- copy->body = replacement.has_value()
- ?
ffi::Optional<SeqStmt>(SeqStmt(std::move(replacement).value()))
- : std::nullopt;
+ copy->body = SeqStmt(std::move(replacement));
}
copy->ty = copy->func_type_annotation();
return ffi::Any(std::move(copy));
@@ -87,7 +85,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMaybeInplaceMut
TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(
ffi::UnchangedOr<Type>, mapped_ret_type,
mutator->MutateExpected(self->ret_type, ffi::InplaceMode::kAllow));
- TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<ffi::Optional<Stmt>>,
mapped_body,
+ TVM_FFI_S_MUTATE_ASSIGN_OR_RETURN(ffi::UnchangedOr<Stmt>, mapped_body,
mutator->MutateExpected(self->body,
ffi::InplaceMode::kAllow));
if (!mapped_params.UnchangedOrSameAs(self->params)) {
self->params = std::move(mapped_params).ValueUnchecked();
@@ -97,9 +95,7 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMaybeInplaceMut
}
if (!(mapped_body.IsUnchanged() ||
ffi::AnyView(mapped_body).same_as(self->body))) {
auto replacement = std::move(mapped_body).ValueUnchecked();
- self->body = replacement.has_value()
- ?
ffi::Optional<SeqStmt>(SeqStmt(std::move(replacement).value()))
- : std::nullopt;
+ self->body = SeqStmt(std::move(replacement));
}
// A nested signature node may change in place without replacing params or
ret_type.
self->ty = self->func_type_annotation();
@@ -109,14 +105,15 @@ TVM_FFI_INLINE ffi::Expected<ffi::UnchangedOr<ffi::Any>>
FunctionMaybeInplaceMut
} // namespace
// Get the function type of a Function
-Function::Function(ffi::Array<tvm::Var> params, ffi::Optional<SeqStmt> body,
Type ret_type,
- DictAttrs attrs, Location loc)
+Function::Function(ffi::Array<tvm::Var> params, SeqStmt body, Type ret_type,
DictAttrs attrs,
+ Location loc)
: BaseFunc(ffi::UnsafeInit{}) {
+ TVM_FFI_ICHECK(body.defined());
if (ret_type.as<MissingType>().has_value()) {
ret_type = VoidType();
}
- auto n = ffi::make_object<FunctionObj>();
+ auto n = ffi::make_object<FunctionObj>(ffi::UnsafeInit{});
n->params = std::move(params);
n->body = std::move(body);
n->ret_type = std::move(ret_type);
@@ -161,9 +158,8 @@ TVM_FFI_STATIC_INIT_BLOCK() {
ffi::FStructuralMutate::FromNative<&FunctionMaybeInplaceMutate>());
refl::GlobalDef().def(
- "tirx.Function",
- [](ffi::Array<tvm::Var> params, ffi::Optional<SeqStmt> body, Type
ret_type, DictAttrs attrs,
- Location loc) { return Function(params, body, ret_type, attrs, loc);
});
+ "tirx.Function", [](ffi::Array<tvm::Var> params, SeqStmt body, Type
ret_type, DictAttrs attrs,
+ Location loc) { return Function(params, body,
ret_type, attrs, loc); });
refl::GlobalDef().def("tirx.RenewDef", RenewDef);
}
diff --git a/src/tirx/script/ir_builder/frame.cc
b/src/tirx/script/ir_builder/frame.cc
index 46362e910..31d83fd9d 100644
--- a/src/tirx/script/ir_builder/frame.cc
+++ b/src/tirx/script/ir_builder/frame.cc
@@ -88,13 +88,14 @@ void FunctionFrameObj::ExitWithScope() {
}
TVM_FFI_CHECK(!is_declaration || stmts.empty(), ValueError)
<< "A function declaration cannot contain body statements";
- ffi::Optional<tvm::SeqStmt> body = std::nullopt;
- if (!is_declaration) body = AsStmt(stmts);
+ // The marker belongs to this signature, not the frame reused for its body.
+ auto function_attrs = attrs;
+ if (is_declaration) function_attrs.Set(tvm::attr::kDeclaration.name, true);
tvm::tirx::Function func(
/*params=*/args,
- /*body=*/body,
+ /*body=*/AsStmt(stmts),
/*ret_type=*/ret_type.value_or(TupleType::Empty()),
- /*attrs=*/attrs.defined() ? DictAttrs(attrs) : DictAttrs(),
+ /*attrs=*/DictAttrs(function_attrs),
/*loc=*/loc);
func = FinalizeFunction(std::move(func));
function = func;
@@ -109,18 +110,12 @@ void FunctionFrameObj::ExitWithScope() {
"function scope, if it's defined in a Module";
const ir::IRModuleFrame& frame = opt_frame.value();
const ffi::String& func_name = name.value_or("");
- if (!frame->global_var_map.count(func_name) ||
- !frame->functions.count(frame->global_var_map.at(func_name))) {
- // Case. First time visiting the function.
+ if (!frame->global_var_map.count(func_name)) {
+ // A complete function can also be registered on its first visit.
global_var = ir::DeclFunction(func_name, func);
- }
- // Define the function.
- // Note we do checks to disallow redefinition of functions inside the
`DefFunction`.
- if (!global_var.has_value()) {
+ } else {
TVM_FFI_CHECK(!is_declaration, ValueError) << "function " << func_name
<< " already exists";
global_var = frame->global_var_map.at(func_name);
- }
- if (!is_declaration) {
ir::DefFunction(func_name, func);
}
} else {
diff --git a/src/tirx/script/printer/function.cc
b/src/tirx/script/printer/function.cc
index 96271e0fe..3f7bf5b3c 100644
--- a/src/tirx/script/printer/function.cc
+++ b/src/tirx/script/printer/function.cc
@@ -94,9 +94,8 @@ void PrintFunction(DocTranslatorObj* d, const
tirx::FunctionObj* func, ExprDoc d
}
if (!decorator_keys.empty())
decorator = decorator->Call({}, decorator_keys, decorator_values);
- ffi::Array<StmtDoc> body;
- if (func->body.has_value())
- body = ToStmtDocArray(d->WithDocScope([&]() {
d->Translate(func->body.value()); }));
+ ffi::Array<StmtDoc> body =
+ ToStmtDocArray(d->WithDocScope([&]() { d->Translate(func->body);
}));
std::vector<std::pair<ffi::String, ffi::Any>> attrs;
for (const auto& [key, value] : func->attrs->dict) {
if (key != tvm::attr::kGlobalSymbol.name &&
diff --git a/src/tirx/tile/lower.cc b/src/tirx/tile/lower.cc
index e6504b709..fa0b4c00a 100644
--- a/src/tirx/tile/lower.cc
+++ b/src/tirx/tile/lower.cc
@@ -632,11 +632,11 @@ namespace transform {
Pass TileDispatch() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
Target target = ResolveTarget(f);
auto* n = f.CopyOnWrite();
- n->body = TileDispatcher::LowerOpCalls(n->body.value(), target);
- if (!NoOpCallVerifier::Verify(n->body.value())) {
+ n->body = TileDispatcher::LowerOpCalls(n->body, target);
+ if (!NoOpCallVerifier::Verify(n->body)) {
LOG(FATAL) << "Failed to lower the TIRx program: " << f;
}
return f;
diff --git a/src/tirx/transform/common_subexpr_elim.cc
b/src/tirx/transform/common_subexpr_elim.cc
index 850080b23..bc3e02114 100644
--- a/src/tirx/transform/common_subexpr_elim.cc
+++ b/src/tirx/transform/common_subexpr_elim.cc
@@ -882,12 +882,12 @@ namespace transform {
*/
Pass CommonSubexprElim() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
- auto [insert_before, expr_remap] = CSEPlanner::Plan(f->body.value());
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
+ auto [insert_before, expr_remap] = CSEPlanner::Plan(f->body);
if (!insert_before.empty()) {
auto* n = f.CopyOnWrite();
n->body = ffi::make_object<CSERewriter>(std::move(insert_before),
std::move(expr_remap))
- ->Rewrite(f->body.value());
+ ->Rewrite(f->body);
}
return f;
};
diff --git a/src/tirx/transform/decorate_device_scope.cc
b/src/tirx/transform/decorate_device_scope.cc
index 2cbad33c0..5d6164907 100644
--- a/src/tirx/transform/decorate_device_scope.cc
+++ b/src/tirx/transform/decorate_device_scope.cc
@@ -33,10 +33,10 @@ namespace transform {
Pass DecorateDeviceScope() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
static const Op device_scope = Op::Get("tirx.device_scope");
auto* n = f.CopyOnWrite();
- n->body = RegionStmt(device_scope, {}, {}, DictAttrs(),
std::move(n->body).value());
+ n->body = RegionStmt(device_scope, {}, {}, DictAttrs(),
std::move(n->body));
return f;
};
return CreateFunctionPass(pass_func, 0, "tirx.DecorateDeviceScope");
diff --git a/src/tirx/transform/flatten_tensor.cc
b/src/tirx/transform/flatten_tensor.cc
index ddb761dcb..0b4e58573 100644
--- a/src/tirx/transform/flatten_tensor.cc
+++ b/src/tirx/transform/flatten_tensor.cc
@@ -65,7 +65,7 @@ class TensorFlattener : public
StmtExprMutatorWithConstraintContext {
using StmtExprMutatorWithConstraintContext::Mutate;
using StmtExprMutatorWithConstraintContext::Mutate_;
static Function Flatten(Function func) {
- if (!func->body.has_value()) return func;
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return func;
auto pass = ffi::make_object<TensorFlattener>(sym::ConstraintContext());
pass->AddParamConstraints(func);
for (const Var& param : func->params) {
@@ -74,8 +74,8 @@ class TensorFlattener : public
StmtExprMutatorWithConstraintContext {
pass->Define(tensor.value());
}
}
- auto body_result = pass->Mutate(func->body, InplaceMode::kDisallow)
- .as_or_throw<UnchangedOr<ffi::Optional<SeqStmt>>>();
+ auto body_result =
+ pass->Mutate(func->body,
InplaceMode::kDisallow).as_or_throw<UnchangedOr<SeqStmt>>();
bool body_unchanged = body_result.UnchangedOrSameAs(func->body);
auto body = std::move(body_result).ValueOrUnchanged(func->body);
@@ -89,7 +89,7 @@ class TensorFlattener : public
StmtExprMutatorWithConstraintContext {
if (!old_buf.value().same_as(new_buf)) {
body = SeqStmt({Bind(new_buf, Call(new_buf.ty(), decl_tensor_op(),
{old_buf.value().data()}, {},
{new_buf.ty()})),
- std::move(body).value()});
+ std::move(body)});
body_unchanged = false;
}
}
diff --git a/src/tirx/transform/hoist_if.cc b/src/tirx/transform/hoist_if.cc
index 94be8ae59..7dcd9c418 100644
--- a/src/tirx/transform/hoist_if.cc
+++ b/src/tirx/transform/hoist_if.cc
@@ -240,11 +240,11 @@ namespace transform {
Pass HoistIf() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
auto* n = f.CopyOnWrite();
Stmt body = ffi::make_object<IfHoister>()
- ->Mutate(n->body.value(), InplaceMode::kAllow)
- .ValueOrUnchanged(n->body.value());
+ ->Mutate(n->body, InplaceMode::kAllow)
+ .ValueOrUnchanged(n->body);
n->body = tirx::ConvertSSA(std::move(body));
return f;
};
diff --git a/src/tirx/transform/lower_intrin.cc
b/src/tirx/transform/lower_intrin.cc
index 645e58170..9a37aa325 100644
--- a/src/tirx/transform/lower_intrin.cc
+++ b/src/tirx/transform/lower_intrin.cc
@@ -411,7 +411,7 @@ Pass LowerIntrin() {
n->body =
ffi::make_object<IntrinInjecter>(sym::ConstraintContext(),
target.value(), enable_fast_math)
->Mutate(n->body, InplaceMode::kAllow)
- .as_or_throw<UnchangedOr<ffi::Optional<SeqStmt>>>()
+ .as_or_throw<UnchangedOr<SeqStmt>>()
.ValueOrUnchanged(n->body);
return f;
};
diff --git a/src/tirx/transform/lower_tensor_layouts.cc
b/src/tirx/transform/lower_tensor_layouts.cc
index ebeb8c70a..fbe259110 100644
--- a/src/tirx/transform/lower_tensor_layouts.cc
+++ b/src/tirx/transform/lower_tensor_layouts.cc
@@ -390,10 +390,10 @@ namespace transform {
Pass LowerTensorLayouts() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
Target target = ResolveTarget(f);
auto* n = f.CopyOnWrite();
- auto [body, params] = LayoutApplier::Flatten(n->body.value(), n->params,
target);
+ auto [body, params] = LayoutApplier::Flatten(n->body, n->params, target);
n->body = std::move(body);
n->params = std::move(params);
return f;
diff --git a/src/tirx/transform/lower_tvm_builtin.cc
b/src/tirx/transform/lower_tvm_builtin.cc
index 6d8af98c6..bf2bc0272 100644
--- a/src/tirx/transform/lower_tvm_builtin.cc
+++ b/src/tirx/transform/lower_tvm_builtin.cc
@@ -61,7 +61,7 @@ class BuiltinLower : public StmtExprMutator {
using StmtExprMutator::Mutate;
using StmtExprMutator::Mutate_;
static Function Build(Function func) {
- if (!func->body.has_value()) return func;
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return func;
ffi::Optional<PrimExpr> device_type = std::nullopt;
bool preserve_ffi_kernel = false;
if (auto target = func->GetAttr(tvm::attr::kTarget)) {
@@ -71,7 +71,7 @@ class BuiltinLower : public StmtExprMutator {
}
auto mutator = ffi::make_object<BuiltinLower>(device_type,
preserve_ffi_kernel);
- func.CopyOnWrite()->body =
mutator->VisitBodyAndRealizeAlloca(func->body.value());
+ func.CopyOnWrite()->body = mutator->VisitBodyAndRealizeAlloca(func->body);
return func;
}
diff --git a/src/tirx/transform/make_packed_api.cc
b/src/tirx/transform/make_packed_api.cc
index c95beaf70..f98a4ad1f 100644
--- a/src/tirx/transform/make_packed_api.cc
+++ b/src/tirx/transform/make_packed_api.cc
@@ -221,7 +221,7 @@ ffi::Optional<ffi::String> RequiresPackedAPI(const
Function& func) {
}
Function MakePackedAPI(Function func) {
- if (!func->body.has_value()) return func;
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) return func;
auto global_symbol = RequiresPackedAPI(func);
if (!global_symbol.has_value()) {
return func;
@@ -272,8 +272,8 @@ Function MakePackedAPI(Function func) {
ffi::symbol::tvm_ffi_symbol_prefix +
global_symbol.value()}});
Stmt body = ffi::make_object<ReturnRewriter>(v_result)
- ->Mutate(func_ptr->body.value(), InplaceMode::kAllow)
- .ValueOrUnchanged(func_ptr->body.value());
+ ->Mutate(func_ptr->body, InplaceMode::kAllow)
+ .ValueOrUnchanged(func_ptr->body);
body = RegionStmt(tirx::compute_scope_op(), {StringImm(name_hint +
"_compute_")}, {}, DictAttrs(),
body);
// Set device context
@@ -311,7 +311,7 @@ Function MakePackedAPI(Function func) {
func_ptr->body = body;
func_ptr->params = args;
- ffi::Array<Var> undefined = UndefinedVars(func_ptr->body.value(),
func_ptr->params);
+ ffi::Array<Var> undefined = UndefinedVars(func_ptr->body, func_ptr->params);
TVM_FFI_ICHECK_EQ(undefined.size(), 0)
<< "In Function " << name_hint << " variables " << undefined
<< " are used, but are not passed in as API arguments";
@@ -342,10 +342,10 @@ Pass MakePackedAPI() {
for (const auto& [gvar, base_func] : mptr->functions) {
if (auto opt = base_func.as<Function>()) {
auto func = opt.value();
- if (!func->body.has_value()) continue;
+ if (func->GetAttr(tvm::attr::kDeclaration).value_or(false)) continue;
auto orig_func = func;
- if (auto body = SubroutineCallRewriter::Apply(packed_func_methods,
func->body.value())) {
+ if (auto body = SubroutineCallRewriter::Apply(packed_func_methods,
func->body)) {
func.CopyOnWrite()->body = body.value();
}
diff --git a/src/tirx/transform/narrow_datatype.cc
b/src/tirx/transform/narrow_datatype.cc
index 8dabe1b45..11abd4d43 100644
--- a/src/tirx/transform/narrow_datatype.cc
+++ b/src/tirx/transform/narrow_datatype.cc
@@ -320,9 +320,9 @@ namespace transform {
Pass NarrowDataType(int target_bits) {
auto pass_func = [target_bits](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
auto* n = f.CopyOnWrite();
- n->body =
ffi::make_object<NarrowDataTypeRewriter>(target_bits)->Rewrite(n->body.value());
+ n->body =
ffi::make_object<NarrowDataTypeRewriter>(target_bits)->Rewrite(n->body);
return f;
};
return CreateFunctionPass(pass_func, 0, "tirx.NarrowDataType");
diff --git a/src/tirx/transform/remove_no_op.cc
b/src/tirx/transform/remove_no_op.cc
index 0692ec16b..ec021dd27 100644
--- a/src/tirx/transform/remove_no_op.cc
+++ b/src/tirx/transform/remove_no_op.cc
@@ -251,16 +251,15 @@ namespace transform {
Pass RemoveNoOp() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
RemoveNoOpConfig config =
ctx->GetConfigOrDefault<RemoveNoOpConfig>("tirx.RemoveNoOp");
bool ignore_profiler_call = config->ignore_profiler_call;
{
auto* write_ptr = f.CopyOnWrite();
- write_ptr->body =
- NoOpRemover::Apply(std::move(write_ptr->body).value(),
sym::ConstraintContext(),
- ignore_profiler_call,
config->max_simplification_steps);
+ write_ptr->body = NoOpRemover::Apply(std::move(write_ptr->body),
sym::ConstraintContext(),
+ ignore_profiler_call,
config->max_simplification_steps);
}
return f;
};
diff --git a/src/tirx/transform/renormalize_split_pattern.cc
b/src/tirx/transform/renormalize_split_pattern.cc
index 1565acb29..a1ac65e08 100644
--- a/src/tirx/transform/renormalize_split_pattern.cc
+++ b/src/tirx/transform/renormalize_split_pattern.cc
@@ -215,7 +215,7 @@ Pass RenormalizeSplitPattern() {
auto* n = f.CopyOnWrite();
n->body =
ffi::make_object<SplitPatternReNormalizer>(sym::ConstraintContext())
->Mutate(n->body, InplaceMode::kAllow)
- .as_or_throw<UnchangedOr<ffi::Optional<SeqStmt>>>()
+ .as_or_throw<UnchangedOr<SeqStmt>>()
.ValueOrUnchanged(std::move(n->body));
return f;
};
diff --git a/src/tirx/transform/split_host_device.cc
b/src/tirx/transform/split_host_device.cc
index da85655db..121fafe96 100644
--- a/src/tirx/transform/split_host_device.cc
+++ b/src/tirx/transform/split_host_device.cc
@@ -691,7 +691,7 @@ class DeviceKernelMutator : public StmtExprMutator {
write_ptr->ret_type = VoidType();
Target target = func->GetAttr(tvm::attr::kTarget).value();
bool preserve_early_returns = target->kind->name == "cuda";
- write_ptr->body = ReturnRemover::Apply(write_ptr->body.value(),
!preserve_early_returns);
+ write_ptr->body = ReturnRemover::Apply(write_ptr->body,
!preserve_early_returns);
}
func = WithAttrs(std::move(func),
diff --git a/src/tirx/transform/stmt_simplify.cc
b/src/tirx/transform/stmt_simplify.cc
index d63b98ce6..3a6b45007 100644
--- a/src/tirx/transform/stmt_simplify.cc
+++ b/src/tirx/transform/stmt_simplify.cc
@@ -98,9 +98,9 @@ Function StmtSimplifier::Run(Function func) {
if (auto value = string_bindings.Get(var)) return
ffi::Any(*value);
return ffi::Unchanged();
})
- .as_or_throw<ffi::Optional<SeqStmt>>();
+ .as_or_throw<SeqStmt>();
n->body = Mutate(n->body, InplaceMode::kAllow)
- .as_or_throw<UnchangedOr<ffi::Optional<SeqStmt>>>()
+ .as_or_throw<UnchangedOr<SeqStmt>>()
.ValueOrUnchanged(n->body);
return func;
}
diff --git a/src/tirx/transform/storage_rewrite.cc
b/src/tirx/transform/storage_rewrite.cc
index 4d3e92190..aaa09d6c7 100644
--- a/src/tirx/transform/storage_rewrite.cc
+++ b/src/tirx/transform/storage_rewrite.cc
@@ -2269,7 +2269,7 @@ namespace transform {
Pass StorageRewrite() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
bool enable_reuse = true;
bool reuse_require_exact_matched_dtype = false;
bool merge_static_smem =
ctx->GetConfig<bool>("tirx.merge_static_smem").value_or(false);
@@ -2287,9 +2287,8 @@ Pass StorageRewrite() {
reuse_require_exact_matched_dtype = true;
}
auto* n = f.CopyOnWrite();
- n->body =
ffi::make_object<StoragePlanRewriter>()->Rewrite(std::move(n->body).value(),
- n->params,
true, enable_reuse,
-
reuse_require_exact_matched_dtype);
+ n->body = ffi::make_object<StoragePlanRewriter>()->Rewrite(
+ std::move(n->body), n->params, true, enable_reuse,
reuse_require_exact_matched_dtype);
// Parameters may not be rewritten, but internal allocations may.
return PointerValueTypeRewrite(std::move(f), true, false, false, true,
true, true, false);
};
@@ -2303,7 +2302,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
Pass PointerValueTypeRewrite() {
auto pass_func = [](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
return PointerValueTypeRewrite(std::move(f));
};
return CreateFunctionPass(pass_func, 0, "tirx.PointerValueTypeRewrite");
diff --git a/src/tirx/transform/unroll_loop.cc
b/src/tirx/transform/unroll_loop.cc
index 53d09f43b..833b41d9b 100644
--- a/src/tirx/transform/unroll_loop.cc
+++ b/src/tirx/transform/unroll_loop.cc
@@ -368,10 +368,10 @@ namespace transform {
Pass UnrollLoop() {
auto pass_func = [=](Function f, IRModule m, PassContext ctx) {
- if (!f->body.has_value()) return f;
+ if (f->GetAttr(tvm::attr::kDeclaration).value_or(false)) return f;
auto* n = f.CopyOnWrite();
auto cfg = ctx->GetConfigOrDefault<UnrollLoopConfig>("tirx.UnrollLoop");
- n->body = UnrollLoop(std::move(f->body).value(), cfg);
+ n->body = UnrollLoop(std::move(f->body), cfg);
return f;
};
return CreateFunctionPass(pass_func, 0, "tirx.UnrollLoop");
diff --git a/tests/cpp/tirx/thread_scope_test.cc
b/tests/cpp/tirx/thread_scope_test.cc
index c3adc3959..1e2e86e6f 100644
--- a/tests/cpp/tirx/thread_scope_test.cc
+++ b/tests/cpp/tirx/thread_scope_test.cc
@@ -97,9 +97,7 @@ TEST(ThreadScope, TypedCallAndDispatchWithoutPython) {
auto result = transform::TileDispatch()(IRModule({{main, func}}));
auto lowered = result->Lookup(main).as_or_throw<Function>();
EXPECT_EQ(*invoked, 1);
- EXPECT_EQ(
-
lowered->body.value()->seq[0].as_or_throw<Evaluate>()->value.as_or_throw<IntImm>()->value,
- 37);
+
EXPECT_EQ(lowered->body->seq[0].as_or_throw<Evaluate>()->value.as_or_throw<IntImm>()->value,
37);
EXPECT_THROW(refl::ObjectCreator("tirx.tensor.ScopeAttrs")({}), ffi::Error);
EXPECT_THROW(refl::ObjectCreator("tirx.tensor.ScopeAttrs")({{"scope",
ffi::String("inner")}}),
ffi::Error);
diff --git a/tests/python/tirx/ir/test_function.py
b/tests/python/tirx/ir/test_function.py
index 5ddc96633..f5bf5a7b2 100644
--- a/tests/python/tirx/ir/test_function.py
+++ b/tests/python/tirx/ir/test_function.py
@@ -53,8 +53,8 @@ def test_structural_map_specializes_function_signature():
# place while a nested TensorType changes. The cached FuncType must follow.
def declaration():
return Function(
- [Var("tensor", TensorType((n,), "float32"))], None,
TensorType((n,), "float32")
- )
+ [Var("tensor", TensorType((n,), "float32"))], [], TensorType((n,),
"float32")
+ ).with_attr("declaration", True)
inplace = declaration()
inplace = tvm_ffi.structural_map(
diff --git a/tests/python/tirx/test_constructor.py
b/tests/python/tirx/test_constructor.py
index 04c507700..40aad8975 100644
--- a/tests/python/tirx/test_constructor.py
+++ b/tests/python/tirx/test_constructor.py
@@ -352,9 +352,9 @@ def test_math_unary_constructor_preserves_bfloat16():
assert y.ty == tvm_tirx.ir.PrimType("bfloat16")
-def test_function_declaration_optional_body_roundtrip():
- declaration = tvm_tirx.tirx.Function([], None)
- assert declaration.body is None
+def test_function_declaration_required_body_roundtrip():
+ declaration = tvm_tirx.tirx.Function([], []).with_attr("declaration", True)
+ assert len(declaration.body) == 0
assert len(declaration.with_body(tvm_tirx.ir.Evaluate(0)).body) == 0
restored = tvm_tirx.ir.load_json(tvm_tirx.ir.save_json(declaration))
tvm_tirx.ir.assert_structural_equal(declaration, restored,
map_free_vars=True)
