You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2026/03/02 12:52:59 UTC
(tvm) branch main updated: [TIR][Refactor] Enhance error reporting with structured AssertStmt and TVMFFIABIBuilder (#18857)
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.git
The following commit(s) were added to refs/heads/main by this push:
new 611a815dc1 [TIR][Refactor] Enhance error reporting with structured AssertStmt and TVMFFIABIBuilder (#18857)
611a815dc1 is described below
commit 611a815dc102ef188f7ca44faee8240552347f4b
Author: Tianqi Chen <tq...@users.noreply.github.com>
AuthorDate: Mon Mar 2 07:52:53 2026 -0500
[TIR][Refactor] Enhance error reporting with structured AssertStmt and TVMFFIABIBuilder (#18857)
---
include/tvm/script/ir_builder/tir/frame.h | 9 +-
include/tvm/script/ir_builder/tir/ir.h | 6 +-
include/tvm/tir/builtin.h | 30 +-
include/tvm/tir/stmt.h | 24 +-
.../tvm/relax/transform/legalize_ops/inspect_op.py | 38 +-
python/tvm/script/ir_builder/tir/ir.py | 15 +-
python/tvm/script/parser/tir/parser.py | 32 +-
python/tvm/tir/stmt.py | 25 +-
src/arith/ir_mutator_with_analyzer.cc | 4 +-
src/arith/ir_visitor_with_analyzer.cc | 1 -
src/relax/op/tensor/inspect.cc | 22 +-
src/s_tir/analysis/estimate_flops.cc | 3 -
src/s_tir/transform/bound_checker.cc | 3 +-
src/script/ir_builder/tir/frame.cc | 4 +-
src/script/ir_builder/tir/ir.cc | 10 +-
src/script/printer/tir/stmt.cc | 20 +-
src/target/llvm/codegen_cpu.cc | 147 +++-
src/target/llvm/codegen_cpu.h | 7 +
src/target/llvm/codegen_llvm.cc | 8 +-
src/target/source/codegen_c.cc | 107 ++-
src/target/source/codegen_c.h | 2 +
src/target/source/codegen_c_host.cc | 15 +-
src/tir/ir/buffer.cc | 2 +-
src/tir/ir/stmt.cc | 16 +-
src/tir/ir/stmt_functor.cc | 14 +-
src/tir/ir/tir_visitor_with_path.cc | 3 +-
src/tir/transform/arg_binder.cc | 329 ---------
src/tir/transform/arg_binder.h | 169 -----
src/tir/transform/lower_tvm_builtin.cc | 27 +-
src/tir/transform/make_packed_api.cc | 148 +---
src/tir/transform/split_host_device.cc | 4 +-
src/tir/transform/tvm_ffi_binder.cc | 810 +++++++++++++++++++++
src/tir/transform/tvm_ffi_binder.h | 431 +++++++++++
tests/python/codegen/test_codegen_assert.py | 124 ++++
.../python/codegen/test_codegen_error_handling.py | 466 ++++++++++++
tests/python/codegen/test_target_codegen_llvm.py | 36 +-
tests/python/relax/test_analysis_well_formed.py | 12 +-
tests/python/tir-base/test_tir_constructor.py | 9 +-
.../test_tir_transform_make_packed_api.py | 294 ++++----
.../tvmscript/test_tvmscript_ir_builder_tir.py | 6 +-
.../python/tvmscript/test_tvmscript_printer_tir.py | 2 +-
tests/python/tvmscript/test_tvmscript_roundtrip.py | 48 ++
42 files changed, 2515 insertions(+), 967 deletions(-)
diff --git a/include/tvm/script/ir_builder/tir/frame.h b/include/tvm/script/ir_builder/tir/frame.h
index d626df38c2..51d67ae963 100644
--- a/include/tvm/script/ir_builder/tir/frame.h
+++ b/include/tvm/script/ir_builder/tir/frame.h
@@ -304,14 +304,17 @@ class AssertFrameNode : public TIRFrameNode {
public:
/*! \brief The PrimExpr to test. */
PrimExpr condition;
- /*! \brief The output error message when the assertion failed. */
- PrimExpr message;
+ /*! \brief The error kind, e.g. "RuntimeError", "TypeError", "ValueError". */
+ tvm::tir::StringImm error_kind;
+ /*! \brief Error message fragments, concatenated at runtime when assertion fails. */
+ ffi::Array<tvm::tir::StringImm> message_parts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AssertFrameNode>()
.def_ro("condition", &AssertFrameNode::condition)
- .def_ro("message", &AssertFrameNode::message);
+ .def_ro("error_kind", &AssertFrameNode::error_kind)
+ .def_ro("message_parts", &AssertFrameNode::message_parts);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("script.ir_builder.tir.AssertFrame", AssertFrameNode,
TIRFrameNode);
diff --git a/include/tvm/script/ir_builder/tir/ir.h b/include/tvm/script/ir_builder/tir/ir.h
index 694bd57aab..caf01081dc 100644
--- a/include/tvm/script/ir_builder/tir/ir.h
+++ b/include/tvm/script/ir_builder/tir/ir.h
@@ -286,10 +286,12 @@ ForFrame Grid(ffi::Array<PrimExpr> extents);
/*!
* \brief The assertion statement.
* \param condition The assertion condition.
- * \param message The error message when the assertion fails.
+ * \param error_kind The error kind (e.g. "RuntimeError", "TypeError", "ValueError").
+ * \param message_parts The error message parts (stored as separate fragments in the IR).
* \return The AssertFrame.
*/
-AssertFrame Assert(PrimExpr condition, ffi::String message);
+AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
+ ffi::Array<ffi::String> message_parts);
/*!
* \brief The let binding.
diff --git a/include/tvm/tir/builtin.h b/include/tvm/tir/builtin.h
index 2e5c10a9d9..4650edc054 100644
--- a/include/tvm/tir/builtin.h
+++ b/include/tvm/tir/builtin.h
@@ -963,25 +963,27 @@ TVM_DLL const Op& ignore_loop_partition();
/*! \brief The kind of structure field info used in intrinsic */
enum TVMStructFieldKind : int {
- // array head address
- kArrAddr,
- kArrData,
- kArrShape,
- kArrStrides,
- kArrNDim,
- kArrTypeCode,
- kArrTypeBits,
- kArrTypeLanes,
- kArrByteOffset,
- kArrDeviceId,
- kArrDeviceType,
- kArrKindBound_,
+ // DLTensor fields
+ kDLTensorAddr,
+ kDLTensorData,
+ kDLTensorShape,
+ kDLTensorStrides,
+ kDLTensorNDim,
+ kDLTensorTypeCode,
+ kDLTensorTypeBits,
+ kDLTensorTypeLanes,
+ kDLTensorByteOffset,
+ kDLTensorDeviceId,
+ kDLTensorDeviceType,
+ kDLTensorKindBound_,
// TVMValue field
kTVMValueContent,
kTVMFFIAnyTypeIndex,
kTVMFFIAnyZeroPadding,
kTVMFFIAnyUnionValue,
- kTVMValueKindBound_
+ kTVMValueKindBound_,
+ // Generic int64 array element access: ((int64_t*)buf)[index]
+ kInt64ArrayElem,
};
} // namespace builtin
} // namespace tir
diff --git a/include/tvm/tir/stmt.h b/include/tvm/tir/stmt.h
index c12e14e0c0..b6424c1da3 100644
--- a/include/tvm/tir/stmt.h
+++ b/include/tvm/tir/stmt.h
@@ -147,25 +147,28 @@ class AttrStmt : public Stmt {
/*!
* \brief Assert condition, if an error occurs, return the error message.
+ *
+ * The error is described by:
+ * - \p error_kind: the error kind (e.g. "RuntimeError", "TypeError", "ValueError")
+ * - \p message_parts: an array of string fragments that are concatenated at runtime
+ * via TVMFFIErrorSetRaisedFromCStrParts. This enables string fragment reuse
+ * across multiple assertions to reduce binary size.
*/
class AssertStmtNode : public StmtNode {
public:
/*! \brief Condition to be checked. */
PrimExpr condition;
- /*! \brief Error message when assertion failed. */
- PrimExpr message;
- /*!
- * \brief Body which this assertion holds true.
- * Will be executed after the assertion.
- */
- Stmt body;
+ /*! \brief The error kind, e.g. "RuntimeError", "TypeError", "ValueError". */
+ StringImm error_kind;
+ /*! \brief Error message fragments, concatenated at runtime when assertion fails. */
+ ffi::Array<StringImm> message_parts;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<AssertStmtNode>()
.def_ro("condition", &AssertStmtNode::condition)
- .def_ro("message", &AssertStmtNode::message)
- .def_ro("body", &AssertStmtNode::body);
+ .def_ro("error_kind", &AssertStmtNode::error_kind)
+ .def_ro("message_parts", &AssertStmtNode::message_parts);
}
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tir.AssertStmt", AssertStmtNode, StmtNode);
};
@@ -176,7 +179,8 @@ class AssertStmtNode : public StmtNode {
*/
class AssertStmt : public Stmt {
public:
- TVM_DLL AssertStmt(PrimExpr condition, PrimExpr message, Span span = Span());
+ TVM_DLL AssertStmt(PrimExpr condition, StringImm error_kind, ffi::Array<StringImm> message_parts,
+ Span span = Span());
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(AssertStmt, Stmt, AssertStmtNode);
TVM_DEFINE_OBJECT_REF_COW_METHOD(AssertStmtNode);
diff --git a/python/tvm/relax/transform/legalize_ops/inspect_op.py b/python/tvm/relax/transform/legalize_ops/inspect_op.py
index 2a9ff0775b..6edba59bc4 100644
--- a/python/tvm/relax/transform/legalize_ops/inspect_op.py
+++ b/python/tvm/relax/transform/legalize_ops/inspect_op.py
@@ -35,18 +35,18 @@ class TVMStructFieldKind(enum.IntEnum):
definition which starts from 0.
"""
- kArrAddr = 0
- kArrData = 1
- kArrShape = 2
- kArrStrides = 3
- kArrNDim = 4
- kArrTypeCode = 5
- kArrTypeBits = 6
- kArrTypeLanes = 7
- kArrByteOffset = 8
- kArrDeviceId = 9
- kArrDeviceType = 10
- kArrKindBound_ = 11
+ kDLTensorAddr = 0
+ kDLTensorData = 1
+ kDLTensorShape = 2
+ kDLTensorStrides = 3
+ kDLTensorNDim = 4
+ kDLTensorTypeCode = 5
+ kDLTensorTypeBits = 6
+ kDLTensorTypeLanes = 7
+ kDLTensorByteOffset = 8
+ kDLTensorDeviceId = 9
+ kDLTensorDeviceType = 10
+ kDLTensorKindBound_ = 11
kTVMValueContent = 12
kTVMValueKindBound_ = 13
@@ -58,18 +58,18 @@ def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
T.func_attr({"tir.is_host": True, "tir.is_scheduled": True})
assert T.int64(0) <= axis, "Specified axis may not be negative"
ndim: T.int32 = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrNDim), "int32"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), "int32"
)
assert axis < T.Cast("int64", ndim), (
"Specified axis may not be larger than the tensor's dimensionality"
)
stride_ptr: T.handle("int64") = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrStrides), "handle"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), "handle"
)
if T.isnullptr(stride_ptr):
shape_ptr: T.handle("int64") = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrShape), "handle"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), "handle"
)
shape = T.decl_buffer(ndim, "int64", data=shape_ptr)
@@ -99,7 +99,7 @@ def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
def _get_tensor_byte_offset(dlpack_handle: T.handle) -> T.int64:
T.func_attr({"tir.is_host": True, "tir.is_scheduled": True})
byte_offset: T.uint64 = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrByteOffset), "uint64"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
)
return byte_offset
@@ -113,13 +113,13 @@ def _tensor_elem_offset(bb: BlockBuilder, call: Call) -> Expr:
def _get_tensor_elem_offset(dlpack_handle: T.handle) -> T.int64:
T.func_attr({"tir.is_host": True, "tir.is_scheduled": True})
byte_offset: T.uint64 = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrByteOffset), "uint64"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
)
scalar_bits: T.uint8 = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrTypeBits), "uint8"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), "uint8"
)
lanes: T.uint16 = T.tvm_struct_get(
- dlpack_handle, 0, int(TVMStructFieldKind.kArrTypeLanes), "uint16"
+ dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), "uint16"
)
bytes_per_element = T.ceildiv(scalar_bits.astype("uint64") * lanes.astype("uint64"), 8)
elem_offset = byte_offset // bytes_per_element
diff --git a/python/tvm/script/ir_builder/tir/ir.py b/python/tvm/script/ir_builder/tir/ir.py
index 1d259e192a..860dae3f4a 100644
--- a/python/tvm/script/ir_builder/tir/ir.py
+++ b/python/tvm/script/ir_builder/tir/ir.py
@@ -920,7 +920,7 @@ def grid(*extents: PrimExpr) -> frame.ForFrame:
return _ffi_api.Grid(extents) # type: ignore[attr-defined] # pylint: disable=no-member
-def Assert(condition: PrimExpr, message: str) -> frame.AssertFrame: # pylint: disable=invalid-name
+def Assert(condition: PrimExpr, message, error_kind: str = "RuntimeError") -> frame.AssertFrame: # pylint: disable=invalid-name
"""Create an assertion statement.
Parameters
@@ -928,8 +928,13 @@ def Assert(condition: PrimExpr, message: str) -> frame.AssertFrame: # pylint: d
condition : PrimExpr
The PrimExpr to test.
- message : str
- The output error message when the assertion fails.
+ message : str or list[str]
+ The error message when the assertion fails. Can be a single string
+ or a list of string parts (fragments stored separately in the IR
+ for binary size reduction through string reuse).
+
+ error_kind : str
+ The error kind (e.g. "RuntimeError", "TypeError", "ValueError").
Returns
-------
@@ -938,7 +943,9 @@ def Assert(condition: PrimExpr, message: str) -> frame.AssertFrame: # pylint: d
"""
if isinstance(condition, bool):
condition = IntImm("bool", condition)
- return _ffi_api.Assert(condition, message) # type: ignore[attr-defined] # pylint: disable=no-member
+ if not isinstance(message, list | tuple):
+ message = [message]
+ return _ffi_api.Assert(condition, error_kind, message) # type: ignore[attr-defined] # pylint: disable=no-member
def LetStmt( # pylint: disable=invalid-name
diff --git a/python/tvm/script/parser/tir/parser.py b/python/tvm/script/parser/tir/parser.py
index 5cad41e216..e20629795c 100644
--- a/python/tvm/script/parser/tir/parser.py
+++ b/python/tvm/script/parser/tir/parser.py
@@ -537,10 +537,40 @@ def visit_assert(self: Parser, node: doc.Assert) -> None:
node : doc.Assert
The doc AST assert node.
+
+ The assert message can be either:
+ - A plain string: ``assert cond, "message"``
+ - A tuple of (kind, [parts...]): ``assert cond, ("ValueError", ["part0", "part1"])``
"""
cond = self.eval_expr(node.test)
msg = self.eval_expr(node.msg)
- frame = T.Assert(cond, msg)
+
+ kind = "RuntimeError"
+ message = msg
+
+ if isinstance(msg, tuple):
+ if len(msg) != 2:
+ self.report_error(
+ node,
+ f"Assert message tuple must have exactly 2 elements (kind, [parts...]), "
+ f"got {len(msg)} elements",
+ )
+ kind_str, parts = msg
+ if isinstance(kind_str, tvm.tir.StringImm):
+ kind_str = kind_str.value
+ if not isinstance(kind_str, str):
+ self.report_error(
+ node,
+ f"Assert message tuple first element must be a string (error kind like "
+ f'"ValueError"), got {type(kind_str).__name__}',
+ )
+ kind = kind_str
+ message = parts
+
+ if isinstance(message, list | tuple):
+ message = [p.value if isinstance(p, tvm.tir.StringImm) else str(p) for p in message]
+
+ frame = T.Assert(cond, message, error_kind=kind)
frame.add_callback(partial(frame.__exit__, None, None, None))
frame.__enter__()
diff --git a/python/tvm/tir/stmt.py b/python/tvm/tir/stmt.py
index 093adc4fed..004d350917 100644
--- a/python/tvm/tir/stmt.py
+++ b/python/tvm/tir/stmt.py
@@ -37,7 +37,7 @@ from tvm.runtime import Object, Scriptable, const
from . import _ffi_api
from .buffer import Buffer
-from .expr import IterVar, Var
+from .expr import IterVar, StringImm, Var
class Stmt(Object, Scriptable):
@@ -84,25 +84,38 @@ class AssertStmt(Stmt):
Parameters
----------
+ kind : StringImm
+ The error kind, e.g. "RuntimeError", "TypeError", "ValueError".
+
condition : PrimExpr
The assert condition.
- message : PrimExpr
- The error message.
+ message_parts : list[StringImm]
+ Error message fragments, concatenated at runtime when assertion fails.
span : Span | None
The location of the stmt in the source code.
"""
+ kind: StringImm
condition: PrimExpr
- message: PrimExpr
+ message_parts: list
span: Span | None
- def __init__(self, condition: PrimExpr, message: PrimExpr, span: Span | None = None) -> None:
+ def __init__(
+ self,
+ kind: StringImm,
+ condition: PrimExpr,
+ message_parts: list | None = None,
+ span: Span | None = None,
+ ) -> None:
+ if message_parts is None:
+ message_parts = []
self.__init_handle_by_constructor__(
_ffi_api.AssertStmt,
+ kind,
condition,
- message,
+ message_parts,
span, # type: ignore
)
diff --git a/src/arith/ir_mutator_with_analyzer.cc b/src/arith/ir_mutator_with_analyzer.cc
index 96a64c5e47..7ba8059336 100644
--- a/src/arith/ir_mutator_with_analyzer.cc
+++ b/src/arith/ir_mutator_with_analyzer.cc
@@ -155,15 +155,13 @@ Stmt IRMutatorWithAnalyzer::VisitStmt_(const AttrStmtNode* op) {
Stmt IRMutatorWithAnalyzer::VisitStmt_(const AssertStmtNode* op) {
PrimExpr condition = this->VisitExpr(op->condition);
- PrimExpr message = this->VisitExpr(op->message);
constraint_scope_.Current().Emplace(analyzer_, condition);
- if (condition.same_as(op->condition) && message.same_as(op->message)) {
+ if (condition.same_as(op->condition)) {
return ffi::GetRef<Stmt>(op);
} else {
auto n = this->CopyOnWrite(op);
n->condition = std::move(condition);
- n->message = std::move(message);
return Stmt(n);
}
}
diff --git a/src/arith/ir_visitor_with_analyzer.cc b/src/arith/ir_visitor_with_analyzer.cc
index 02c194b14e..54464ad70f 100644
--- a/src/arith/ir_visitor_with_analyzer.cc
+++ b/src/arith/ir_visitor_with_analyzer.cc
@@ -86,7 +86,6 @@ void IRVisitorWithAnalyzer::VisitStmt_(const AttrStmtNode* op) {
void IRVisitorWithAnalyzer::VisitStmt_(const AssertStmtNode* op) {
this->VisitExpr(op->condition);
- this->VisitExpr(op->message);
constraint_scope_.Current().Emplace(&analyzer_, op->condition);
}
diff --git a/src/relax/op/tensor/inspect.cc b/src/relax/op/tensor/inspect.cc
index bd0f963c49..0130ee1063 100644
--- a/src/relax/op/tensor/inspect.cc
+++ b/src/relax/op/tensor/inspect.cc
@@ -141,7 +141,7 @@ Expr LegalizeTensorDtypeCode(const BlockBuilder& bb, const Call& call) {
Expr arg = call->args[0];
tir::PrimFunc getter =
- GetDLTensorField(tir::builtin::TVMStructFieldKind::kArrTypeCode, field_dtype);
+ GetDLTensorField(tir::builtin::TVMStructFieldKind::kDLTensorTypeCode, field_dtype);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_code");
return Call(gvar_getter, {arg});
@@ -179,7 +179,7 @@ Expr LegalizeTensorDtypeBits(const BlockBuilder& bb, const Call& call) {
Expr arg = call->args[0];
tir::PrimFunc getter =
- GetDLTensorField(tir::builtin::TVMStructFieldKind::kArrTypeBits, field_dtype);
+ GetDLTensorField(tir::builtin::TVMStructFieldKind::kDLTensorTypeBits, field_dtype);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_bits");
return Call(gvar_getter, {arg});
@@ -217,7 +217,7 @@ Expr LegalizeTensorDtypeLanes(const BlockBuilder& bb, const Call& call) {
Expr arg = call->args[0];
tir::PrimFunc getter =
- GetDLTensorField(tir::builtin::TVMStructFieldKind::kArrTypeLanes, field_dtype);
+ GetDLTensorField(tir::builtin::TVMStructFieldKind::kDLTensorTypeLanes, field_dtype);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_dtype_lanes");
return Call(gvar_getter, {arg});
@@ -254,7 +254,8 @@ Expr LegalizeTensorNDim(const BlockBuilder& bb, const Call& call) {
auto field_dtype = Downcast<PrimStructInfo>(call->struct_info_)->dtype;
Expr arg = call->args[0];
- tir::PrimFunc getter = GetDLTensorField(tir::builtin::TVMStructFieldKind::kArrNDim, field_dtype);
+ tir::PrimFunc getter =
+ GetDLTensorField(tir::builtin::TVMStructFieldKind::kDLTensorNDim, field_dtype);
GlobalVar gvar_getter = bb->AddFunction(getter, "_get_tensor_ndim");
return Call(gvar_getter, {arg});
@@ -312,24 +313,25 @@ Expr LegalizeTensorShape(const BlockBuilder& bb, const Call& call) {
shape_buffer->data,
tir::Call(DataType::Handle(), tir::builtin::tvm_struct_get(),
{dlpack_handle, IntImm(DataType::Int(32), 0),
- IntImm(DataType::Int(32), tir::builtin::TVMStructFieldKind::kArrShape)}),
+ IntImm(DataType::Int(32), tir::builtin::TVMStructFieldKind::kDLTensorShape)}),
body);
body = tir::SeqStmt(
{tir::AssertStmt(
- axis < tvm::cast(axis->dtype, ndim),
- tir::StringImm("Specified axis may not be larger than the tensor's dimensionality")),
+ axis < tvm::cast(axis->dtype, ndim), tir::StringImm("RuntimeError"),
+ {tir::StringImm("Specified axis may not be larger than the tensor's dimensionality")}),
body});
body = tir::LetStmt(
ndim,
tir::Call(ndim->dtype, tir::builtin::tvm_struct_get(),
{dlpack_handle, IntImm(DataType::Int(32), 0),
- IntImm(DataType::Int(32), tir::builtin::TVMStructFieldKind::kArrNDim)}),
+ IntImm(DataType::Int(32), tir::builtin::TVMStructFieldKind::kDLTensorNDim)}),
body);
- body = tir::SeqStmt(
- {tir::AssertStmt(0 <= axis, tir::StringImm("Specified axis may not be negative")), body});
+ body = tir::SeqStmt({tir::AssertStmt(0 <= axis, tir::StringImm("RuntimeError"),
+ {tir::StringImm("Specified axis may not be negative")}),
+ body});
DictAttrs attrs({{"tir.is_scheduled", true}, {"tir.is_host", true}});
diff --git a/src/s_tir/analysis/estimate_flops.cc b/src/s_tir/analysis/estimate_flops.cc
index 0f106f5854..c7c9eb2abb 100644
--- a/src/s_tir/analysis/estimate_flops.cc
+++ b/src/s_tir/analysis/estimate_flops.cc
@@ -196,9 +196,6 @@ class FlopEstimator : private ExprFunctor<TResult(const PrimExpr& n)>,
TResult VisitStmt_(const AssertStmtNode* op) override {
TResult result = VisitExpr(op->condition);
- if (op->message.defined()) {
- result += VisitExpr(op->message);
- }
return result;
}
diff --git a/src/s_tir/transform/bound_checker.cc b/src/s_tir/transform/bound_checker.cc
index 62b76fa806..d7692ae19e 100644
--- a/src/s_tir/transform/bound_checker.cc
+++ b/src/s_tir/transform/bound_checker.cc
@@ -97,7 +97,8 @@ class BoundChecker : public StmtExprMutator {
PrimExpr condition = MakeCondition();
if (!condition.as<StringImmNode>()) {
Stmt then_case = ffi::GetRef<Stmt>(op);
- Stmt else_case = AssertStmt(condition, StringImm(error_message_));
+ Stmt else_case =
+ AssertStmt(condition, StringImm("RuntimeError"), {StringImm(error_message_)});
Stmt body = IfThenElse(condition, then_case, else_case);
return body;
}
diff --git a/src/script/ir_builder/tir/frame.cc b/src/script/ir_builder/tir/frame.cc
index e65be1d45f..ede14b3829 100644
--- a/src/script/ir_builder/tir/frame.cc
+++ b/src/script/ir_builder/tir/frame.cc
@@ -129,10 +129,10 @@ void ForFrameNode::ExitWithScope() {
void AssertFrameNode::ExitWithScope() {
TIRFrameNode::ExitWithScope();
if (stmts.empty()) {
- AddToParent(tvm::tir::AssertStmt(condition, message));
+ AddToParent(tvm::tir::AssertStmt(condition, error_kind, message_parts));
} else {
ffi::Array<tvm::tir::Stmt> seq;
- seq.push_back(tvm::tir::AssertStmt(condition, message));
+ seq.push_back(tvm::tir::AssertStmt(condition, error_kind, message_parts));
for (const auto& stmt : stmts) {
seq.push_back(stmt);
}
diff --git a/src/script/ir_builder/tir/ir.cc b/src/script/ir_builder/tir/ir.cc
index a343981355..5579c72abe 100644
--- a/src/script/ir_builder/tir/ir.cc
+++ b/src/script/ir_builder/tir/ir.cc
@@ -446,10 +446,16 @@ ForFrame Grid(ffi::Array<PrimExpr> extents) {
return ForFrame(n);
}
-AssertFrame Assert(PrimExpr condition, ffi::String message) {
+AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
+ ffi::Array<ffi::String> message_parts) {
ObjectPtr<AssertFrameNode> n = ffi::make_object<AssertFrameNode>();
n->condition = condition;
- n->message = tvm::tir::StringImm(message);
+ n->error_kind = tvm::tir::StringImm(error_kind);
+ ffi::Array<tvm::tir::StringImm> parts;
+ for (const auto& p : message_parts) {
+ parts.push_back(tvm::tir::StringImm(p));
+ }
+ n->message_parts = parts;
return AssertFrame(n);
}
diff --git a/src/script/printer/tir/stmt.cc b/src/script/printer/tir/stmt.cc
index 633704c164..ff1bcd00cd 100644
--- a/src/script/printer/tir/stmt.cc
+++ b/src/script/printer/tir/stmt.cc
@@ -131,14 +131,18 @@ TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
});
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
- .set_dispatch<tir::AssertStmt>("",
- [](tir::AssertStmt stmt, AccessPath p, IRDocsifier d) -> Doc {
- ExprDoc cond =
- d->AsDoc<ExprDoc>(stmt->condition, p->Attr("condition"));
- ExprDoc msg =
- d->AsDoc<ExprDoc>(stmt->message, p->Attr("message"));
- return AssertDoc(cond, msg);
- });
+ .set_dispatch<tir::AssertStmt>(
+ "", [](tir::AssertStmt stmt, AccessPath p, IRDocsifier d) -> Doc {
+ ExprDoc cond = d->AsDoc<ExprDoc>(stmt->condition, p->Attr("condition"));
+ // Always emit the canonical tuple form: assert cond, ("Kind", ["part0", "part1", ...])
+ ffi::Array<ExprDoc> parts;
+ auto parts_path = p->Attr("message_parts");
+ for (size_t i = 0; i < stmt->message_parts.size(); ++i) {
+ parts.push_back(d->AsDoc<ExprDoc>(stmt->message_parts[i], parts_path->ArrayItem(i)));
+ }
+ ExprDoc kind_doc = d->AsDoc<ExprDoc>(stmt->error_kind, p->Attr("error_kind"));
+ return AssertDoc(cond, TupleDoc({kind_doc, ListDoc(parts)}));
+ });
TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
.set_dispatch<tir::While>("", [](tir::While stmt, AccessPath p, IRDocsifier d) -> Doc {
diff --git a/src/target/llvm/codegen_cpu.cc b/src/target/llvm/codegen_cpu.cc
index 8a9895f72b..6ffb4fc030 100644
--- a/src/target/llvm/codegen_cpu.cc
+++ b/src/target/llvm/codegen_cpu.cc
@@ -121,6 +121,12 @@ void CodeGenCPU::Init(const std::string& module_name, LLVMTarget* llvm_target,
// void TVMFFIErrorSetRaisedFromCStr(const char *kind, const char* msg);
ftype_tvm_ffi_error_set_raised_by_c_str_ = llvm::FunctionType::get(
t_void_, {llvmGetPointerTo(t_char_, 0), llvmGetPointerTo(t_char_, 0)}, false);
+ // Defined in include/tvm/ffi/c_api.h:
+ // void TVMFFIErrorSetRaisedFromCStrParts(const char* kind, const char** msg_parts, int32_t n);
+ ftype_tvm_ffi_error_set_raised_from_c_str_parts_ = llvm::FunctionType::get(
+ t_void_,
+ {llvmGetPointerTo(t_char_, 0), llvmGetPointerTo(llvmGetPointerTo(t_char_, 0), 0), t_int_},
+ false);
// Defined in include/tvm/runtime/c_backend_api.h:
// int TVMBackendGetFuncFromEnv(void* mod_node, const char* func_name, TVMFunctionHandle* out);
ftype_tvm_get_func_from_env_ = llvm::FunctionType::get(
@@ -158,6 +164,9 @@ void CodeGenCPU::Init(const std::string& module_name, LLVMTarget* llvm_target,
f_tvm_ffi_set_raised_by_c_str_ = llvm::Function::Create(
ftype_tvm_ffi_error_set_raised_by_c_str_, llvm::Function::ExternalLinkage,
"TVMFFIErrorSetRaisedFromCStr", module_.get());
+ f_tvm_ffi_set_raised_from_c_str_parts_ = llvm::Function::Create(
+ ftype_tvm_ffi_error_set_raised_from_c_str_parts_, llvm::Function::ExternalLinkage,
+ "TVMFFIErrorSetRaisedFromCStrParts", module_.get());
f_tvm_get_func_from_env_ =
llvm::Function::Create(ftype_tvm_get_func_from_env_, llvm::Function::ExternalLinkage,
"TVMBackendGetFuncFromEnv", module_.get());
@@ -270,7 +279,7 @@ std::unique_ptr<llvm::Module> CodeGenCPU::Finish() {
CodeGenLLVM::TypedPointer CodeGenCPU::CreateStructRefPtr(DataType t, llvm::Value* buf,
llvm::Value* index, int kind) {
- if (kind < builtin::kArrKindBound_) {
+ if (kind < builtin::kDLTensorKindBound_) {
if (buf->getType() == t_void_p_) {
buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_tvm_array_, 0));
} else {
@@ -278,64 +287,64 @@ CodeGenLLVM::TypedPointer CodeGenCPU::CreateStructRefPtr(DataType t, llvm::Value
}
}
switch (kind) {
- case builtin::kArrAddr: {
+ case builtin::kDLTensorAddr: {
return TypedPointer(t_tvm_array_, builder_->CreateInBoundsGEP(t_tvm_array_, buf, index));
}
- case builtin::kArrData: {
+ case builtin::kDLTensorData: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(0);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(0)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrShape: {
+ case builtin::kDLTensorShape: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(4);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(4)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrStrides: {
+ case builtin::kDLTensorStrides: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(5);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(5)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrNDim: {
+ case builtin::kDLTensorNDim: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(2);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(2)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrTypeCode: {
+ case builtin::kDLTensorTypeCode: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(0);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(0)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrTypeBits: {
+ case builtin::kDLTensorTypeBits: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(1);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(1)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrTypeLanes: {
+ case builtin::kDLTensorTypeLanes: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(3)->getStructElementType(2);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(3), ConstInt32(2)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrByteOffset: {
+ case builtin::kDLTensorByteOffset: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(6);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(6)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrDeviceId: {
+ case builtin::kDLTensorDeviceId: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(1)->getStructElementType(1);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(1), ConstInt32(1)});
return TypedPointer(member_type, member_addr);
}
- case builtin::kArrDeviceType: {
+ case builtin::kDLTensorDeviceType: {
llvm::Type* member_type = t_tvm_array_->getStructElementType(1)->getStructElementType(0);
llvm::Value* member_addr =
builder_->CreateInBoundsGEP(t_tvm_array_, buf, {index, ConstInt32(1), ConstInt32(0)});
@@ -373,6 +382,11 @@ CodeGenLLVM::TypedPointer CodeGenCPU::CreateStructRefPtr(DataType t, llvm::Value
LOG(DEBUG) << "DataType " << t << " cannot be stored into a TVMFFIAny's value field";
}
}
+ case builtin::kInt64ArrayElem: {
+ buf = builder_->CreatePointerCast(buf, llvmGetPointerTo(t_int64_, 0));
+ llvm::Value* elem_addr = builder_->CreateInBoundsGEP(t_int64_, buf, index);
+ return TypedPointer(t_int64_, elem_addr);
+ }
default:
TVM_FFI_THROW(InternalError) << "unknown field code";
}
@@ -890,6 +904,82 @@ llvm::Value* CodeGenCPU::RuntimeTVMFFIErrorSetRaisedFromCStr() {
if (f_tvm_ffi_set_raised_by_c_str_ != nullptr) return f_tvm_ffi_set_raised_by_c_str_;
return GetContextPtr(gv_tvm_ffi_set_last_error_c_str_);
}
+llvm::Value* CodeGenCPU::RuntimeTVMFFIErrorSetRaisedFromCStrParts() {
+ if (f_tvm_ffi_set_raised_from_c_str_parts_ != nullptr) {
+ return f_tvm_ffi_set_raised_from_c_str_parts_;
+ }
+ // fallback: create external linkage
+ f_tvm_ffi_set_raised_from_c_str_parts_ = llvm::Function::Create(
+ ftype_tvm_ffi_error_set_raised_from_c_str_parts_, llvm::Function::ExternalLinkage,
+ "TVMFFIErrorSetRaisedFromCStrParts", module_.get());
+ return f_tvm_ffi_set_raised_from_c_str_parts_;
+}
+/*!
+ * \brief Get or create a noinline LLVM helper function for raising errors from TIR assertions.
+ *
+ * Creates an internal LLVM function `__tvm_set_raised_<max_n>` that:
+ * 1. Stack-allocates a `const char**` array of size `max_n`
+ * 2. Stores each message part argument into the array
+ * 3. Calls `TVMFFIErrorSetRaisedFromCStrParts(kind, parts, n)`
+ *
+ * The function is marked `noinline` to keep error-raising code out of the hot path.
+ * Assert codegen emits a conditional branch to a block that calls this helper,
+ * so the error-handling logic does not bloat the main function body.
+ *
+ * Results are cached in `set_raised_helpers_` keyed by `max_n`, so each distinct
+ * bucket size (6, 12, 18, ...) is created at most once per module.
+ *
+ * \param max_n The maximum number of message part slots (rounded up to a bucket size).
+ * \return The LLVM function for raising errors with up to `max_n` message parts.
+ */
+llvm::Function* CodeGenCPU::GetOrCreateSetRaisedHelper(int max_n) {
+ auto it = set_raised_helpers_.find(max_n);
+ if (it != set_raised_helpers_.end()) return it->second;
+
+ llvm::LLVMContext* ctx = llvm_target_->GetContext();
+ // Build function type: void(i8* kind, i32 n, i8* msg0, i8* msg1, ..., i8* msg_{max_n-1})
+ std::vector<llvm::Type*> param_types;
+ param_types.push_back(llvmGetPointerTo(t_char_, 0)); // kind
+ param_types.push_back(t_int_); // n
+ for (int i = 0; i < max_n; ++i) {
+ param_types.push_back(llvmGetPointerTo(t_char_, 0)); // msg_i
+ }
+ auto* ftype = llvm::FunctionType::get(t_void_, param_types, false);
+ std::string fname = "__tvm_set_raised_" + std::to_string(max_n);
+ auto* func = llvm::Function::Create(ftype, llvm::Function::InternalLinkage, fname, module_.get());
+ func->addFnAttr(llvm::Attribute::NoInline);
+
+ // Create the function body
+ auto* entry = llvm::BasicBlock::Create(*ctx, "entry", func);
+ llvm::IRBuilder<> b(entry);
+
+ // Alloca array of char* on the stack
+ auto* arr_type = llvm::ArrayType::get(llvmGetPointerTo(t_char_, 0), max_n);
+ auto* parts_alloca = b.CreateAlloca(arr_type);
+
+ // Store each msg_i into parts[i]
+ auto arg_it = func->arg_begin();
+ llvm::Value* kind_arg = &*arg_it++;
+ llvm::Value* n_arg = &*arg_it++;
+ for (int i = 0; i < max_n; ++i) {
+ llvm::Value* msg_i = &*arg_it++;
+ llvm::Value* gep = b.CreateInBoundsGEP(arr_type, parts_alloca, {ConstInt32(0), ConstInt32(i)});
+ b.CreateStore(msg_i, gep);
+ }
+
+ // Cast alloca to const char**
+ llvm::Value* parts_ptr =
+ b.CreateInBoundsGEP(arr_type, parts_alloca, {ConstInt32(0), ConstInt32(0)});
+
+ // Call TVMFFIErrorSetRaisedFromCStrParts(kind, parts, n)
+ auto callee = llvm::FunctionCallee(ftype_tvm_ffi_error_set_raised_from_c_str_parts_,
+ RuntimeTVMFFIErrorSetRaisedFromCStrParts());
+ b.CreateCall(callee, {kind_arg, parts_ptr, n_arg});
+ b.CreateRetVoid();
+
+ set_raised_helpers_[max_n] = func;
+ return func;
+}
llvm::Value* CodeGenCPU::RuntimeTVMParallelLaunch() {
if (f_tvm_parallel_launch_ != nullptr) return f_tvm_parallel_launch_;
return GetContextPtr(gv_tvm_parallel_launch_);
@@ -940,7 +1030,7 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
int kind = op->args[2].as<IntImm>().value()->value;
TypedPointer ref =
CreateStructRefPtr(op->dtype, MakeValue(op->args[0]), MakeValue(op->args[1]), kind);
- if (kind == builtin::kArrAddr) {
+ if (kind == builtin::kDLTensorAddr) {
return builder_->CreatePointerCast(ref.addr, t_void_p_);
}
@@ -957,7 +1047,7 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
llvm::Value* value = MakeValue(op->args[3]);
TypedPointer ref = CreateStructRefPtr(op->args[3].dtype(), MakeValue(op->args[0]),
MakeValue(op->args[1]), kind);
- TVM_FFI_ICHECK(kind != builtin::kArrAddr);
+ TVM_FFI_ICHECK(kind != builtin::kDLTensorAddr);
if (value->getType()->isPointerTy()) {
value = builder_->CreatePointerCast(value, ref.type);
}
@@ -1002,12 +1092,6 @@ llvm::Value* CodeGenCPU::CreateIntrinsic(const CallNode* op) {
void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
EmitDebugLocation(op);
llvm::Value* cond = MakeValue(op->condition);
- std::ostringstream os;
- os << "Assert fail: " << op->condition;
- if (op->message.as<StringImmNode>()) {
- os << ", " << op->message.as<StringImmNode>()->value;
- }
- llvm::Value* msg = GetConstString(os.str());
llvm::LLVMContext* ctx = llvm_target_->GetContext();
auto* fail_block = llvm::BasicBlock::Create(*ctx, "assert_fail", function_);
auto* end_block = llvm::BasicBlock::Create(*ctx, "assert_end", function_);
@@ -1015,9 +1099,26 @@ void CodeGenCPU::VisitStmt_(const AssertStmtNode* op) {
// fail condition.
builder_->SetInsertPoint(fail_block);
- auto err_callee = llvm::FunctionCallee(ftype_tvm_ffi_error_set_raised_by_c_str_,
- RuntimeTVMFFIErrorSetRaisedFromCStr());
- builder_->CreateCall(err_callee, {GetConstString("RuntimeError"), msg});
+ int num_parts = static_cast<int>(op->message_parts.size());
+ // Round up to the nearest helper bucket: 6, 12, 18, ...
+ int max_n = ((num_parts + 5) / 6) * 6;
+ if (max_n == 0) max_n = 6;
+
+ llvm::Function* helper = GetOrCreateSetRaisedHelper(max_n);
+ // Build arguments: kind, n, msg0, msg1, ..., msg_{max_n-1}
+ std::vector<llvm::Value*> call_args;
+ call_args.push_back(GetConstString(op->error_kind->value));
+ call_args.push_back(ConstInt32(num_parts));
+ for (int i = 0; i < max_n; ++i) {
+ if (i < num_parts) {
+ call_args.push_back(GetConstString(op->message_parts[i]->value));
+ } else {
+ // Pad with null
+ call_args.push_back(llvm::ConstantPointerNull::get(
+ llvm::cast<llvm::PointerType>(llvmGetPointerTo(t_char_, 0))));
+ }
+ }
+ builder_->CreateCall(helper, call_args);
builder_->CreateRet(ConstInt32(-1));
// otherwise set it to be new end.
builder_->SetInsertPoint(end_block);
diff --git a/src/target/llvm/codegen_cpu.h b/src/target/llvm/codegen_cpu.h
index d5401b9662..6096cc1405 100644
--- a/src/target/llvm/codegen_cpu.h
+++ b/src/target/llvm/codegen_cpu.h
@@ -94,6 +94,7 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::FunctionType* ftype_tvm_ffi_func_call_{nullptr};
llvm::FunctionType* ftype_tvm_get_func_from_env_{nullptr};
llvm::FunctionType* ftype_tvm_ffi_error_set_raised_by_c_str_{nullptr};
+ llvm::FunctionType* ftype_tvm_ffi_error_set_raised_from_c_str_parts_{nullptr};
llvm::FunctionType* ftype_tvm_parallel_launch_{nullptr};
llvm::FunctionType* ftype_tvm_parallel_barrier_{nullptr};
llvm::FunctionType* ftype_tvm_register_system_symbol_{nullptr};
@@ -118,6 +119,9 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::Value* RuntimeTVMFFIFunctionCall();
llvm::Value* RuntimeTVMGetFuncFromEnv();
llvm::Value* RuntimeTVMFFIErrorSetRaisedFromCStr();
+ llvm::Value* RuntimeTVMFFIErrorSetRaisedFromCStrParts();
+ // Create a temp function to simplify error raising.
+ llvm::Function* GetOrCreateSetRaisedHelper(int max_n);
llvm::Value* RuntimeTVMParallelLaunch();
llvm::Value* RuntimeTVMParallelBarrier();
llvm::Value* CreateStaticHandle();
@@ -166,11 +170,14 @@ class CodeGenCPU : public CodeGenLLVM {
llvm::Function* f_tvm_ffi_func_call_{nullptr};
llvm::Function* f_tvm_get_func_from_env_{nullptr};
llvm::Function* f_tvm_ffi_set_raised_by_c_str_{nullptr};
+ llvm::Function* f_tvm_ffi_set_raised_from_c_str_parts_{nullptr};
llvm::Function* f_tvm_parallel_launch_{nullptr};
llvm::Function* f_tvm_parallel_barrier_{nullptr};
llvm::Function* f_tvm_register_system_symbol_{nullptr};
// Current parallel environment scope.
ParallelEnv parallel_env_;
+ // cached noinline helper functions for SetRaisedFromCStrParts
+ std::unordered_map<int, llvm::Function*> set_raised_helpers_;
// global to packed function handle
std::unordered_map<std::string, llvm::GlobalVariable*> func_handle_map_;
// List of symbols to be exported to TVM system lib.
diff --git a/src/target/llvm/codegen_llvm.cc b/src/target/llvm/codegen_llvm.cc
index 6daa2d49a6..388971a843 100644
--- a/src/target/llvm/codegen_llvm.cc
+++ b/src/target/llvm/codegen_llvm.cc
@@ -1428,7 +1428,13 @@ llvm::Value* CodeGenLLVM::CreateIntrinsic(const CallNode* op) {
return post_dummy;
} else if (op->op.same_as(builtin::reinterpret())) {
llvm::Type* target = DTypeToLLVMType(op->dtype);
- return builder_->CreateBitCast(MakeValue(op->args[0]), target);
+ llvm::Value* value = MakeValue(op->args[0]);
+ if (value->getType()->isPointerTy() && target->isIntegerTy()) {
+ return builder_->CreatePtrToInt(value, target);
+ } else if (value->getType()->isIntegerTy() && target->isPointerTy()) {
+ return builder_->CreateIntToPtr(value, target);
+ }
+ return builder_->CreateBitCast(value, target);
} else if (op->op.same_as(builtin::isnan())) {
// TODO(hgt312): set fast math flag
llvm::Value* a = MakeValue(op->args[0]);
diff --git a/src/target/source/codegen_c.cc b/src/target/source/codegen_c.cc
index 95b1260e45..322bbeea40 100644
--- a/src/target/source/codegen_c.cc
+++ b/src/target/source/codegen_c.cc
@@ -280,12 +280,12 @@ std::string CodeGenC::GetBufferRef(DataType t, const BufferNode* buffer, PrimExp
// Print a reference expression to a buffer.
std::string CodeGenC::GetStructRef(DataType t, const PrimExpr& buffer, const PrimExpr& index,
int kind) {
- if (kind < builtin::kArrKindBound_) {
+ if (kind < builtin::kDLTensorKindBound_) {
std::ostringstream os;
os << "(((DLTensor*)";
this->PrintExpr(buffer, os);
os << ")";
- if (kind == builtin::kArrAddr) {
+ if (kind == builtin::kDLTensorAddr) {
os << " + ";
this->PrintExpr(index, os);
os << ")";
@@ -296,34 +296,34 @@ std::string CodeGenC::GetStructRef(DataType t, const PrimExpr& buffer, const Pri
os << "].";
// other case: get fields.
switch (kind) {
- case builtin::kArrData:
+ case builtin::kDLTensorData:
os << "data";
break;
- case builtin::kArrShape:
+ case builtin::kDLTensorShape:
os << "shape";
break;
- case builtin::kArrStrides:
+ case builtin::kDLTensorStrides:
os << "strides";
break;
- case builtin::kArrNDim:
+ case builtin::kDLTensorNDim:
os << "ndim";
break;
- case builtin::kArrTypeCode:
+ case builtin::kDLTensorTypeCode:
os << "dtype.code";
break;
- case builtin::kArrTypeBits:
+ case builtin::kDLTensorTypeBits:
os << "dtype.bits";
break;
- case builtin::kArrByteOffset:
+ case builtin::kDLTensorByteOffset:
os << "byte_offset";
break;
- case builtin::kArrTypeLanes:
+ case builtin::kDLTensorTypeLanes:
os << "dtype.lanes";
break;
- case builtin::kArrDeviceId:
+ case builtin::kDLTensorDeviceId:
os << "device.device_id";
break;
- case builtin::kArrDeviceType:
+ case builtin::kDLTensorDeviceType:
os << "device.device_type";
break;
default:
@@ -359,6 +359,14 @@ std::string CodeGenC::GetStructRef(DataType t, const PrimExpr& buffer, const Pri
}
os << ")";
return os.str();
+ } else if (kind == builtin::kInt64ArrayElem) {
+ std::ostringstream os;
+ os << "(((int64_t*)";
+ this->PrintExpr(buffer, os);
+ os << ")[";
+ this->PrintExpr(index, os);
+ os << "])";
+ return os.str();
} else {
TVM_FFI_THROW(RuntimeError) << "Unsupported type index: " << kind;
}
@@ -1079,12 +1087,77 @@ void CodeGenC::VisitStmt_(const AttrStmtNode* op) {
this->PrintStmt(op->body);
}
+void CodeGenC::PrintEscapedCString(const std::string& str, std::ostream& os) {
+ os << "\"";
+ for (unsigned char c : str) {
+ switch (c) {
+ case '"':
+ os << "\\\"";
+ break;
+ case '\\':
+ os << "\\\\";
+ break;
+ case '\n':
+ os << "\\n";
+ break;
+ case '\t':
+ os << "\\t";
+ break;
+ case '\r':
+ os << "\\r";
+ break;
+ case '\a':
+ os << "\\a";
+ break;
+ case '\b':
+ os << "\\b";
+ break;
+ case '\f':
+ os << "\\f";
+ break;
+ case '\v':
+ os << "\\v";
+ break;
+ case '\0':
+ os << "\\0";
+ break;
+ default:
+ if (c < 0x20 || c >= 0x7f) {
+ // Non-printable: emit as hex escape
+ os << "\\x" << std::setfill('0') << std::setw(2) << std::hex << static_cast<int>(c)
+ << std::dec;
+ } else {
+ os << static_cast<char>(c);
+ }
+ break;
+ }
+ }
+ os << "\"";
+}
+
void CodeGenC::VisitStmt_(const AssertStmtNode* op) {
std::string cond = PrintExpr(op->condition);
PrintIndent();
- if (const auto* str = op->message.as<StringImmNode>()) {
- // GLOG style check
- stream << "TVM_FFI_ICHECK(" << cond << ") << \"" << str->value << "\";\n";
+ int num_parts = static_cast<int>(op->message_parts.size());
+ if (num_parts > 0) {
+ stream << "if (!(" << cond << ")) {\n";
+ int assert_if_scope = this->BeginScope();
+ PrintIndent();
+ stream << "const char* __tvm_assert_parts[" << num_parts << "] = {";
+ for (int i = 0; i < num_parts; ++i) {
+ if (i > 0) stream << ", ";
+ PrintEscapedCString(op->message_parts[i]->value, stream);
+ }
+ stream << "};\n";
+ PrintIndent();
+ stream << "TVMFFIErrorSetRaisedFromCStrParts(";
+ PrintEscapedCString(op->error_kind->value, stream);
+ stream << ", __tvm_assert_parts, " << num_parts << ");\n";
+ PrintIndent();
+ stream << "return -1;\n";
+ this->EndScope(assert_if_scope);
+ PrintIndent();
+ stream << "}\n";
} else {
stream << "assert(" << cond << ");\n";
}
@@ -1179,10 +1252,10 @@ void CodeGenC::VisitStmt_(const EvaluateNode* op) {
<< " = 0;\n";
}
- if (kind == builtin::kArrStrides) {
+ if (kind == builtin::kDLTensorStrides) {
// cast void* to int64_t*
cast = call->args[3]->dtype.is_handle() ? "(int64_t*)" : "";
- } else if (kind == builtin::kArrDeviceType) {
+ } else if (kind == builtin::kDLTensorDeviceType) {
// cast int to enum
cast = "(DLDeviceType)";
}
diff --git a/src/target/source/codegen_c.h b/src/target/source/codegen_c.h
index 791fe4a015..d699c4a791 100644
--- a/src/target/source/codegen_c.h
+++ b/src/target/source/codegen_c.h
@@ -234,6 +234,8 @@ class CodeGenC : public ExprFunctor<void(const PrimExpr&, std::ostream&)>,
}
protected:
+ /*! \brief Print a C string literal with proper escaping of special chars. */
+ void PrintEscapedCString(const std::string& str, std::ostream& os);
// Print reference to struct location
std::string GetStructRef(DataType t, const PrimExpr& buffer, const PrimExpr& index, int kind);
// Print reference to a buffer as type t in index.
diff --git a/src/target/source/codegen_c_host.cc b/src/target/source/codegen_c_host.cc
index ca6b71b4d8..8e7fb11b64 100644
--- a/src/target/source/codegen_c_host.cc
+++ b/src/target/source/codegen_c_host.cc
@@ -323,9 +323,18 @@ void CodeGenCHost::VisitStmt_(const AssertStmtNode* op) { // NOLINT(*)
PrintIndent();
stream << "if (!(" << cond << ")) {\n";
int assert_if_scope = this->BeginScope();
+ int num_parts = static_cast<int>(op->message_parts.size());
PrintIndent();
- stream << "TVMFFIErrorSetRaisedFromCStr(\"RuntimeError\", \""
- << op->message.as<StringImmNode>()->value << "\", NULL);\n";
+ stream << "const char* __tvm_assert_parts[" << num_parts << "] = {";
+ for (int i = 0; i < num_parts; ++i) {
+ if (i > 0) stream << ", ";
+ PrintEscapedCString(op->message_parts[i]->value, stream);
+ }
+ stream << "};\n";
+ PrintIndent();
+ stream << "TVMFFIErrorSetRaisedFromCStrParts(";
+ PrintEscapedCString(op->error_kind->value, stream);
+ stream << ", __tvm_assert_parts, " << num_parts << ");\n";
PrintIndent();
stream << "return -1;\n";
this->EndScope(assert_if_scope);
@@ -358,7 +367,7 @@ inline void CodeGenCHost::PrintTernaryCondExpr(const T* op, const char* compare,
ffi::Module BuildCHost(IRModule mod, Target target) {
bool output_ssa = false;
- bool emit_asserts = false;
+ bool emit_asserts = true;
bool emit_fwd_func_decl = true;
std::unordered_set<std::string> devices;
diff --git a/src/tir/ir/buffer.cc b/src/tir/ir/buffer.cc
index b569a1cb07..c1ea74f3d2 100644
--- a/src/tir/ir/buffer.cc
+++ b/src/tir/ir/buffer.cc
@@ -515,7 +515,7 @@ Buffer Buffer::MakeSlice(ffi::Array<PrimExpr> begins, ffi::Array<PrimExpr> exten
// Buffer must be constructed with a singular element offset which means there is no
// support for n-dimensional buffers where n > 1. Insert sentinel value for
- // ArgBinder::BindBuffer to state that any usage of element offset is invalid
+ // TVMFFIABIBuilder::BindBuffer to state that any usage of element offset is invalid
// in this case. This allows for construction of a Buffer with multiple element offsets
// but disallows any usage of those element offsets. See PR #10816 for discussion on
// supporting multiple element offsets in TIR Buffer.
diff --git a/src/tir/ir/stmt.cc b/src/tir/ir/stmt.cc
index 5eaff0c314..f08b0b0ba9 100644
--- a/src/tir/ir/stmt.cc
+++ b/src/tir/ir/stmt.cc
@@ -104,25 +104,27 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
// AssertStmt
-AssertStmt::AssertStmt(PrimExpr condition, PrimExpr message, Span span) {
+AssertStmt::AssertStmt(PrimExpr condition, StringImm error_kind,
+ ffi::Array<StringImm> message_parts, Span span) {
TVM_FFI_ICHECK(condition.defined());
- TVM_FFI_ICHECK(condition.dtype().is_bool())
+ TVM_FFI_ICHECK(condition.dtype().is_predicate_dtype())
<< "AssertStmt should have boolean condition, "
<< "but received " << condition << " with dtype " << condition.dtype();
- TVM_FFI_CHECK(message.dtype() == DataType::Int(32) || message.as<StringImmNode>(), TypeError)
- << "AssertStmt message must be an int or string:" << message << "\n";
+ TVM_FFI_ICHECK(error_kind.defined());
ObjectPtr<AssertStmtNode> node = ffi::make_object<AssertStmtNode>();
node->condition = std::move(condition);
- node->message = std::move(message);
+ node->error_kind = std::move(error_kind);
+ node->message_parts = std::move(message_parts);
node->span = std::move(span);
data_ = std::move(node);
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("tir.AssertStmt", [](PrimExpr condition, StringImm message, Span span) {
- return AssertStmt(condition, message, span);
+ refl::GlobalDef().def("tir.AssertStmt", [](PrimExpr condition, StringImm error_kind,
+ ffi::Array<StringImm> message_parts, Span span) {
+ return AssertStmt(condition, error_kind, message_parts, span);
});
}
diff --git a/src/tir/ir/stmt_functor.cc b/src/tir/ir/stmt_functor.cc
index 659056a3a2..401b7e0c04 100644
--- a/src/tir/ir/stmt_functor.cc
+++ b/src/tir/ir/stmt_functor.cc
@@ -80,7 +80,8 @@ void StmtVisitor::VisitStmt_(const IfThenElseNode* op) {
void StmtVisitor::VisitStmt_(const AssertStmtNode* op) {
this->VisitExpr(op->condition);
- this->VisitExpr(op->message);
+ this->VisitExpr(op->error_kind);
+ VisitArray(op->message_parts, [this](const StringImm& e) { this->VisitExpr(e); });
}
void StmtVisitor::VisitStmt_(const SeqStmtNode* op) {
@@ -396,14 +397,19 @@ Stmt StmtMutator::VisitSeqStmt_(const SeqStmtNode* op, bool flatten_before_visit
Stmt StmtMutator::VisitStmt_(const AssertStmtNode* op) {
PrimExpr condition = this->VisitExpr(op->condition);
- PrimExpr message = this->VisitExpr(op->message);
+ PrimExpr error_kind = this->VisitExpr(op->error_kind);
+ ffi::Array<StringImm> message_parts = Internal::MutateArray(
+ this, op->message_parts,
+ [this](const StringImm& e) { return Downcast<StringImm>(this->VisitExpr(e)); });
- if (condition.same_as(op->condition) && message.same_as(op->message)) {
+ if (condition.same_as(op->condition) && error_kind.same_as(op->error_kind) &&
+ message_parts.same_as(op->message_parts)) {
return ffi::GetRef<Stmt>(op);
} else {
auto n = CopyOnWrite(op);
n->condition = std::move(condition);
- n->message = std::move(message);
+ n->error_kind = Downcast<StringImm>(std::move(error_kind));
+ n->message_parts = std::move(message_parts);
return Stmt(n);
}
}
diff --git a/src/tir/ir/tir_visitor_with_path.cc b/src/tir/ir/tir_visitor_with_path.cc
index e126ebf986..eca926f547 100644
--- a/src/tir/ir/tir_visitor_with_path.cc
+++ b/src/tir/ir/tir_visitor_with_path.cc
@@ -252,7 +252,8 @@ void TIRVisitorWithPath::VisitStmt_(const IfThenElseNode* op, AccessPath path) {
void TIRVisitorWithPath::VisitStmt_(const AssertStmtNode* op, AccessPath path) {
Visit(op->condition, path->Attr("condition"));
- Visit(op->message, path->Attr("message"));
+ Visit(op->error_kind, path->Attr("error_kind"));
+ Visit(op->message_parts, path->Attr("message_parts"));
}
void TIRVisitorWithPath::VisitStmt_(const SeqStmtNode* op, AccessPath path) {
diff --git a/src/tir/transform/arg_binder.cc b/src/tir/transform/arg_binder.cc
deleted file mode 100644
index bab4cdb3da..0000000000
--- a/src/tir/transform/arg_binder.cc
+++ /dev/null
@@ -1,329 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file arg_binder.cc
- * \brief Helper utility to match and bind arguments.
- */
-#include "arg_binder.h"
-
-#include <tvm/runtime/device_api.h>
-#include <tvm/tir/builtin.h>
-#include <tvm/tir/expr.h>
-#include <tvm/tir/op.h>
-
-#include "ir_utils.h"
-
-namespace tvm {
-namespace tir {
-
-void BinderAddAssert(arith::Analyzer* ana, PrimExpr cond, const std::string& arg_name,
- std::vector<Stmt>* asserts) {
- PrimExpr scond = ana->Simplify(cond);
- if (is_zero(scond)) {
- TVM_FFI_THROW(InternalError) << "Bind have an unmet assertion: " << cond << ", "
- << " on argument " << arg_name;
- }
- if (!is_one(scond)) {
- std::ostringstream os;
- os << "Argument " << arg_name << " has an unsatisfied constraint: " << cond;
- asserts->emplace_back(AssertStmt(scond, tvm::tir::StringImm(os.str())));
- }
-}
-
-bool ArgBinder::Bind_(const PrimExpr& arg, const PrimExpr& value, const std::string& arg_name,
- bool with_lets) {
- TVM_FFI_ICHECK_EQ(arg.dtype(), value.dtype());
- if (const VarNode* v = arg.as<VarNode>()) {
- auto it = def_map_->find(v);
- if (it == def_map_->end()) {
- Var v_arg = Downcast<Var>(arg);
- defs_.emplace_back(v_arg);
- if (with_lets) {
- (*def_map_)[v] = arg;
- init_nest_.emplace_back(LetStmt(v_arg, value, Evaluate(0)));
- } else {
- (*def_map_)[v] = value;
- }
- return true;
- } else {
- BinderAddAssert(&analyzer_, it->second == value, arg_name, &asserts_);
- }
- } else {
- BinderAddAssert(&analyzer_, arg == value, arg_name, &asserts_);
- }
- return false;
-}
-
-void ArgBinder::Bind(const PrimExpr& arg, const PrimExpr& value, const std::string& arg_name,
- bool with_let) {
- Bind_(arg, value, arg_name, with_let);
-}
-
-void ArgBinder::BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
- const std::string& arg_name) {
- TVM_FFI_ICHECK_EQ(arg.size(), value.size()) << "Argument " << arg_name << " array size mismatch";
- for (size_t i = 0; i < arg.size(); ++i) {
- std::ostringstream os;
- os << arg_name << "[" << i << "]";
- this->Bind(arg[i], value[i], os.str());
- }
-}
-
-void ArgBinder::BindBuffer(const Buffer& arg, const Buffer& value, const std::string& arg_name,
- bool fuzzy_match) {
- TVM_FFI_ICHECK_EQ(arg.scope(), value.scope())
- << "Argument " << arg_name << " Buffer bind scope mismatch";
- TVM_FFI_ICHECK_EQ(arg->dtype, value->dtype)
- << "Argument " << arg_name << " Buffer bind data type mismatch";
- if (value->data_alignment % arg->data_alignment != 0) {
- LOG(WARNING) << "Trying to bind buffer to another one with lower alignment requirement "
- << " required alignment=" << arg->data_alignment
- << ", provided alignment=" << value->data_alignment;
- }
-
- if (value->elem_offset.defined()) {
- // bind pointer and offset.
- if (is_zero(arg->elem_offset)) {
- TVM_FFI_ICHECK(is_zero(value->elem_offset))
- << "Trying to bind a Buffer with offset into one without offset "
- << " required elem_offset=" << arg->elem_offset
- << ", provided elem_offset=" << value->elem_offset;
- }
-
- this->Bind(arg->data, value->data, arg_name + ".data");
- if (Bind_(arg->elem_offset, value->elem_offset, arg_name + ".elem_offset", false)) {
- if (arg->offset_factor > 1) {
- PrimExpr offset = value->elem_offset;
- PrimExpr factor = make_const(offset.dtype(), arg->offset_factor);
- PrimExpr zero = make_zero(offset.dtype());
- BinderAddAssert(&analyzer_, truncmod(offset, factor) == zero, arg_name + ".elem_offset",
- &asserts_);
- }
- }
- }
-
- if (arg->shape.size() < value->shape.size()) {
- TVM_FFI_ICHECK(fuzzy_match) << "Argument " << arg_name << " size mismatch";
- size_t diff = value->shape.size() - arg->shape.size();
- for (size_t i = 0; i < diff; ++i) {
- TVM_FFI_ICHECK(is_one(analyzer_.Simplify(value->shape[i])))
- << "Argument " << arg_name << " shape mismatch" << arg->shape << " vs " << value->shape;
- }
- for (size_t i = 0; i < arg->shape.size(); ++i) {
- std::ostringstream os;
- os << arg_name << ".shape[" << i << "]";
- this->Bind(arg->shape[i], value->shape[i + diff], os.str());
- }
- if (value->strides.size() != 0) {
- TVM_FFI_ICHECK_EQ(arg->strides.size(), arg->shape.size());
- TVM_FFI_ICHECK_EQ(value->strides.size(), value->shape.size());
- for (size_t i = 0; i < arg->strides.size(); ++i) {
- std::ostringstream os;
- os << arg_name << ".strides[" << i << "]";
- this->Bind(arg->strides[i], value->strides[i + diff], os.str());
- }
- }
- } else {
- this->BindArray(arg->shape, value->shape, arg_name + ".shape");
- this->BindArray(arg->strides, value->strides, arg_name + ".strides");
- }
-}
-
-inline PrimExpr TVMArrayGet(DataType t, Var arr, builtin::TVMStructFieldKind kind) {
- return TVMStructGet(t, arr, 0, kind);
-}
-
-void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
- const PrimExpr& device_id, const Var& handle,
- const std::string& arg_name) {
- const DataType tvm_shape_type = DataType::ShapeIndex();
- const DataType tvm_ndim_type = DataType::Int(32);
- const Stmt nop = Evaluate(0);
-
- init_nest_.emplace_back(AssertStmt(
- !Call(DataType::Bool(), builtin::isnullptr(), {handle}),
- tvm::tir::StringImm(arg_name + " is expected to have non-NULL DLTensor* pointer")));
-
- // dimension checks
- PrimExpr v_ndim = TVMArrayGet(tvm_ndim_type, handle, builtin::kArrNDim);
-
- // Helper functions for shape/stride name formatting
- auto shape_handle_name = [&]() { return arg_name + ".shape"; };
- auto stride_handle_name = [&]() { return arg_name + ".strides"; };
- auto array_element_name = [&](const std::string& arr_name, size_t k) {
- std::stringstream ss;
- ss << arr_name << '[' << k << ']';
- return ss.str();
- };
- auto shape_element_name = [&](size_t k) { return array_element_name(shape_handle_name(), k); };
- auto stride_element_name = [&](size_t k) { return array_element_name(stride_handle_name(), k); };
-
- PrimExpr a_ndim = make_const(tvm_ndim_type, static_cast<int64_t>(buffer->shape.size()));
- std::ostringstream ndim_err_msg;
- ndim_err_msg << arg_name << ".ndim is expected to equal " << buffer->shape.size();
- auto msg = tvm::tir::StringImm(ndim_err_msg.str());
- init_nest_.emplace_back(AssertStmt(a_ndim == v_ndim, msg));
- // type checks
- std::ostringstream type_err_msg;
- type_err_msg << arg_name << ".dtype is expected to be " << buffer->dtype;
- PrimExpr cond = (TVMArrayGet(DataType::UInt(8), handle, builtin::kArrTypeCode) ==
- IntImm(DataType::UInt(8), buffer->dtype.code()) &&
- TVMArrayGet(DataType::UInt(8), handle, builtin::kArrTypeBits) ==
- IntImm(DataType::UInt(8), buffer->dtype.bits()) &&
- TVMArrayGet(DataType::UInt(16), handle, builtin::kArrTypeLanes) ==
- IntImm(DataType::UInt(16), buffer->dtype.lanes()));
- if (!(buffer->dtype == DataType::Int(1) || buffer->dtype == DataType::Int(4) ||
- buffer->dtype == DataType::UInt(4))) {
- auto type_msg = tvm::tir::StringImm(type_err_msg.str());
- asserts_.emplace_back(AssertStmt(cond, type_msg));
- }
-
- // shape field
- Buffer buf_shape = decl_buffer({IntImm(DataType::Int(32), buffer->shape.size())}, tvm_shape_type,
- shape_handle_name());
- Var v_shape(shape_handle_name(), DataType::Handle());
- def_handle_dtype_.Set(v_shape, make_const(tvm_shape_type, 0));
- init_nest_.emplace_back(
- LetStmt(buf_shape->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrShape), nop));
- init_nest_.emplace_back(DeclBuffer(buf_shape, nop));
- for (size_t k = 0; k < buffer->shape.size(); ++k) {
- if (buffer->dtype == DataType::Int(4) || buffer->dtype == DataType::UInt(4) ||
- buffer->dtype == DataType::Int(1)) {
- break;
- }
- Bind_(buffer->shape[k],
- cast(buffer->shape[k].dtype(), BufferLoad(buf_shape, {IntImm(DataType::Int(32), k)})),
- shape_element_name(k), true);
- }
- // strides field
- Buffer buf_strides = decl_buffer({IntImm(DataType::Int(32), buffer->strides.size())},
- tvm_shape_type, arg_name + ".strides");
- def_handle_dtype_.Set(buf_strides->data, tir::TypeAnnotation(tvm_shape_type));
- init_nest_.emplace_back(LetStmt(
- buf_strides->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrStrides), nop));
- init_nest_.emplace_back(DeclBuffer(buf_strides, nop));
- PrimExpr v_strides_is_null = Call(DataType::Bool(), builtin::isnullptr(), {buf_strides->data});
- if (buffer->strides.size() == 0) {
- // Assert the buffer is compact
- DataType stype = buffer->DefaultIndexType();
- PrimExpr expect_stride = make_const(stype, 1);
- ffi::Array<PrimExpr> conds;
- for (size_t i = buffer->shape.size(); i != 0; --i) {
- size_t k = i - 1;
- PrimExpr svalue = cast(stype, BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
- conds.push_back(buffer->shape[k] == 1 || expect_stride == svalue);
- expect_stride = expect_stride * buffer->shape[k];
- }
- std::ostringstream stride_err_msg;
- stride_err_msg << stride_handle_name() << ": expected to be compact array";
- if (conds.size() != 0) {
- auto stride_msg = tvm::tir::StringImm(stride_err_msg.str());
- Stmt check = AssertStmt(
- foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); },
- const_true(1), conds),
- stride_msg);
- check = IfThenElse(Not(v_strides_is_null), check);
- asserts_.emplace_back(SeqStmt({check, Evaluate(0)}));
- }
- } else if (buffer->buffer_type == kAutoBroadcast) {
- DataType stype = buffer->DefaultIndexType();
- PrimExpr stride = make_const(stype, 1);
- for (size_t i = buffer->shape.size(); i != 0; --i) {
- size_t k = i - 1;
- PrimExpr value =
- cast(buffer->shape[k].dtype(), BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
- value = tvm::if_then_else(v_strides_is_null, stride, value);
- value = tvm::if_then_else(buffer->shape[k] == 1, 0, value);
- Bind_(buffer->strides[k], value, stride_element_name(k), true);
- stride = analyzer_.Simplify(stride * buffer->shape[k]);
- }
- } else {
- PrimExpr stride_from_shape = 1;
-
- for (int k = buffer->strides.size() - 1; k >= 0; k--) {
- PrimExpr explicit_stride =
- cast(buffer->shape[k].dtype(), BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
-
- Bind_(buffer->strides[k],
- tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride),
- stride_element_name(k), true);
-
- stride_from_shape *=
- cast(buffer->shape[k].dtype(), BufferLoad(buf_shape, {IntImm(DataType::Int(32), k)}));
- }
- }
- // Byte_offset field.
- int data_bytes = GetVectorBytes(buffer->dtype);
-
- if (const auto* const_offset = buffer->elem_offset.as<IntImmNode>()) {
- Bind_(make_const(DataType::UInt(64), const_offset->value * data_bytes),
- TVMArrayGet(DataType::UInt(64), handle, builtin::kArrByteOffset),
- arg_name + ".byte_offset", true);
- } else {
- if (Bind_(buffer->elem_offset,
- cast(buffer->elem_offset.dtype(),
- (TVMArrayGet(DataType::UInt(64), handle, builtin::kArrByteOffset) /
- make_const(DataType::UInt(64), data_bytes))),
- arg_name + ".elem_offset", true)) {
- if (buffer->offset_factor > 1) {
- PrimExpr offset = buffer->elem_offset;
- PrimExpr factor = make_const(offset.dtype(), buffer->offset_factor);
- PrimExpr zero = make_zero(offset.dtype());
- BinderAddAssert(&analyzer_, truncmod(offset, factor) == zero, arg_name + ".elem_offset",
- &asserts_);
- }
- }
- }
- // device info.
- Bind_(device_type, TVMArrayGet(DataType::Int(32), handle, builtin::kArrDeviceType),
- arg_name + ".device_type", true);
- Bind_(device_id, TVMArrayGet(DataType::Int(32), handle, builtin::kArrDeviceId),
- arg_name + ".device_id", true);
-
- // Data field. Because the validation of the data field may depend
- // on a dynamic size defined by the other DLTensor* parameters, this
- // field must be generated last.
- if (Bind_(buffer->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrData),
- arg_name + ".data", true)) {
- Var vptr(buffer->data);
-
- // Check if the data pointer is NULL. This check is skipped for
- // size-0 arrays, since CUDA provides a NULL pointer for size-zero
- // allocations.
- auto alloc_size = [&]() -> PrimExpr {
- PrimExpr product = IntImm(buffer->DefaultIndexType(), 1);
- for (const auto& dim : buffer->shape) {
- product *= dim;
- }
- return product;
- }();
- asserts_.emplace_back(
- AssertStmt(alloc_size == 0 || !Call(DataType::Bool(), builtin::isnullptr(), {vptr}),
- tvm::tir::StringImm(arg_name + " is expected to have non-NULL data pointer")));
-
- def_handle_dtype_.Set(vptr, tir::TypeAnnotation(buffer->dtype));
- // mark alignment of external bufs
- init_nest_.emplace_back(AttrStmt(vptr, tir::attr::storage_alignment,
- IntImm(DataType::Int(32), buffer->data_alignment), nop));
- }
-}
-
-} // namespace tir
-} // namespace tvm
diff --git a/src/tir/transform/arg_binder.h b/src/tir/transform/arg_binder.h
deleted file mode 100644
index e9d08fafb0..0000000000
--- a/src/tir/transform/arg_binder.h
+++ /dev/null
@@ -1,169 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied. See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/*!
- * \file arg_binder.h
- * \brief Helper utility to match and bind arguments.
- */
-#ifndef TVM_TIR_TRANSFORM_ARG_BINDER_H_
-#define TVM_TIR_TRANSFORM_ARG_BINDER_H_
-
-#include <tvm/arith/analyzer.h>
-#include <tvm/tir/buffer.h>
-#include <tvm/tir/expr.h>
-
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-namespace tvm {
-namespace tir {
-
-/*!
- * \brief Helper utility to generate match and bind of arguments.
- *
- * \note There is many places in TVM IR where we need argument bindings.
- *
- * Consider a function f(tA(shape=var(n)), tB(shape=3), tC(shape=(n+2)).
- * Here n is a undefined variable that is decided by the outside, tB imposes
- * a constraint such that it can only take tensor with shape 3, tC imposes
- * another constraint that it's shape must equals n + 2.
- * So if we call it with f(bufferA, bufferB, bufferC), we need to generate
- * the following binding sequence:
- * - define n = bufferA.shape[0]
- * - assert bufferB.shape[0] == 3
- * - assert bufferB.shape[1] == n + 3
- *
- * In general, this is a constraint solving problem. We have simplified assumption
- * over the binding declaration, such that we require the variable occurred in
- * constraint must be declared in argument list. So it is illegal to have signature
- * f(tA(shape=(n+3))) without any argument variable corresponds to n, even though
- * it is already enough to derive n from the input argument.
- */
-class ArgBinder {
- public:
- /*!
- * \brief Constructor
- * \param def_map A definition map that contains definition of known variables.
- * ArgBinder will update this def_map when adding new definitions.
- */
- explicit ArgBinder(std::unordered_map<const VarNode*, PrimExpr>* def_map) : def_map_(def_map) {}
- /*!
- * \brief Try to bind arg to value, generate constraint if necessary.
- * \param arg The argument to be binded.
- * \param value The target expression value
- * \param arg_name argument name.
- * \param with_let Whether add lets during bind
- */
- void Bind(const PrimExpr& arg, const PrimExpr& value, const std::string& arg_name,
- bool with_let = false);
- /*!
- * \brief Bind array to array
- * \param arg The argument to be binded.
- * \param value The target expression value
- * \param arg_name argument name.
- */
- void BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
- const std::string& arg_name);
- /*!
- * \brief Bind symbolic buffer to another symbolic buffer
- * \param arg The argument to be binded.
- * \param value The target expression value
- * \param arg_name argument name.
- * \param fuzzy_match If enabled, we allow value's dimension to be smaller than arg, as long as
- * arg's higher dimensions are of 1.
- */
- void BindBuffer(const Buffer& arg, const Buffer& value, const std::string& arg_name,
- bool fuzzy_match);
- /*!
- * \brief Bind symbolic buffer to a DLTensor handle.
- * \param buffer The argument buffer to be binded.
- * \param device_type The device id to be binded.
- * \param device_id The device id to be binded.
- * \param handle The DLTensor handle.
- * \param arg_name argument name.
- */
- void BindDLTensor(const Buffer& buffer, const PrimExpr& device_type, const PrimExpr& device_id,
- const Var& handle, const std::string& arg_name);
-
- /*! \return The defs generated in binding. */
- const std::vector<Var>& defs() const { return defs_; }
-
- /*! \return The asserts generated in binding
- *
- * This contains statements that assert the correct value has been
- * bound. For example, `binder.Bind(var, expr_1)` will produce an
- * entry mapping `var` to `expr_1` in the `binder.defs()`. If
- * `binder.Bind(var, expr_2)` is called later, then this will
- * produce an assert statemtn that `expr_1 == expr_2`.
- *
- * Note: Some assert statements produced by BindDLTensor are located
- * in `binder.init_nest()`, not within `binder.asserts()`. This is
- * deliberate, as some values may require checks prior to
- * initialization. (e.g. Intializing `m = dl_tensor->shape[3]`
- * requires first asserting that `3 < dl_tensor->ndim`.)
- */
- const std::vector<Stmt>& asserts() const { return asserts_; }
-
- /*!
- * \brief Initialization nest generated
- *
- * This contains both variable bindings and any assert statements
- * that are required in order to safely produce those variable
- * bindings.
- *
- * \note Variable bindings may be implemented either as a `LetStmt`
- * that defines the variable, or as a variable replacement. Any
- * bindings implemented as a `LetStmt` will be in the
- * initialization list. Any bindings implemented as a variable
- * replacement will be stored in the `var_def` map.
- *
- * A `tir::LetStmt` is usually generated when binding to a
- * `DLTensor`. This requires loading values from memory, which
- * should only be performed once. If the binding to a
- * `DLTensor` were implemented as a variable replacement, it
- * would load values from memory once for each usage of the
- * variable.
- *
- * \return The initialization nest generated during binding.
- */
- const std::vector<Stmt>& init_nest() const { return init_nest_; }
- /*! \return Handle data type of the data */
- const ffi::Map<Var, PrimExpr>& def_handle_dtype() const { return def_handle_dtype_; }
-
- private:
- // Internal bind function
- bool Bind_(const PrimExpr& arg, const PrimExpr& value, const std::string& arg_name,
- bool with_lets);
- /*! \brief The definition map, can be uses to substitute */
- std::unordered_map<const VarNode*, PrimExpr>* def_map_;
- /*! \brief defs generated in the current binder */
- std::vector<Var> defs_;
- /*! \brief Initialize nest */
- std::vector<Stmt> init_nest_;
- /*! \brief handle data type in the defintiions */
- ffi::Map<Var, PrimExpr> def_handle_dtype_;
- /*! \brief asserts generated */
- std::vector<Stmt> asserts_;
- /*! \brief internal analyzer. */
- arith::Analyzer analyzer_;
-};
-} // namespace tir
-} // namespace tvm
-#endif // TVM_TIR_TRANSFORM_ARG_BINDER_H_
diff --git a/src/tir/transform/lower_tvm_builtin.cc b/src/tir/transform/lower_tvm_builtin.cc
index 2520caeb7b..2e9855e223 100644
--- a/src/tir/transform/lower_tvm_builtin.cc
+++ b/src/tir/transform/lower_tvm_builtin.cc
@@ -444,21 +444,24 @@ class BuiltinLower : public StmtExprMutator {
PrimExpr expr = StmtExprMutator::VisitExpr_(op);
op = expr.as<CallNode>();
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrData, op->args[0]));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrShape, op->args[1]));
+ prep_seq.emplace_back(
+ TVMStructSet(scope.stack_array, idx, builtin::kDLTensorData, op->args[0]));
+ prep_seq.emplace_back(
+ TVMStructSet(scope.stack_array, idx, builtin::kDLTensorShape, op->args[1]));
PrimExpr strides = op->args[2];
if (!strides.defined() || is_zero(strides)) {
strides = make_zero(DataType::Handle());
}
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrStrides, strides));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrNDim, op->args[3]));
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorStrides, strides));
+ prep_seq.emplace_back(
+ TVMStructSet(scope.stack_array, idx, builtin::kDLTensorNDim, op->args[3]));
DataType dtype = op->args[4].dtype();
prep_seq.emplace_back(
- TVMStructSet(scope.stack_array, idx, builtin::kArrTypeCode,
+ TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeCode,
make_const(DataType::UInt(8), static_cast<int>(dtype.code()))));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrTypeBits,
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeBits,
make_const(DataType::UInt(8), dtype.bits())));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrTypeLanes,
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeLanes,
make_const(DataType::UInt(16), dtype.lanes())));
// set byte offset
int data_bytes = GetVectorBytes(dtype);
@@ -469,15 +472,15 @@ class BuiltinLower : public StmtExprMutator {
} else {
byte_offset = elem_offset;
}
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrByteOffset,
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorByteOffset,
cast(DataType::UInt(64), byte_offset)));
TVM_FFI_ICHECK(device_type_) << "Unknown device type in current IR";
TVM_FFI_ICHECK(device_id_) << "Unknown device id in current IR";
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrDeviceId,
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorDeviceId,
cast(DataType::Int(32), device_id_.value())));
- prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kArrDeviceType,
+ prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorDeviceType,
cast(DataType::Int(32), device_type_.value())));
- return TVMStructGet(DataType::Handle(), scope.stack_array, idx, builtin::kArrAddr);
+ return TVMStructGet(DataType::Handle(), scope.stack_array, idx, builtin::kDLTensorAddr);
}
void SetPackedArg(PrimExpr arg, const Var& args_stack, size_t stack_offset,
@@ -648,7 +651,7 @@ class BuiltinLower : public StmtExprMutator {
// specially set array handle.
if (const CallNode* buf = arg.as<CallNode>()) {
if (buf->op.same_as(builtin::tvm_struct_get()) &&
- buf->args[2].as<IntImmNode>()->value == builtin::kArrAddr) {
+ buf->args[2].as<IntImmNode>()->value == builtin::kDLTensorAddr) {
return true;
}
}
diff --git a/src/tir/transform/make_packed_api.cc b/src/tir/transform/make_packed_api.cc
index af63d9ba54..e94c355c7a 100644
--- a/src/tir/transform/make_packed_api.cc
+++ b/src/tir/transform/make_packed_api.cc
@@ -22,6 +22,7 @@
*/
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
+#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/module.h>
@@ -37,8 +38,8 @@
#include <utility>
#include <vector>
-#include "arg_binder.h"
#include "ir_utils.h"
+#include "tvm_ffi_binder.h"
namespace tvm {
namespace tir {
@@ -167,15 +168,6 @@ class SubroutineCallRewriter : public StmtExprMutator {
} // namespace
-inline Stmt MakeAssertEQ(PrimExpr lhs, PrimExpr rhs, std::string msg) {
- return AssertStmt(lhs == rhs, tvm::tir::StringImm(msg));
-}
-
-inline Stmt MakeAssertNotNull(PrimExpr ptr, std::string msg) {
- Call isnull(DataType::Bool(), builtin::isnullptr(), {ptr});
- return AssertStmt(!isnull, tvm::tir::StringImm(msg));
-}
-
/* \brief Return the global_symbol of the function, if it should be updated
*
* \param func The function to be inspected
@@ -226,12 +218,9 @@ PrimFunc MakePackedAPI(PrimFunc func) {
}
auto* func_ptr = func.CopyOnWrite();
- // set the global symbol to the packed function name
const Stmt nop = Evaluate(0);
- int num_args = static_cast<int>(func_ptr->params.size());
// Data field definitions
- // The packed fields
Var v_self_handle("self_handle", DataType::Handle());
Var v_packed_args("args", DataType::Handle());
Var v_num_packed_args("num_args", DataType::Int(32));
@@ -240,132 +229,20 @@ PrimFunc MakePackedAPI(PrimFunc func) {
// The device context
Var device_id("dev_id");
Integer device_type(target_device_type);
- // seq_init gives sequence of initialization
- // seq_check gives sequence of later checks after init
- std::vector<Stmt> seq_init, seq_check, arg_buffer_declarations;
- std::unordered_map<const VarNode*, PrimExpr> vmap;
- ArgBinder binder(&vmap);
-
- // ---------------------------
- // local function definitions
- // load i-th argument as type t
- auto f_load_arg_value = [&](DataType arg_type, int i) {
- ffi::Array<PrimExpr> call_args{v_packed_args, IntImm(DataType::Int(32), i),
- IntImm(DataType::Int(32), builtin::kTVMFFIAnyUnionValue)};
- // load 64 bit version
- DataType api_type = APIType(arg_type);
- PrimExpr res = Call(api_type, builtin::tvm_struct_get(), call_args);
- // cast to the target version.
- if (api_type != arg_type) {
- res = Cast(arg_type, res);
- }
- return res;
- };
- // Assert correct type codes for each argument. This must be done
- // *before* any initialization steps produced by
- // `binder.BindDLTensor()`. The validity of those initialization
- // steps depends on the correct types being present, and must not
- // occur before the type codes are actually checked.
- seq_init.push_back(MakeAssertEQ(v_num_packed_args, num_args, [&]() -> std::string {
- std::ostringstream error_message;
- error_message << name_hint << ": num_args should be " << num_args;
- return error_message.str();
- }()));
-
- if (num_args > 0) {
- seq_init.push_back(MakeAssertNotNull(v_packed_args, name_hint + ": args pointer is NULL"));
- }
+ // Create TVMFFIABIBuilder and decode all packed args
+ TVMFFIABIBuilder binder(name_hint, func_ptr->params, func_ptr->buffer_map, v_packed_args,
+ v_num_packed_args, device_type, device_id);
+ binder.DecodeAllParams();
- // Need to delay binding of the buffers, in case some arguments also
- // appear in the buffer.
- std::vector<std::pair<PrimExpr, Var>> var_def;
- std::vector<std::pair<Var, Buffer>> buffer_def;
-
- for (int i = 0; i < static_cast<int>(func_ptr->params.size()); ++i) {
- Var param = func_ptr->params[i];
- PrimExpr arg_value;
- // type index checks
- Var type_index(param->name_hint + ".type_index", DataType::Int(32));
- seq_init.push_back(LetStmt(type_index,
- tir::Call(DataType::Int(32), builtin::tvm_struct_get(),
- {v_packed_args, IntImm(DataType::Int(32), i),
- IntImm(DataType::Int(32), builtin::kTVMFFIAnyTypeIndex)}),
- nop));
- DataType dtype = param.dtype();
- if (dtype.is_handle()) {
- std::ostringstream msg;
- msg << name_hint << ": Expect arg[" << i << "] to be pointer";
- seq_init.emplace_back(AssertStmt(type_index == ffi::TypeIndex::kTVMFFINone ||
- type_index == ffi::TypeIndex::kTVMFFIOpaquePtr ||
- type_index == ffi::TypeIndex::kTVMFFIDLTensorPtr ||
- type_index >= ffi::TypeIndex::kTVMFFIStaticObjectBegin,
- tvm::tir::StringImm(msg.str())));
- // if type_index is Tensor, we need to add the offset of the DLTensor header
- // which always equals 16 bytes, this ensures that T.handle always shows up as a DLTensor*
- const int64_t object_cell_offset = sizeof(TVMFFIObject);
- static_assert(object_cell_offset == 24);
- arg_value = f_load_arg_value(param.dtype(), i);
- PrimExpr handle_from_tensor =
- Call(DataType::Handle(), tir::builtin::handle_add_byte_offset(),
- {arg_value, IntImm(DataType::Int(32), object_cell_offset)});
- arg_value =
- Select(type_index == ffi::TypeIndex::kTVMFFITensor, handle_from_tensor, arg_value);
- } else if (dtype.is_bool()) {
- std::ostringstream msg;
- msg << name_hint << ": Expect arg[" << i << "] to be boolean";
- seq_init.emplace_back(AssertStmt(
- type_index == ffi::TypeIndex::kTVMFFIBool || type_index == ffi::TypeIndex::kTVMFFIInt,
- tvm::tir::StringImm(msg.str())));
- arg_value = Cast(DataType::Bool(), f_load_arg_value(DataType::Int(64), i));
+ auto result = binder.Finalize();
+ bool need_set_device = result.var_defs.count(device_id.get());
- } else if (dtype.is_int() || dtype.is_uint()) {
- std::ostringstream msg;
- msg << name_hint << ": Expect arg[" << i << "] to be int";
- seq_init.emplace_back(AssertStmt(
- type_index == ffi::TypeIndex::kTVMFFIInt || type_index == ffi::TypeIndex::kTVMFFIBool,
- tvm::tir::StringImm(msg.str())));
- arg_value = f_load_arg_value(param.dtype(), i);
- } else {
- TVM_FFI_ICHECK(dtype.is_float());
- std::ostringstream msg;
- msg << name_hint << ": Expect arg[" << i << "] to be float";
- seq_init.emplace_back(AssertStmt(type_index == ffi::TypeIndex::kTVMFFIFloat ||
- type_index == ffi::TypeIndex::kTVMFFIInt ||
- type_index == ffi::TypeIndex::kTVMFFIBool,
- tvm::tir::StringImm(msg.str())));
- // use select so we can also handle int conversion to bool
- arg_value = tir::Select(
- type_index == ffi::TypeIndex::kTVMFFIFloat,
- /* true_value = */ f_load_arg_value(param.dtype(), i),
- /* false_value = */ Cast(param.dtype(), f_load_arg_value(DataType::Int(64), i)));
- }
- var_def.emplace_back(arg_value, param);
- if (func_ptr->buffer_map.count(param)) {
- // buffer binding now depends on type index
- // if the index is Tensor handle, we need to offset to get the DLTensor*
- buffer_def.emplace_back(param, func_ptr->buffer_map[param]);
- }
- }
+ std::vector<Stmt> seq_check;
// signature: (void* handle, TVMFFIAny* packed_args, int num_args, TVMFFIAny* v_result)
ffi::Array<Var> args{v_self_handle, v_packed_args, v_num_packed_args, v_result};
- // Arg definitions are defined before buffer binding to avoid the use before
- // def errors.
- //
- // For example, for auto broadcasting, checks are required to guarantee that
- // either 0 or the original stride will be correctly used. Checks here have
- // to use the args that may have no let binding yet. Therefore, hoisting let
- // binding for args before buffer declaration is needed.
- for (const auto& [expr, param] : var_def) {
- binder.Bind(param, expr, name_hint + "." + param->name_hint, true);
- }
-
- for (const auto& [var, buffer] : buffer_def) {
- binder.BindDLTensor(buffer, device_type, device_id, var, name_hint + "." + var->name_hint);
- arg_buffer_declarations.push_back(DeclBuffer(buffer, nop));
- }
// reset global symbol to attach prefix
func = WithAttrs(
std::move(func),
@@ -377,7 +254,7 @@ PrimFunc MakePackedAPI(PrimFunc func) {
body = AttrStmt(make_zero(DataType::Int(32)), attr::compute_scope,
StringImm(name_hint + "_compute_"), body);
// Set device context
- if (vmap.count(device_id.get())) {
+ if (need_set_device) {
ffi::Any node = ffi::String("default");
seq_check.push_back(AttrStmt(node, attr::device_id, device_id, nop));
seq_check.push_back(AttrStmt(node, attr::device_type, device_type, nop));
@@ -393,8 +270,9 @@ PrimFunc MakePackedAPI(PrimFunc func) {
// Return error code of zero on success
body = SeqStmt({body, Evaluate(ret(Integer(0)))});
- body = MergeNest(
- {seq_init, binder.init_nest(), seq_check, binder.asserts(), arg_buffer_declarations}, body);
+ body = MergeNest({std::move(result.init_nest), seq_check, std::move(result.asserts),
+ std::move(result.decl_buffers)},
+ body);
func_ptr->body = body;
func_ptr->params = args;
diff --git a/src/tir/transform/split_host_device.cc b/src/tir/transform/split_host_device.cc
index 43b6701f1f..e832e9d1ca 100644
--- a/src/tir/transform/split_host_device.cc
+++ b/src/tir/transform/split_host_device.cc
@@ -103,8 +103,8 @@ class HostDeviceSplitter : public StmtMutator {
if (can_propagate_errors) {
Var kernel_error_code("kernel_error_code", success->dtype);
Call kernel_call(success->dtype, kernel_symbol_global, args);
- AssertStmt assert_success(kernel_error_code == success,
- StringImm("Error executing compute kernel"));
+ AssertStmt assert_success(kernel_error_code == success, StringImm("RuntimeError"),
+ {StringImm("Error executing compute kernel")});
LetStmt let_check(kernel_error_code, kernel_call, assert_success);
return let_check;
diff --git a/src/tir/transform/tvm_ffi_binder.cc b/src/tir/transform/tvm_ffi_binder.cc
new file mode 100644
index 0000000000..02058970e1
--- /dev/null
+++ b/src/tir/transform/tvm_ffi_binder.cc
@@ -0,0 +1,810 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm_ffi_binder.cc
+ * \brief Helper utility to match and bind packed function arguments.
+ */
+#include "tvm_ffi_binder.h"
+
+#include <tvm/runtime/device_api.h>
+#include <tvm/tir/builtin.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/expr_functor.h>
+#include <tvm/tir/op.h>
+
+#include "ir_utils.h"
+
+namespace tvm {
+namespace tir {
+
+using ffi::reflection::AccessPath;
+using ffi::reflection::AccessStep;
+
+// ============================================================
+// Constructor
+// ============================================================
+
+TVMFFIABIBuilder::TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array<Var>& params,
+ const ffi::Map<Var, Buffer>& buffer_map,
+ const Var& v_packed_args, const Var& v_num_packed_args,
+ const PrimExpr& device_type, const PrimExpr& device_id)
+ : func_name_(func_name),
+ params_(params),
+ buffer_map_(buffer_map),
+ v_packed_args_(v_packed_args),
+ device_type_(device_type),
+ device_id_(device_id) {
+ // Build function signature string
+ std::ostringstream os;
+ os << func_name << "(";
+ for (size_t i = 0; i < params.size(); ++i) {
+ if (i > 0) os << ", ";
+ Var param = params[i];
+ if (buffer_map.count(param)) {
+ Buffer buf = buffer_map[param];
+ std::string buf_name = buf->name;
+ os << buf_name << ": Tensor([";
+ for (size_t j = 0; j < buf->shape.size(); ++j) {
+ if (j > 0) os << ", ";
+ std::ostringstream shape_os;
+ shape_os << buf->shape[j];
+ os << shape_os.str();
+ }
+ os << "], " << buf->dtype << ")";
+ param_names_[static_cast<int>(i)] = buf_name;
+ } else {
+ os << param->name_hint << ": " << param.dtype();
+ param_names_[static_cast<int>(i)] = param->name_hint;
+ }
+ }
+ os << ")";
+ func_signature_ = os.str();
+ sig_imm_ = StringImm(func_signature_);
+
+ // Emit argument count check (early check — must execute before any loads)
+ int num_args = static_cast<int>(params.size());
+ EmitAssert(v_num_packed_args == num_args, "TypeError", //
+ "Expected ", std::to_string(num_args), " arguments", when_calling_imm_, sig_imm_, "`");
+
+ // Emit null-pointer check for packed args (early check)
+ if (num_args > 0) {
+ EmitAssert(!Call(DataType::Bool(), builtin::isnullptr(), {v_packed_args}),
+ "TypeError", //
+ "args pointer is NULL", when_calling_imm_, sig_imm_, "`");
+ }
+}
+
+// ============================================================
+// EmitTypeIndexCheck
+// ============================================================
+
+void TVMFFIABIBuilder::EmitTypeIndexCheck(int param_index, const PrimExpr& cond,
+ const std::string& expected_type) {
+ EmitAssert(cond, "TypeError", //
+ "Mismatched type on argument #", std::to_string(param_index), when_calling_imm_,
+ sig_imm_, "`,\n expected ", expected_type);
+}
+
+// ============================================================
+// RenderAccessPath
+// ============================================================
+
+ffi::String TVMFFIABIBuilder::RenderAccessPath(const AccessPath& path) const {
+ ffi::Array<AccessStep> steps = path->ToSteps();
+ std::ostringstream os;
+ bool first_printed = false;
+ bool consumed_root_array_item = false;
+ for (size_t i = 0; i < steps.size(); ++i) {
+ const AccessStep& step = steps[i];
+ if (step->kind == ffi::reflection::AccessKind::kAttr) {
+ std::string field_name = step->key.cast<ffi::String>();
+ if (!first_printed) {
+ os << field_name;
+ first_printed = true;
+ } else {
+ os << "." << field_name;
+ }
+ } else if (step->kind == ffi::reflection::AccessKind::kArrayItem) {
+ int64_t index = step->key.cast<int64_t>();
+ if (!first_printed && !consumed_root_array_item) {
+ consumed_root_array_item = true;
+ if (i + 1 < steps.size() && steps[i + 1]->kind == ffi::reflection::AccessKind::kAttr) {
+ continue;
+ }
+ auto it = param_names_.find(static_cast<int>(index));
+ if (it != param_names_.end()) {
+ os << it->second;
+ first_printed = true;
+ }
+ } else {
+ os << "[" << index << "]";
+ }
+ } else if (step->kind == ffi::reflection::AccessKind::kMapItem) {
+ os << "[" << step->key << "]";
+ }
+ }
+ return os.str();
+}
+
+// ============================================================
+// GetParamIndex
+// ============================================================
+
+int TVMFFIABIBuilder::GetParamIndex(const AccessPath& path) const {
+ ffi::Array<AccessStep> steps = path->ToSteps();
+ if (steps.size() >= 1 && steps[0]->kind == ffi::reflection::AccessKind::kArrayItem) {
+ return static_cast<int>(steps[0]->key.cast<int64_t>());
+ }
+ return -1;
+}
+
+// ============================================================
+// BindScalar (scalar bind with AccessPath)
+// ============================================================
+
+bool TVMFFIABIBuilder::BindScalar(const PrimExpr& arg, const PrimExpr& value,
+ const AccessPath& path, bool with_lets) {
+ TVM_FFI_ICHECK_EQ(arg.dtype(), value.dtype());
+ if (arg.as<VarNode>()) {
+ Var v_arg = Downcast<Var>(arg);
+ auto it = var_defs_.find(v_arg.get());
+ if (it == var_defs_.end()) {
+ // First bind: define the variable
+ if (with_lets) {
+ var_defs_.emplace(v_arg.get(), VarDefInfo{arg, path});
+ init_nest_.emplace_back(LetStmt(v_arg, value, Evaluate(0)));
+ } else {
+ var_defs_.emplace(v_arg.get(), VarDefInfo{value, path});
+ }
+ return true;
+ } else {
+ // Duplicate bind: create rich assertion with both paths
+ PrimExpr prev_value = it->second.value;
+ PrimExpr scond = analyzer_.Simplify(prev_value == value);
+ if (is_zero(scond)) {
+ TVM_FFI_THROW(InternalError) << "Bind have an unmet assertion: " << prev_value
+ << " == " << value << " at " << RenderAccessPath(path);
+ }
+ if (!is_one(scond)) {
+ ffi::String current_path_str = RenderAccessPath(path);
+ ffi::String first_path_str = RenderAccessPath(it->second.first_def_path);
+ int param_index = GetParamIndex(path);
+ ffi::Array<StringImm> parts;
+ parts.push_back(StringImm("Mismatched "));
+ parts.push_back(StringImm(current_path_str));
+ if (param_index >= 0) {
+ parts.push_back(StringImm(" on argument #"));
+ parts.push_back(StringImm(std::to_string(param_index)));
+ }
+ parts.push_back(when_calling_imm_);
+ parts.push_back(sig_imm_);
+ if (!first_path_str.empty() && first_path_str != current_path_str) {
+ parts.push_back(StringImm("`,\n expected to match "));
+ parts.push_back(StringImm(first_path_str));
+ } else {
+ parts.push_back(StringImm("`,\n expected matching value"));
+ }
+ asserts_.emplace_back(AssertStmt(scond, StringImm("ValueError"), parts));
+ }
+ }
+ } else {
+ // Non-Var expression (e.g. batch_size + 1): defer assertion to Finalize()
+ // so display-var substitution can render human-readable names.
+ PrimExpr scond = analyzer_.Simplify(arg == value);
+ if (is_zero(scond)) {
+ TVM_FFI_THROW(InternalError) << "Bind have an unmet assertion: " << arg << " == " << value
+ << " at " << RenderAccessPath(path);
+ }
+ if (!is_one(scond)) {
+ pending_const_asserts_.push_back({scond, path, arg});
+ }
+ }
+ return false;
+}
+
+// ============================================================
+// ExprPathRenderer (generic expr-to-string via ExprFunctor)
+// ============================================================
+
+/*!
+ * \brief Render PrimExpr to string with variable names replaced by AccessPath names.
+ *
+ * Uses ExprFunctor for generic dispatch over all expression types.
+ * The default TIR printer sanitizes Var name_hints (e.g. "B.shape[0]" -> "B_shape_0_")
+ * and adds type annotations (e.g. T.int64(1)). This functor preserves original path
+ * names and uses plain integer formatting for human-readable error messages.
+ */
+class ExprPathRenderer : public ExprFunctor<std::string(const PrimExpr&)> {
+ public:
+ using FVarName = std::function<std::string(const VarNode*)>;
+ explicit ExprPathRenderer(FVarName f_var_name) : f_var_name_(std::move(f_var_name)) {}
+
+ protected:
+ std::string VisitExpr_(const VarNode* op) final { return f_var_name_(op); }
+ std::string VisitExpr_(const IntImmNode* op) final { return std::to_string(op->value); }
+ std::string VisitExpr_(const FloatImmNode* op) final {
+ std::ostringstream os;
+ os << op->value;
+ return os.str();
+ }
+ std::string VisitExpr_(const CastNode* op) final { return VisitExpr(op->value); }
+ std::string VisitExpr_(const AddNode* op) final { return BinOp(op->a, " + ", op->b); }
+ std::string VisitExpr_(const SubNode* op) final { return BinOp(op->a, " - ", op->b); }
+ std::string VisitExpr_(const MulNode* op) final { return BinOp(op->a, " * ", op->b); }
+ std::string VisitExpr_(const DivNode* op) final { return BinOp(op->a, " / ", op->b); }
+ std::string VisitExpr_(const ModNode* op) final { return BinOp(op->a, " % ", op->b); }
+ std::string VisitExpr_(const FloorDivNode* op) final { return FuncOp("floordiv", op->a, op->b); }
+ std::string VisitExpr_(const FloorModNode* op) final { return FuncOp("floormod", op->a, op->b); }
+ std::string VisitExpr_(const MinNode* op) final { return FuncOp("min", op->a, op->b); }
+ std::string VisitExpr_(const MaxNode* op) final { return FuncOp("max", op->a, op->b); }
+ // Fallback: use operator<< for unhandled expression types.
+ std::string VisitExprDefault_(const Object* op) final {
+ std::ostringstream os;
+ os << ffi::GetRef<PrimExpr>(static_cast<const PrimExprNode*>(op));
+ return os.str();
+ }
+
+ private:
+ std::string BinOp(const PrimExpr& a, const char* op, const PrimExpr& b) {
+ return VisitExpr(a) + op + VisitExpr(b);
+ }
+ std::string FuncOp(const char* name, const PrimExpr& a, const PrimExpr& b) {
+ return std::string(name) + "(" + VisitExpr(a) + ", " + VisitExpr(b) + ")";
+ }
+ FVarName f_var_name_;
+};
+
+// ============================================================
+// RenderPendingAsserts
+// ============================================================
+
+void TVMFFIABIBuilder::RenderPendingAsserts() {
+ if (pending_const_asserts_.empty()) return;
+
+ ExprPathRenderer renderer([this](const VarNode* v) -> std::string {
+ auto it = var_defs_.find(v);
+ if (it != var_defs_.end()) {
+ ffi::String path = RenderAccessPath(it->second.first_def_path);
+ if (!path.empty()) return std::string(path);
+ }
+ return std::string(v->name_hint);
+ });
+
+ for (auto& pending : pending_const_asserts_) {
+ std::string display_str = renderer.VisitExpr(pending.expected_expr);
+
+ ffi::String path_str = RenderAccessPath(pending.path);
+ int param_index = GetParamIndex(pending.path);
+ ffi::Array<StringImm> parts;
+ parts.push_back(StringImm("Invalid "));
+ parts.push_back(StringImm(path_str));
+ if (param_index >= 0) {
+ parts.push_back(StringImm(" on argument #"));
+ parts.push_back(StringImm(std::to_string(param_index)));
+ }
+ parts.push_back(when_calling_imm_);
+ parts.push_back(sig_imm_);
+ parts.push_back(StringImm("`,\n expected "));
+ parts.push_back(StringImm(display_str));
+ asserts_.emplace_back(AssertStmt(pending.condition, StringImm("ValueError"), parts));
+ }
+ pending_const_asserts_.clear();
+}
+
+// ============================================================
+// Finalize
+// ============================================================
+
+TVMFFIABIBuilder::Result TVMFFIABIBuilder::Finalize() {
+ RenderPendingAsserts();
+ Result result;
+ result.var_defs = std::move(var_defs_);
+ result.init_nest = std::move(init_nest_);
+ result.asserts = std::move(asserts_);
+ result.decl_buffers = std::move(decl_buffers_);
+ return result;
+}
+
+// ============================================================
+// BindArray (array bind with AccessPath)
+// ============================================================
+
+void TVMFFIABIBuilder::BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
+ const AccessPath& base_path) {
+ TVM_FFI_ICHECK_EQ(arg.size(), value.size())
+ << "Array size mismatch at " << RenderAccessPath(base_path);
+ for (size_t i = 0; i < arg.size(); ++i) {
+ AccessPath elem_path = base_path->ArrayItem(i);
+ BindScalar(arg[i], value[i], elem_path, false);
+ }
+}
+
+// ============================================================
+// BindBuffer (buffer-to-buffer bind with AccessPath)
+// ============================================================
+
+void TVMFFIABIBuilder::BindBuffer(const Buffer& arg, const Buffer& value, AccessPath base_path,
+ bool fuzzy_match) {
+ TVM_FFI_ICHECK_EQ(arg.scope(), value.scope())
+ << "Argument " << arg->name << " Buffer bind scope mismatch";
+ TVM_FFI_ICHECK_EQ(arg->dtype, value->dtype)
+ << "Argument " << arg->name << " Buffer bind data type mismatch";
+ if (value->data_alignment % arg->data_alignment != 0) {
+ LOG(WARNING) << "Trying to bind buffer to another one with lower alignment requirement "
+ << " required alignment=" << arg->data_alignment
+ << ", provided alignment=" << value->data_alignment;
+ }
+
+ if (value->elem_offset.defined()) {
+ if (is_zero(arg->elem_offset)) {
+ TVM_FFI_ICHECK(is_zero(value->elem_offset))
+ << "Trying to bind a Buffer with offset into one without offset "
+ << " required elem_offset=" << arg->elem_offset
+ << ", provided elem_offset=" << value->elem_offset;
+ }
+ AccessPath data_path = base_path->Attr(ffi::String("data"));
+ BindScalar(arg->data, value->data, data_path, false);
+ AccessPath offset_path = base_path->Attr(ffi::String("elem_offset"));
+ if (BindScalar(arg->elem_offset, value->elem_offset, offset_path, false)) {
+ if (arg->offset_factor > 1) {
+ PrimExpr offset = value->elem_offset;
+ PrimExpr factor = make_const(offset.dtype(), arg->offset_factor);
+ PrimExpr zero = make_zero(offset.dtype());
+ PrimExpr acond = analyzer_.Simplify(truncmod(offset, factor) == zero);
+ if (is_zero(acond)) {
+ TVM_FFI_THROW(InternalError)
+ << "Bind have an unmet assertion at " << RenderAccessPath(offset_path);
+ }
+ if (!is_one(acond)) {
+ int param_index = GetParamIndex(base_path);
+ int data_bytes = GetVectorBytes(arg->dtype);
+ EmitAssert(acond, "ValueError", //
+ "Misaligned buffer data on argument #", std::to_string(param_index),
+ when_calling_imm_, sig_imm_, "`,\n expected data alignment=",
+ std::to_string(arg->offset_factor * data_bytes), " bytes");
+ }
+ }
+ }
+ }
+
+ AccessPath shape_path = base_path->Attr(ffi::String("shape"));
+ AccessPath strides_path = base_path->Attr(ffi::String("strides"));
+
+ if (arg->shape.size() < value->shape.size()) {
+ TVM_FFI_ICHECK(fuzzy_match) << "Buffer size mismatch at " << RenderAccessPath(base_path);
+ size_t diff = value->shape.size() - arg->shape.size();
+ for (size_t i = 0; i < diff; ++i) {
+ TVM_FFI_ICHECK(is_one(analyzer_.Simplify(value->shape[i])))
+ << "Buffer shape mismatch at " << RenderAccessPath(base_path) << ": " << arg->shape
+ << " vs " << value->shape;
+ }
+ for (size_t i = 0; i < arg->shape.size(); ++i) {
+ AccessPath shape_k_path = shape_path->ArrayItem(i);
+ BindScalar(arg->shape[i], value->shape[i + diff], shape_k_path, false);
+ }
+ if (value->strides.size() != 0) {
+ TVM_FFI_ICHECK_EQ(arg->strides.size(), arg->shape.size());
+ TVM_FFI_ICHECK_EQ(value->strides.size(), value->shape.size());
+ for (size_t i = 0; i < arg->strides.size(); ++i) {
+ AccessPath strides_k_path = strides_path->ArrayItem(i);
+ BindScalar(arg->strides[i], value->strides[i + diff], strides_k_path, false);
+ }
+ }
+ } else {
+ BindArray(arg->shape, value->shape, shape_path);
+ BindArray(arg->strides, value->strides, strides_path);
+ }
+}
+
+// ============================================================
+// Per-dtype type-check + value-load methods
+// ============================================================
+
+/*! \brief Load the i-th packed argument as the given type. */
+PrimExpr TVMFFIABIBuilder::LoadTVMFFIAnyUnionValue(const Var& v_packed_args, int param_index,
+ DataType arg_type) {
+ ffi::Array<PrimExpr> call_args{v_packed_args, IntImm(DataType::Int(32), param_index),
+ IntImm(DataType::Int(32), builtin::kTVMFFIAnyUnionValue)};
+ DataType api_type = APIType(arg_type);
+ PrimExpr res = Call(api_type, builtin::tvm_struct_get(), call_args);
+ if (api_type != arg_type) {
+ res = Cast(arg_type, res);
+ }
+ return res;
+}
+
+PrimExpr TVMFFIABIBuilder::DecodeParamOpaqueHandle(int param_index, const Var& type_index) {
+ // ── Type check: accept handle-like types ───────────────────
+ std::string expected_type = buffer_map_.count(params_[param_index]) ? "Tensor" : "pointer";
+ EmitTypeIndexCheck(param_index,
+ type_index == ffi::TypeIndex::kTVMFFINone ||
+ type_index == ffi::TypeIndex::kTVMFFIOpaquePtr ||
+ type_index == ffi::TypeIndex::kTVMFFIDLTensorPtr ||
+ type_index >= ffi::TypeIndex::kTVMFFIStaticObjectBegin,
+ expected_type);
+
+ // ── Load value and apply tensor offset ─────────────────────
+ const int64_t object_cell_offset = sizeof(TVMFFIObject);
+ static_assert(sizeof(TVMFFIObject) == 24);
+ PrimExpr arg_value =
+ LoadTVMFFIAnyUnionValue(v_packed_args_, param_index, params_[param_index].dtype());
+ PrimExpr handle_from_tensor = Call(DataType::Handle(), tir::builtin::handle_add_byte_offset(),
+ {arg_value, IntImm(DataType::Int(32), object_cell_offset)});
+ return Select(type_index == ffi::TypeIndex::kTVMFFITensor, handle_from_tensor, arg_value);
+}
+
+PrimExpr TVMFFIABIBuilder::DecodeParamBool(int param_index, const Var& type_index) {
+ // ── Type check: accept bool or int ─────────────────────────
+ EmitTypeIndexCheck(
+ param_index,
+ type_index == ffi::TypeIndex::kTVMFFIBool || type_index == ffi::TypeIndex::kTVMFFIInt,
+ "boolean");
+ return Cast(DataType::Bool(),
+ LoadTVMFFIAnyUnionValue(v_packed_args_, param_index, DataType::Int(64)));
+}
+
+PrimExpr TVMFFIABIBuilder::DecodeParamInt(int param_index, const Var& type_index, DataType dtype) {
+ // ── Type check: accept int or bool ─────────────────────────
+ EmitTypeIndexCheck(
+ param_index,
+ type_index == ffi::TypeIndex::kTVMFFIInt || type_index == ffi::TypeIndex::kTVMFFIBool, "int");
+ return LoadTVMFFIAnyUnionValue(v_packed_args_, param_index, dtype);
+}
+
+PrimExpr TVMFFIABIBuilder::DecodeParamFloat(int param_index, const Var& type_index,
+ DataType dtype) {
+ // ── Type check: accept float, int, or bool ─────────────────
+ EmitTypeIndexCheck(param_index,
+ type_index == ffi::TypeIndex::kTVMFFIFloat ||
+ type_index == ffi::TypeIndex::kTVMFFIInt ||
+ type_index == ffi::TypeIndex::kTVMFFIBool,
+ "float");
+ return tir::Select(
+ type_index == ffi::TypeIndex::kTVMFFIFloat,
+ /* true_value = */ LoadTVMFFIAnyUnionValue(v_packed_args_, param_index, dtype),
+ /* false_value = */
+ Cast(dtype, LoadTVMFFIAnyUnionValue(v_packed_args_, param_index, DataType::Int(64))));
+}
+
+// ============================================================
+// DecodeParam (primary public method)
+// ============================================================
+
+void TVMFFIABIBuilder::DecodeParam(int param_index) {
+ const Stmt nop = Evaluate(0);
+ Var param = params_[param_index];
+ DataType dtype = param.dtype();
+
+ // Extract type_index from packed_args
+ Var type_index(param->name_hint + ".type_index", DataType::Int(32));
+ init_nest_.push_back(LetStmt(type_index,
+ tir::Call(DataType::Int(32), builtin::tvm_struct_get(),
+ {v_packed_args_, IntImm(DataType::Int(32), param_index),
+ IntImm(DataType::Int(32), builtin::kTVMFFIAnyTypeIndex)}),
+ nop));
+
+ // Type-check and load value via per-dtype dispatch
+ PrimExpr arg_value;
+ if (dtype.is_handle()) {
+ arg_value = DecodeParamOpaqueHandle(param_index, type_index);
+ } else if (dtype.is_bool()) {
+ arg_value = DecodeParamBool(param_index, type_index);
+ } else if (dtype.is_int() || dtype.is_uint()) {
+ arg_value = DecodeParamInt(param_index, type_index, dtype);
+ } else {
+ TVM_FFI_ICHECK(dtype.is_float());
+ arg_value = DecodeParamFloat(param_index, type_index, dtype);
+ }
+
+ // Bind scalar param to loaded value (defines vars before buffer binds reference them)
+ AccessPath param_path = AccessPath::Root()->Extend(AccessStep::ArrayItem(param_index));
+ BindScalar(param, arg_value, param_path, true);
+}
+
+// ============================================================
+// DecodeAllParams (primary public method)
+// ============================================================
+
+void TVMFFIABIBuilder::DecodeAllParams() {
+ const Stmt nop = Evaluate(0);
+ int num_args = static_cast<int>(params_.size());
+
+ // Phase 1: Decode each packed argument (type-check, value load, scalar bind)
+ for (int i = 0; i < num_args; ++i) {
+ DecodeParam(i);
+ }
+
+ // Phase 2: Bind DLTensor buffers (shape, strides, dtype, device checks)
+ for (int i = 0; i < num_args; ++i) {
+ Var param = params_[i];
+ if (buffer_map_.count(param)) {
+ Buffer buffer = buffer_map_[param];
+ AccessPath param_path =
+ AccessPath::Root()->Extend(AccessStep::ArrayItem(i))->Attr(ffi::String(buffer->name));
+ DecodeParamDLTensor(buffer, device_type_, device_id_, param,
+ func_name_ + "." + param->name_hint, param_path);
+ decl_buffers_.push_back(DeclBuffer(buffer, nop));
+ }
+ }
+}
+
+// ============================================================
+// DLTensorGetFieldPtr (helper for shape/strides pointer extraction)
+// ============================================================
+
+Var TVMFFIABIBuilder::DLTensorGetFieldPtr(const Var& handle, int field_kind,
+ const std::string& var_name) {
+ Var ptr(var_name, DataType::Handle());
+ init_nest_.emplace_back(
+ LetStmt(ptr,
+ TVMStructGet(DataType::Handle(), handle, 0,
+ static_cast<builtin::TVMStructFieldKind>(field_kind)),
+ Evaluate(0)));
+ return ptr;
+}
+
+// ============================================================
+// LoadInt64ArrayElem (helper for shape/strides element access)
+// ============================================================
+
+PrimExpr TVMFFIABIBuilder::LoadInt64ArrayElem(const Var& ptr, int index) {
+ return TVMStructGet(DataType::ShapeIndex(), ptr, index, builtin::kInt64ArrayElem);
+}
+
+// ============================================================
+// Strides validation subfunctions
+// ============================================================
+
+void TVMFFIABIBuilder::BindCompactStrides(const Buffer& buffer, const Var& strides_ptr,
+ const PrimExpr& v_strides_is_null,
+ const AccessPath& param_path) {
+ DataType stype = buffer->DefaultIndexType();
+ PrimExpr expect_stride = make_const(stype, 1);
+ ffi::Array<PrimExpr> conds;
+ for (size_t i = buffer->shape.size(); i != 0; --i) {
+ size_t k = i - 1;
+ PrimExpr svalue = cast(stype, LoadInt64ArrayElem(strides_ptr, k));
+ conds.push_back(buffer->shape[k] == 1 || expect_stride == svalue);
+ expect_stride = expect_stride * buffer->shape[k];
+ }
+ if (conds.size() != 0) {
+ int param_index = GetParamIndex(param_path);
+ Stmt check = AssertStmt(
+ foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); },
+ const_true(1), conds),
+ StringImm("ValueError"),
+ ffi::Array<StringImm>({StringImm("Mismatched "), StringImm(buffer->name),
+ StringImm(".strides on argument #"),
+ StringImm(std::to_string(param_index)), when_calling_imm_, sig_imm_,
+ StringImm("`,\n expected to be compact array")}));
+ check = IfThenElse(Not(v_strides_is_null), check);
+ asserts_.emplace_back(SeqStmt({check, Evaluate(0)}));
+ }
+}
+
+void TVMFFIABIBuilder::BindAutoBroadcastStrides(const Buffer& buffer, const Var& strides_ptr,
+ const PrimExpr& v_strides_is_null,
+ const AccessPath& param_path) {
+ DataType stype = buffer->DefaultIndexType();
+ PrimExpr stride = make_const(stype, 1);
+ for (size_t i = buffer->shape.size(); i != 0; --i) {
+ size_t k = i - 1;
+ PrimExpr value = cast(buffer->shape[k].dtype(), LoadInt64ArrayElem(strides_ptr, k));
+ value = tvm::if_then_else(v_strides_is_null, stride, value);
+ value = tvm::if_then_else(buffer->shape[k] == 1, 0, value);
+ AccessPath strides_k_path = param_path->Attr(ffi::String("strides"))->ArrayItem(k);
+ BindScalar(buffer->strides[k], value, strides_k_path, true);
+ stride = analyzer_.Simplify(stride * buffer->shape[k]);
+ }
+}
+
+void TVMFFIABIBuilder::BindRegularStrides(const Buffer& buffer, const Var& strides_ptr,
+ const Var& shape_ptr, const PrimExpr& v_strides_is_null,
+ const AccessPath& param_path) {
+ PrimExpr stride_from_shape = 1;
+ for (int k = buffer->strides.size() - 1; k >= 0; k--) {
+ PrimExpr explicit_stride = cast(buffer->shape[k].dtype(), LoadInt64ArrayElem(strides_ptr, k));
+ AccessPath strides_k_path = param_path->Attr(ffi::String("strides"))->ArrayItem(k);
+ BindScalar(buffer->strides[k],
+ tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride),
+ strides_k_path, true);
+ stride_from_shape *= cast(buffer->shape[k].dtype(), LoadInt64ArrayElem(shape_ptr, k));
+ }
+}
+
+// ============================================================
+// DecodeParamDLTensor (private)
+// ============================================================
+
+void TVMFFIABIBuilder::DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& device_type,
+ const PrimExpr& device_id, const Var& handle,
+ const std::string& arg_name, AccessPath base_path) {
+ const DataType tvm_ndim_type = DataType::Int(32);
+
+ std::string buf_name = buffer->name;
+ AccessPath param_path = base_path;
+ int param_index = GetParamIndex(base_path);
+
+ // ── Section: Null pointer check ──────────────────────────────
+ EmitTypeIndexCheck(param_index, !Call(DataType::Bool(), builtin::isnullptr(), {handle}),
+ "Tensor");
+
+ // ── Section: ndim ────────────────────────────────────────────
+ PrimExpr v_ndim = TVMStructGet(tvm_ndim_type, handle, 0, builtin::kDLTensorNDim);
+ PrimExpr a_ndim = make_const(tvm_ndim_type, static_cast<int64_t>(buffer->shape.size()));
+ EmitAssert(a_ndim == v_ndim, "ValueError", //
+ "Mismatched ", buf_name, ".ndim on argument #", std::to_string(param_index),
+ when_calling_imm_, sig_imm_, "`,\n expected ", std::to_string(buffer->shape.size()));
+
+ // ── Section: dtype ───────────────────────────────────────────
+ {
+ PrimExpr cond = (TVMStructGet(DataType::UInt(8), handle, 0, builtin::kDLTensorTypeCode) ==
+ IntImm(DataType::UInt(8), buffer->dtype.code()) &&
+ TVMStructGet(DataType::UInt(8), handle, 0, builtin::kDLTensorTypeBits) ==
+ IntImm(DataType::UInt(8), buffer->dtype.bits()) &&
+ TVMStructGet(DataType::UInt(16), handle, 0, builtin::kDLTensorTypeLanes) ==
+ IntImm(DataType::UInt(16), buffer->dtype.lanes()));
+ if (!(buffer->dtype == DataType::Int(1) || buffer->dtype == DataType::Int(4) ||
+ buffer->dtype == DataType::UInt(4))) {
+ std::ostringstream dtype_os;
+ dtype_os << buffer->dtype;
+ EmitAssert(cond, "TypeError", //
+ "Mismatched ", buf_name, ".dtype on argument #", std::to_string(param_index),
+ when_calling_imm_, sig_imm_, "`,\n expected ", dtype_os.str());
+ }
+ }
+
+ // ── Section: shape ───────────────────────────────────────────
+ Var shape_ptr = DLTensorGetFieldPtr(handle, builtin::kDLTensorShape, arg_name + "_shape");
+ for (size_t k = 0; k < buffer->shape.size(); ++k) {
+ if (buffer->dtype == DataType::Int(4) || buffer->dtype == DataType::UInt(4) ||
+ buffer->dtype == DataType::Int(1)) {
+ break;
+ }
+ AccessPath shape_k_path = param_path->Attr(ffi::String("shape"))->ArrayItem(k);
+ BindScalar(buffer->shape[k], cast(buffer->shape[k].dtype(), LoadInt64ArrayElem(shape_ptr, k)),
+ shape_k_path, true);
+ }
+
+ // ── Section: strides ─────────────────────────────────────────
+ Var strides_ptr = DLTensorGetFieldPtr(handle, builtin::kDLTensorStrides, arg_name + "_strides");
+ PrimExpr v_strides_is_null = Call(DataType::Bool(), builtin::isnullptr(), {strides_ptr});
+ if (buffer->strides.size() == 0) {
+ BindCompactStrides(buffer, strides_ptr, v_strides_is_null, param_path);
+ } else if (buffer->buffer_type == kAutoBroadcast) {
+ BindAutoBroadcastStrides(buffer, strides_ptr, v_strides_is_null, param_path);
+ } else {
+ BindRegularStrides(buffer, strides_ptr, shape_ptr, v_strides_is_null, param_path);
+ }
+
+ // ── Section: byte_offset ─────────────────────────────────────
+ int data_bytes = GetVectorBytes(buffer->dtype);
+ AccessPath byte_offset_path = param_path->Attr(ffi::String("byte_offset"));
+ if (const auto* const_offset = buffer->elem_offset.as<IntImmNode>()) {
+ BindScalar(make_const(DataType::UInt(64), const_offset->value * data_bytes),
+ TVMStructGet(DataType::UInt(64), handle, 0, builtin::kDLTensorByteOffset),
+ byte_offset_path, true);
+ } else {
+ if (BindScalar(buffer->elem_offset,
+ cast(buffer->elem_offset.dtype(),
+ (TVMStructGet(DataType::UInt(64), handle, 0, builtin::kDLTensorByteOffset) /
+ make_const(DataType::UInt(64), data_bytes))),
+ byte_offset_path, true)) {
+ if (buffer->offset_factor > 1) {
+ PrimExpr offset = buffer->elem_offset;
+ PrimExpr factor = make_const(offset.dtype(), buffer->offset_factor);
+ PrimExpr zero = make_zero(offset.dtype());
+ PrimExpr acond = analyzer_.Simplify(truncmod(offset, factor) == zero);
+ if (is_zero(acond)) {
+ TVM_FFI_THROW(InternalError)
+ << "Bind have an unmet assertion at " << RenderAccessPath(byte_offset_path);
+ }
+ if (!is_one(acond)) {
+ EmitAssert(acond, "ValueError", //
+ "Misaligned Tensor data on argument #", std::to_string(param_index),
+ when_calling_imm_, sig_imm_, "`,\n expected data alignment=",
+ std::to_string(buffer->offset_factor * data_bytes), " bytes");
+ }
+ }
+ }
+ }
+
+ // ── Section: device ──────────────────────────────────────────
+ {
+ PrimExpr actual_device_type =
+ TVMStructGet(DataType::Int(32), handle, 0, builtin::kDLTensorDeviceType);
+ // Use custom assertion for device_type to show human-readable device name
+ if (const auto* const_dt = device_type_.as<IntImmNode>()) {
+ PrimExpr cond =
+ analyzer_.Simplify(make_const(DataType::Int(32), const_dt->value) == actual_device_type);
+ if (!is_one(cond)) {
+ std::string device_name = runtime::DLDeviceType2Str(static_cast<int>(const_dt->value));
+ EmitAssert(cond, "ValueError", //
+ "Mismatched ", buf_name, ".device_type on argument #",
+ std::to_string(param_index), when_calling_imm_, sig_imm_, "`,\n expected ",
+ device_name);
+ }
+ } else {
+ AccessPath device_type_path = param_path->Attr(ffi::String("device_type"));
+ BindScalar(device_type_, actual_device_type, device_type_path, true);
+ }
+ AccessPath device_id_path = param_path->Attr(ffi::String("device_id"));
+ BindScalar(device_id_, TVMStructGet(DataType::Int(32), handle, 0, builtin::kDLTensorDeviceId),
+ device_id_path, true);
+ }
+
+ // ── Section: data pointer ────────────────────────────────────
+ {
+ AccessPath data_path = param_path->Attr(ffi::String("data"));
+ if (BindScalar(buffer->data,
+ TVMStructGet(DataType::Handle(), handle, 0, builtin::kDLTensorData), data_path,
+ true)) {
+ Var vptr(buffer->data);
+
+ auto alloc_size = [&]() -> PrimExpr {
+ PrimExpr product = IntImm(buffer->DefaultIndexType(), 1);
+ for (const auto& dim : buffer->shape) {
+ product *= dim;
+ }
+ return product;
+ }();
+ // Data pointer null and alignment checks go to asserts_ because alloc_size
+ // references buffer->shape which may contain forward-referenced symbolic vars.
+ asserts_.emplace_back(AssertStmt(
+ alloc_size == 0 || !Call(DataType::Bool(), builtin::isnullptr(), {vptr}),
+ StringImm("ValueError"),
+ ffi::Array<StringImm>({StringImm(buf_name),
+ StringImm(" data pointer is NULL on argument #"),
+ StringImm(std::to_string(param_index)), when_calling_imm_,
+ sig_imm_, StringImm("`,\n expected non-NULL data pointer")})));
+
+ if (check_alignment_) {
+ // Check data pointer alignment
+ if (buffer->data_alignment > 1) {
+ PrimExpr ptr_as_int =
+ Call(DataType::UInt(64), builtin::reinterpret(), {cast(DataType::Handle(), vptr)});
+ PrimExpr align_cond =
+ truncmod(ptr_as_int, make_const(DataType::UInt(64), buffer->data_alignment)) ==
+ make_const(DataType::UInt(64), 0);
+ asserts_.emplace_back(AssertStmt(
+ alloc_size == 0 || align_cond, StringImm("ValueError"),
+ ffi::Array<StringImm>({StringImm("Misaligned Tensor data on argument #"),
+ StringImm(std::to_string(param_index)), when_calling_imm_,
+ sig_imm_, StringImm("`,\n expected data alignment="),
+ StringImm(std::to_string(buffer->data_alignment)),
+ StringImm(" bytes")})));
+ }
+ // mark alignment of external bufs — must be after the alignment assertion
+ // so the compiler does not emit aligned loads before the check fires.
+ asserts_.emplace_back(AttrStmt(vptr, tir::attr::storage_alignment,
+ IntImm(DataType::Int(32), buffer->data_alignment),
+ Evaluate(0)));
+ } else {
+ // Even without alignment check, mark alignment for the compiler.
+ init_nest_.emplace_back(AttrStmt(vptr, tir::attr::storage_alignment,
+ IntImm(DataType::Int(32), buffer->data_alignment),
+ Evaluate(0)));
+ }
+ }
+ }
+}
+
+} // namespace tir
+} // namespace tvm
diff --git a/src/tir/transform/tvm_ffi_binder.h b/src/tir/transform/tvm_ffi_binder.h
new file mode 100644
index 0000000000..2daa320087
--- /dev/null
+++ b/src/tir/transform/tvm_ffi_binder.h
@@ -0,0 +1,431 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm_ffi_binder.h
+ * \brief Helper utility to match and bind packed function arguments.
+ *
+ * TVMFFIABIBuilder owns ALL error message generation and type checking for
+ * packed function arguments. The caller (make_packed_api.cc) creates
+ * an TVMFFIABIBuilder and calls DecodeParam(i) for each parameter.
+ */
+#ifndef TVM_TIR_TRANSFORM_TVM_FFI_BINDER_H_
+#define TVM_TIR_TRANSFORM_TVM_FFI_BINDER_H_
+
+#include <tvm/arith/analyzer.h>
+#include <tvm/ffi/reflection/access_path.h>
+#include <tvm/tir/buffer.h>
+#include <tvm/tir/expr.h>
+#include <tvm/tir/stmt.h>
+
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace tir {
+
+/*!
+ * \brief Helper utility to generate match and bind of packed function arguments.
+ *
+ * TVMFFIABIBuilder consolidates all binding logic, type checking, and error message
+ * generation for packed function parameters. The primary public method is
+ * DecodeAllParams(), which handles everything: type index extraction,
+ * type checking (TypeError), value loading, scalar binding, buffer
+ * binding, and rich error message generation with AccessPath.
+ *
+ * ## Generated statement ordering
+ *
+ * Statements are separated into three sequences to ensure correct variable
+ * scoping. Symbolic shape variables (e.g. batch_size) may appear in an
+ * earlier buffer's shape expression (batch_size + 1) before being defined
+ * by a later buffer's shape (batch_size). Separating definitions from
+ * checks guarantees all variables are in scope when assertions reference them.
+ *
+ * - init_nest: LetStmts, DeclBuffers for shape/strides arrays, AttrStmts —
+ * all value-loading code that defines variables.
+ * - asserts: AssertStmts — all validation checks.
+ * - decl_buffers: DeclBuffer for buffer_map entries — buffer declarations.
+ *
+ * ## Calling Protocol
+ *
+ * 1. Construct with function metadata (func_name, params, buffer_map, v_packed_args,
+ * v_num_packed_args). The constructor emits arg count and null-pointer checks.
+ * 2. Call DecodeAllParams(device_type, device_id)
+ * - Decodes, type-checks, and binds all packed arguments
+ * - Binds DLTensor buffers (shape, strides, dtype, device checks)
+ * 3. Call Finalize() to get the three statement sequences and definition map.
+ *
+ * \note Consider a function f(tA(shape=var(n)), tB(shape=3), tC(shape=(n+2)).
+ * Here n is an undefined variable decided by the outside, tB imposes
+ * a constraint such that it can only take tensor with shape 3, tC imposes
+ * another constraint that its shape must equal n + 2.
+ * So if we call it with f(bufferA, bufferB, bufferC), we need to generate
+ * the following binding sequence:
+ * - define n = bufferA.shape[0]
+ * - assert bufferB.shape[0] == 3
+ * - assert bufferB.shape[1] == n + 3
+ */
+class TVMFFIABIBuilder {
+ public:
+ /*! \brief Variable definition info: bound value and the AccessPath where first defined. */
+ struct VarDefInfo {
+ PrimExpr value;
+ ffi::reflection::AccessPath first_def_path;
+ };
+
+ /*! \brief Result of Finalize(): definition map and three statement sequences. */
+ struct Result {
+ /*! \brief Var -> VarDefInfo map for defined variables. */
+ std::unordered_map<const VarNode*, VarDefInfo> var_defs;
+ /*! \brief Variable definitions (LetStmts, shape/strides DeclBuffers, AttrStmts). */
+ std::vector<Stmt> init_nest;
+ /*! \brief Validation checks (all AssertStmts). */
+ std::vector<Stmt> asserts;
+ /*! \brief Buffer declarations for buffer_map entries. */
+ std::vector<Stmt> decl_buffers;
+ };
+
+ /*!
+ * \brief Constructor with function signature info for rich error messages.
+ *
+ * Builds a human-readable function signature string, records parameter
+ * name mappings, and emits argument count and null-pointer checks into init_nest_.
+ *
+ * \param func_name The function name.
+ * \param params The function parameters.
+ * \param buffer_map The buffer map from parameters to buffers.
+ * \param v_packed_args The packed args variable (used for struct_get calls).
+ * \param v_num_packed_args The variable holding the actual number of packed args.
+ * \param device_type The expected device type expression.
+ * \param device_id The device id variable (may be defined during buffer binding).
+ */
+ TVMFFIABIBuilder(const ffi::String& func_name, const ffi::Array<Var>& params,
+ const ffi::Map<Var, Buffer>& buffer_map, const Var& v_packed_args,
+ const Var& v_num_packed_args, const PrimExpr& device_type,
+ const PrimExpr& device_id);
+
+ /*!
+ * \brief Decode all packed arguments: type-check, load values, bind buffers.
+ *
+ * This is the primary public method after construction. It:
+ * 1. Calls DecodeParam(i) for each parameter to type-check, load values,
+ * and bind scalar params
+ * 2. Binds DLTensor buffers to handles (shape, strides, dtype, device checks)
+ *
+ * Statements are routed to init_nest_, asserts_, or decl_buffers_.
+ */
+ void DecodeAllParams();
+
+ /*!
+ * \brief Finalize and consume the binder, returning generated statements and definitions.
+ *
+ * Renders any deferred assertions (with display-var substitution for
+ * human-readable error messages), then moves all data out of the binder.
+ * The binder should not be used after this call.
+ *
+ * \return A Result containing def_map, init_nest, asserts, decl_buffers.
+ */
+ Result Finalize();
+
+ private:
+ // ── Assert helpers ────────────────────────────────────────────
+
+ /*! \brief Convert various string types to StringImm for EmitAssert. */
+ static StringImm ToMsgPart(StringImm s) { return s; }
+ static StringImm ToMsgPart(const char* s) { return StringImm(s); }
+ static StringImm ToMsgPart(std::string s) { return StringImm(std::move(s)); }
+ static StringImm ToMsgPart(ffi::String s) { return StringImm(std::move(s)); }
+
+ /*!
+ * \brief Emit an assertion into init_nest_ with auto-converted message parts.
+ *
+ * These assertions check per-field properties (type index, ndim, dtype, device)
+ * and must execute before subsequent value loads from the same argument.
+ * Only BindScalar value-match and BindCompactStrides assertions are deferred
+ * to asserts_ (they may reference symbolic vars with forward definitions).
+ *
+ * Each variadic argument is converted to StringImm automatically.
+ * Accepts StringImm, const char*, std::string, or ffi::String.
+ */
+ template <typename... Args>
+ void EmitAssert(const PrimExpr& cond, const char* error_kind, Args&&... args) {
+ ffi::Array<StringImm> parts;
+ (parts.push_back(ToMsgPart(std::forward<Args>(args))), ...);
+ init_nest_.emplace_back(AssertStmt(cond, StringImm(error_kind), parts));
+ }
+
+ // ── Binding submethods ─────────────────────────────────────────
+
+ /*! \brief Type-check, load value, and bind one packed param. */
+ void DecodeParam(int param_index);
+
+ /*! \brief Load the i-th packed argument as the given type from the union value. */
+ static PrimExpr LoadTVMFFIAnyUnionValue(const Var& v_packed_args, int param_index,
+ DataType arg_type);
+
+ // ── Per-dtype type-check + value-load methods ──────────────────
+ //
+ // Each method combines type checking (emitting TypeError on mismatch)
+ // with value loading from the packed args array. DecodeParam
+ // dispatches to one of these based on the parameter dtype.
+
+ /*!
+ * \brief Type-check and load a handle argument (DLTensor or opaque pointer).
+ * \param param_index Parameter index.
+ * \param type_index The variable holding the FFI type index.
+ * \return The loaded argument value.
+ */
+ PrimExpr DecodeParamOpaqueHandle(int param_index, const Var& type_index);
+
+ /*!
+ * \brief Type-check and load a boolean argument.
+ * \param param_index Parameter index.
+ * \param type_index The variable holding the FFI type index.
+ * \return The loaded argument value.
+ */
+ PrimExpr DecodeParamBool(int param_index, const Var& type_index);
+
+ /*!
+ * \brief Type-check and load an integer argument.
+ * \param param_index Parameter index.
+ * \param type_index The variable holding the FFI type index.
+ * \param dtype The expected data type for this parameter.
+ * \return The loaded argument value.
+ */
+ PrimExpr DecodeParamInt(int param_index, const Var& type_index, DataType dtype);
+
+ /*!
+ * \brief Type-check and load a float argument.
+ * \param param_index Parameter index.
+ * \param type_index The variable holding the FFI type index.
+ * \param dtype The expected data type for this parameter.
+ * \return The loaded argument value.
+ */
+ PrimExpr DecodeParamFloat(int param_index, const Var& type_index, DataType dtype);
+
+ // ── Private binding submethods (all take AccessPath) ───────────
+
+ /*!
+ * \brief Internal scalar bind with AccessPath tracking and rich error messages.
+ *
+ * Binds \p arg to \p value. If arg is a Var not yet in var_defs_, creates a
+ * new definition (LetStmt to init_nest_); otherwise emits a rich assertion
+ * (to asserts_) that the existing value matches the new one.
+ *
+ * When arg is a non-Var expression (e.g. batch_size + 1), the assertion is
+ * deferred to Finalize() so display-var substitution can render the expression
+ * using AccessPath names (e.g. "k.shape[0] + 1" instead of "batch_size + 1").
+ *
+ * \param arg The argument expression to bind (typically a Var or constant).
+ * \param value The value expression to bind to the argument.
+ * \param with_lets If true, emit LetStmt bindings into init_nest_.
+ * \param path AccessPath for rich error message rendering.
+ * \return True if this was the first bind (definition created), false otherwise.
+ */
+ bool BindScalar(const PrimExpr& arg, const PrimExpr& value,
+ const ffi::reflection::AccessPath& path, bool with_lets);
+
+ /*!
+ * \brief Array bind: binds element-wise with AccessPath[k] for each element.
+ *
+ * \param arg The expected array of expressions.
+ * \param value The actual array of expressions to bind against.
+ * \param base_path Base AccessPath; each element appends ArrayItem(k).
+ */
+ void BindArray(const ffi::Array<PrimExpr>& arg, const ffi::Array<PrimExpr>& value,
+ const ffi::reflection::AccessPath& base_path);
+
+ /*!
+ * \brief Buffer-to-buffer bind with AccessPath.
+ *
+ * Binds data, elem_offset, shape, and strides of \p arg against \p value,
+ * emitting assertions for any mismatches.
+ *
+ * \param arg The expected buffer definition.
+ * \param value The actual buffer to bind against.
+ * \param base_path Base AccessPath for the buffer parameter.
+ * \param fuzzy_match If true, allow value to have more dimensions than arg.
+ */
+ void BindBuffer(const Buffer& arg, const Buffer& value, ffi::reflection::AccessPath base_path,
+ bool fuzzy_match);
+
+ /*!
+ * \brief DLTensor bind: ndim/dtype/shape/strides/data/device assertions.
+ *
+ * \param buffer The buffer definition to bind against.
+ * \param device_type The expected device type expression.
+ * \param device_id The expected device id expression.
+ * \param handle The variable holding the DLTensor handle.
+ * \param arg_name Human-readable name for error messages.
+ * \param base_path Base AccessPath for the tensor parameter.
+ */
+ void DecodeParamDLTensor(const Buffer& buffer, const PrimExpr& device_type,
+ const PrimExpr& device_id, const Var& handle,
+ const std::string& arg_name, ffi::reflection::AccessPath base_path);
+
+ // ── DLTensor sub-helpers ───────────────────────────────────────
+
+ /*!
+ * \brief Get a DLTensor field pointer (shape or strides) and store it in a LetStmt.
+ *
+ * \param handle The DLTensor handle variable.
+ * \param field_kind kDLTensorShape or kDLTensorStrides.
+ * \param var_name Name for the pointer variable.
+ * \return A Var holding the field pointer.
+ */
+ Var DLTensorGetFieldPtr(const Var& handle, int field_kind, const std::string& var_name);
+
+ /*!
+ * \brief Load an element from a shape/strides int64 array pointer.
+ *
+ * Uses tvm_struct_get with kInt64ArrayElem to access ptr[index].
+ *
+ * \param ptr The int64 array pointer variable (from DLTensorGetFieldPtr).
+ * \param index The element index to load.
+ * \return The loaded int64 value expression.
+ */
+ static PrimExpr LoadInt64ArrayElem(const Var& ptr, int index);
+
+ /*!
+ * \brief Assert strides form a compact (C-contiguous) layout.
+ *
+ * \param buffer The expected buffer definition.
+ * \param strides_ptr The strides pointer variable.
+ * \param v_strides_is_null Expression checking if strides pointer is NULL.
+ * \param param_path AccessPath for the tensor parameter.
+ */
+ void BindCompactStrides(const Buffer& buffer, const Var& strides_ptr,
+ const PrimExpr& v_strides_is_null,
+ const ffi::reflection::AccessPath& param_path);
+
+ /*!
+ * \brief Bind strides for auto-broadcast buffers: stride=0 for shape==1 dims.
+ *
+ * \param buffer The expected buffer definition.
+ * \param strides_ptr The strides pointer variable.
+ * \param v_strides_is_null Expression checking if strides pointer is NULL.
+ * \param param_path AccessPath for the tensor parameter.
+ */
+ void BindAutoBroadcastStrides(const Buffer& buffer, const Var& strides_ptr,
+ const PrimExpr& v_strides_is_null,
+ const ffi::reflection::AccessPath& param_path);
+
+ /*!
+ * \brief Bind strides with C-contiguous fallback when strides pointer is NULL.
+ *
+ * \param buffer The expected buffer definition.
+ * \param strides_ptr The strides pointer variable.
+ * \param shape_ptr The shape pointer variable (for computing C-contiguous strides).
+ * \param v_strides_is_null Expression checking if strides pointer is NULL.
+ * \param param_path AccessPath for the tensor parameter.
+ */
+ void BindRegularStrides(const Buffer& buffer, const Var& strides_ptr, const Var& shape_ptr,
+ const PrimExpr& v_strides_is_null,
+ const ffi::reflection::AccessPath& param_path);
+
+ // ── Error message helpers ──────────────────────────────────────
+
+ /*!
+ * \brief Emit a type-mismatch TypeError for a given parameter.
+ *
+ * Emits an assertion with kind "TypeError" and message parts:
+ * "Mismatched type on argument #N when calling:\n `<sig>`,\n expected <type>"
+ *
+ * \param param_index The parameter index.
+ * \param cond The boolean condition; assertion fails when cond is false.
+ * \param expected_type Human-readable expected type (e.g. "Tensor", "int").
+ */
+ void EmitTypeIndexCheck(int param_index, const PrimExpr& cond, const std::string& expected_type);
+
+ /*!
+ * \brief Render an AccessPath as a human-readable string (e.g. "a.shape[0]").
+ * \param path The AccessPath to render.
+ * \return A human-readable string representation of the path.
+ */
+ ffi::String RenderAccessPath(const ffi::reflection::AccessPath& path) const;
+
+ /*!
+ * \brief Extract param_index from the root ArrayItem step of a path.
+ * \param path The AccessPath to extract the index from.
+ * \return The param index, or -1 if not found.
+ */
+ int GetParamIndex(const ffi::reflection::AccessPath& path) const;
+
+ /*!
+ * \brief Render pending constant-expression assertions with display-var substitution.
+ *
+ * For each pending assertion, substitutes known variable names with their
+ * AccessPath-rendered names (e.g. batch_size → "k.shape[0]") so error messages
+ * show human-readable expressions like "k.shape[0] + 1" instead of "batch_size + 1".
+ */
+ void RenderPendingAsserts();
+
+ // ── Data members ───────────────────────────────────────────────
+
+ /*! \brief Pending constant-expression assertion: deferred until Finalize() for display-var
+ * substitution. */
+ struct PendingConstAssert {
+ PrimExpr condition;
+ ffi::reflection::AccessPath path;
+ PrimExpr expected_expr;
+ };
+
+ /*! \brief The definition map: VarNode* -> VarDefInfo (value + first_def_path). */
+ std::unordered_map<const VarNode*, VarDefInfo> var_defs_;
+ /*! \brief Variable definitions: LetStmts, shape/strides DeclBuffers, AttrStmts. */
+ std::vector<Stmt> init_nest_;
+ /*! \brief Validation checks: all AssertStmts. */
+ std::vector<Stmt> asserts_;
+ /*! \brief Buffer declarations for buffer_map entries. */
+ std::vector<Stmt> decl_buffers_;
+ /*! \brief Deferred constant-expression assertions for display-var substitution. */
+ std::vector<PendingConstAssert> pending_const_asserts_;
+ /*! \brief internal analyzer. */
+ arith::Analyzer analyzer_;
+
+ // Function metadata
+ /*! \brief function name for error messages. */
+ ffi::String func_name_;
+ /*! \brief function signature string for error messages. */
+ std::string func_signature_;
+ /*! \brief The function parameters. */
+ ffi::Array<Var> params_;
+ /*! \brief The buffer map from parameters to buffers. */
+ ffi::Map<Var, Buffer> buffer_map_;
+ /*! \brief The packed args variable. */
+ Var v_packed_args_;
+ /*! \brief The expected device type expression. */
+ PrimExpr device_type_;
+ /*! \brief The device id variable. */
+ PrimExpr device_id_;
+ /*! \brief Map from param_index to param_name for AccessPath rendering. */
+ std::unordered_map<int, std::string> param_names_;
+
+ // Pre-cached common message fragments for string sharing across assertions
+ StringImm sig_imm_; // func_signature_ (set in constructor)
+ StringImm when_calling_imm_ = StringImm(" when calling:\n `");
+ /*! \brief Whether to emit data pointer alignment checks (disabled for now). */
+ bool check_alignment_ = false;
+};
+
+} // namespace tir
+} // namespace tvm
+#endif // TVM_TIR_TRANSFORM_TVM_FFI_BINDER_H_
diff --git a/tests/python/codegen/test_codegen_assert.py b/tests/python/codegen/test_codegen_assert.py
new file mode 100644
index 0000000000..21c3ac5cde
--- /dev/null
+++ b/tests/python/codegen/test_codegen_assert.py
@@ -0,0 +1,124 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""AssertStmt codegen tests: verify kind and message_parts produce correct exceptions."""
+
+import pytest
+
+import tvm
+import tvm.testing
+from tvm.script import tir as T
+
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+def test_assert_runtime_error(codegen_target):
+ """AssertStmt with RuntimeError kind produces RuntimeError."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("RuntimeError", ["Expected non-null input"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(RuntimeError, match="Expected non-null input"):
+ lib(0)
+
+
+def test_assert_value_error(codegen_target):
+ """AssertStmt with ValueError kind produces ValueError."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("ValueError", ["Shape mismatch: expected 4 got 8"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(ValueError, match="Shape mismatch"):
+ lib(0)
+
+
+def test_assert_type_error(codegen_target):
+ """AssertStmt with TypeError kind produces TypeError."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("TypeError", ["Expected Tensor but got int"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(TypeError, match="Expected Tensor but got int"):
+ lib(0)
+
+
+def test_assert_multi_part_message(codegen_target):
+ """Multi-part messages are correctly concatenated at runtime."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("ValueError", ["Expected shape ", "4", " but got ", "8"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(ValueError, match="Expected shape 4 but got 8"):
+ lib(0)
+
+
+def test_assert_passing_condition(codegen_target):
+ """Passing assertion does not raise."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("RuntimeError", ["This should not be raised"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ lib(1) # should pass without error
+
+
+def test_assert_many_parts(codegen_target):
+ """Assertion with 8 parts concatenated correctly."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("RuntimeError", ["p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(RuntimeError, match="p0p1p2p3p4p5p6p7"):
+ lib(0)
+
+
+def test_tvmscript_assert_preserves_kind(codegen_target):
+ """Regression: TVMScript structured assert preserves kind at runtime."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("ValueError", ["x must be positive"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(ValueError, match="x must be positive"):
+ lib(0)
+
+
+def test_tvmscript_assert_preserves_parts(codegen_target):
+ """Regression: TVMScript structured assert with separate parts."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("ValueError", ["x must be ", "positive"])
+
+ lib = tvm.compile(func, target=codegen_target)
+ with pytest.raises(ValueError, match="x must be positive"):
+ lib(0)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/codegen/test_codegen_error_handling.py b/tests/python/codegen/test_codegen_error_handling.py
new file mode 100644
index 0000000000..b882c3a795
--- /dev/null
+++ b/tests/python/codegen/test_codegen_error_handling.py
@@ -0,0 +1,466 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Runtime error message tests for MakePackedAPI + TVMFFIABIBuilder.
+
+All tests compile TVMScript functions and verify the correct Python exception
+type and exact error message at runtime.
+"""
+
+import re
+
+import numpy as np
+import pytest
+import tvm_ffi
+
+import tvm
+import tvm.testing
+from tvm.script import tir as T
+
+# Parameterize over both LLVM and C backends
+codegen_target = tvm.testing.parameter("llvm", "c")
+
+
+# ── Argument count errors ────────────────────────────────────
+
+
+def test_wrong_argument_count_error(codegen_target):
+ """Wrong argument count produces TypeError with function signature."""
+
+ @T.prim_func
+ def func(a: T.handle, b: T.handle):
+ n0 = T.int64()
+ A = T.match_buffer(a, (n0,), "float32")
+ B = T.match_buffer(b, (n0,), "float32")
+ for i in range(n0):
+ B[i] = A[i] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a, b) # correct input should pass
+
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Expected 2 arguments when calling:\n"
+ " `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`"
+ ),
+ ):
+ lib()
+
+
+# ── Type mismatch errors (tensor parameters) ────────────────
+
+
+def test_type_mismatch_non_tensor(codegen_target):
+ """Passing a non-tensor where a tensor is expected raises TypeError."""
+
+ @T.prim_func
+ def func(a: T.handle, b: T.handle):
+ n0 = T.int64()
+ A = T.match_buffer(a, (n0,), "float32")
+ B = T.match_buffer(b, (n0,), "float32")
+ for i in range(n0):
+ B[i] = A[i] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a, b) # correct input should pass
+
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched type on argument #1 when calling:\n"
+ " `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+ " expected Tensor"
+ ),
+ ):
+ lib(a, 1)
+
+
+# ── Shape mismatch errors ───────────────────────────────────
+
+
+def test_shape_mismatch_shared_variable(codegen_target):
+ """b has different shape than a when they share symbolic variable n0."""
+
+ @T.prim_func
+ def func(a: T.handle, b: T.handle):
+ n0 = T.int64()
+ A = T.match_buffer(a, (n0,), "float32")
+ B = T.match_buffer(b, (n0,), "float32")
+ for i in range(n0):
+ B[i] = A[i] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a, b) # correct input should pass
+
+ b_short = tvm.runtime.tensor(np.zeros(126, dtype="float32"))
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Mismatched B.shape[0] on argument #1 when calling:\n"
+ " `func(A: Tensor([n0], float32), B: Tensor([n0], float32))`,\n"
+ " expected to match A.shape[0]"
+ ),
+ ):
+ lib(a, b_short)
+
+
+def test_invalid_shape_fixed(codegen_target):
+ """Passing wrong shape for a fixed buffer dimension raises ValueError."""
+
+ @T.prim_func
+ def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+ for i in range(128):
+ b[i] = a[i] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a, b) # correct input should pass
+
+ a_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+ b_wrong = tvm.runtime.tensor(np.zeros(256, dtype="float32"))
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Invalid a.shape[0] on argument #0 when calling:\n"
+ " `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+ " expected 128"
+ ),
+ ):
+ lib(a_wrong, b_wrong)
+
+
+# ── ndim mismatch errors ────────────────────────────────────
+
+
+def test_ndim_mismatch_error(codegen_target):
+ """ndim mismatch produces ValueError with function signature."""
+
+ @T.prim_func
+ def func(a: T.Buffer((4, 8), "float32"), b: T.Buffer((4, 8), "float32")):
+ for i, j in T.grid(4, 8):
+ b[i, j] = a[i, j]
+
+ lib = tvm.compile(func, target=codegen_target)
+ a_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+ b_ok = tvm.runtime.tensor(np.zeros((4, 8), dtype="float32"))
+ lib(a_ok, b_ok) # correct input should pass
+
+ a = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+ b = tvm.runtime.tensor(np.zeros(4, dtype="float32"))
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Mismatched a.ndim on argument #0 when calling:\n"
+ " `func(a: Tensor([4, 8], float32), b: Tensor([4, 8], float32))`,\n"
+ " expected 2"
+ ),
+ ):
+ lib(a, b)
+
+
+# ── dtype mismatch errors ───────────────────────────────────
+
+
+def test_dtype_mismatch_error(codegen_target):
+ """dtype mismatch produces TypeError with function signature."""
+
+ @T.prim_func
+ def func(a: T.Buffer((8,), "float32"), b: T.Buffer((8,), "float32")):
+ for i in range(8):
+ b[i] = a[i]
+
+ lib = tvm.compile(func, target=codegen_target)
+ a_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ b_ok = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ lib(a_ok, b_ok) # correct input should pass
+
+ a = tvm.runtime.tensor(np.zeros(8, dtype="int32"))
+ b = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched a.dtype on argument #0 when calling:\n"
+ " `func(a: Tensor([8], float32), b: Tensor([8], float32))`,\n"
+ " expected float32"
+ ),
+ ):
+ lib(a, b)
+
+
+# ── Data alignment errors ──────────────────────────────────
+
+
+@pytest.mark.skip(reason="alignment check disabled for now, revisit after merge")
+def test_data_alignment_error(codegen_target):
+ """Misaligned buffer data pointer raises ValueError."""
+
+ @T.prim_func
+ def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+ for i in range(128):
+ b[i] = a[i] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a_ok, b_ok) # correct input should pass
+
+ # Slice off first element of a 129-element array to create misaligned data pointer
+ np_arr = np.zeros(129, dtype="float32")
+ a_misaligned = tvm_ffi.from_dlpack(np_arr[1:])
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Misaligned Tensor data on argument #0 when calling:\n"
+ " `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+ " expected data alignment=64 bytes"
+ ),
+ ):
+ lib(a_misaligned, b)
+
+
+# ── Compact strides mismatch errors ────────────────────────
+
+
+def test_strides_mismatch_transposed(codegen_target):
+ """Transposed (non-compact) strides raise ValueError."""
+
+ @T.prim_func
+ def func(a: T.Buffer((128, 128), "float32"), b: T.Buffer((128, 128), "float32")):
+ for i, j in T.grid(128, 128):
+ b[i, j] = a[i, j] + T.float32(1)
+
+ lib = tvm.compile(func, target=codegen_target)
+ a_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+ b_ok = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+ lib(a_ok, b_ok) # correct input should pass
+
+ # Use Fortran-order array to get non-compact (non-C-contiguous) strides
+ np_arr = np.asfortranarray(np.zeros((128, 128), dtype="float32"))
+ a_transposed = tvm_ffi.from_dlpack(np_arr)
+ b = tvm.runtime.tensor(np.zeros((128, 128), dtype="float32"))
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Mismatched a.strides on argument #0 when calling:\n"
+ " `func(a: Tensor([128, 128], float32), b: Tensor([128, 128], float32))`,\n"
+ " expected to be compact array"
+ ),
+ ):
+ lib(a_transposed, b)
+
+
+# ── Device mismatch errors ─────────────────────────────────
+
+
+@tvm.testing.requires_cuda
+def test_device_mismatch_error():
+ """Passing GPU tensor to CPU function raises ValueError."""
+
+ @T.prim_func
+ def func(a: T.Buffer((128,), "float32"), b: T.Buffer((128,), "float32")):
+ for i in range(128):
+ b[i] = a[i] + T.float32(1)
+
+ lib = tvm.compile(func, target="llvm")
+ a_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ b_ok = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+ lib(a_ok, b_ok) # correct input should pass
+
+ a_gpu = tvm.runtime.tensor(np.zeros(128, dtype="float32"), device=tvm.cuda(0))
+ b = tvm.runtime.tensor(np.zeros(128, dtype="float32"))
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Mismatched a.device_type on argument #0 when calling:\n"
+ " `func(a: Tensor([128], float32), b: Tensor([128], float32))`,\n"
+ " expected cpu"
+ ),
+ ):
+ lib(a_gpu, b)
+
+
+# ── Scalar type mismatch errors ─────────────────────────────
+
+
+def test_type_mismatch_int_parameter(codegen_target):
+ """Passing a tensor where an int is expected raises TypeError."""
+
+ @T.prim_func
+ def func(x: T.int32) -> T.int32:
+ if x > 0:
+ return 10
+ else:
+ return 20
+
+ lib = tvm.compile(func, target=codegen_target)
+ assert lib(5) == 10 # correct input should pass
+
+ a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched type on argument #0 when calling:\n `func(x: int32)`,\n expected int"
+ ),
+ ):
+ lib(a)
+
+
+def test_type_mismatch_float_parameter(codegen_target):
+ """Passing a tensor where a float is expected raises TypeError."""
+
+ @T.prim_func
+ def func(x: T.float32) -> T.int32:
+ if x > T.float32(0):
+ return 1
+ else:
+ return 0
+
+ lib = tvm.compile(func, target=codegen_target)
+ assert lib(1.0) == 1 # correct input should pass
+
+ a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched type on argument #0 when calling:\n `func(x: float32)`,\n expected float"
+ ),
+ ):
+ lib(a)
+
+
+def test_type_mismatch_bool_parameter(codegen_target):
+ """Passing a tensor where a bool is expected raises TypeError."""
+
+ @T.prim_func
+ def func(x: T.bool) -> T.int32:
+ if x:
+ return 1
+ else:
+ return 0
+
+ lib = tvm.compile(func, target=codegen_target)
+ assert lib(True) == 1 # correct input should pass
+
+ a = tvm.runtime.tensor(np.zeros(8, dtype="float32"))
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched type on argument #0 when calling:\n `func(x: bool)`,\n expected boolean"
+ ),
+ ):
+ lib(a)
+
+
+# ── Forward-reference symbolic variable ────────────────────
+
+
+def test_forward_reference_symbolic_shape(codegen_target):
+ """Buffers sharing a symbolic var with forward reference compile and run correctly.
+
+ When buffer A has shape (batch_size+1,) and buffer B has shape (batch_size,),
+ batch_size is referenced in A's shape assertion before it is defined from B.
+ The three-sequence separation ensures this works. Also verifies the error
+ message uses rendered access paths (e.g. "B.shape[0] + 1") for shape checks.
+ """
+
+ @T.prim_func
+ def func(a: T.handle, b: T.handle):
+ batch_size = T.int64()
+ A = T.match_buffer(a, (batch_size + 1,), "int32")
+ B = T.match_buffer(b, (batch_size,), "int32")
+ for i in range(batch_size):
+ B[i] = A[i] + A[i + 1]
+
+ lib = tvm.compile(func, target=codegen_target)
+ # Correct inputs: A has shape (5,), B has shape (4,)
+ a = tvm.runtime.tensor(np.array([1, 2, 3, 4, 5], dtype="int32"))
+ b = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+ lib(a, b)
+ np.testing.assert_array_equal(b.numpy(), [3, 5, 7, 9])
+
+ # Wrong shape: A has shape (10,) but B has shape (4,), so batch_size=4 but A needs 5
+ a_wrong = tvm.runtime.tensor(np.zeros(10, dtype="int32"))
+ b_ok = tvm.runtime.tensor(np.zeros(4, dtype="int32"))
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Invalid A.shape[0] on argument #0 when calling:\n"
+ " `func(A: Tensor([batch_size + T.int64(1)], int32),"
+ " B: Tensor([batch_size], int32))`,\n"
+ " expected B.shape[0] + 1"
+ ),
+ ):
+ lib(a_wrong, b_ok)
+
+
+# ── Mixed parameter type errors ────────────────────────────
+
+
+def test_invalid_arguments_mixed_params(codegen_target):
+ """Mixed bool + tensor function: type, dtype, and shape errors."""
+
+ @T.prim_func
+ def func(a0: T.bool, a1: T.Buffer([10], "float32")) -> T.int32:
+ return 0
+
+ lib = tvm.compile(func, target=codegen_target)
+ lib(True, tvm.runtime.tensor(np.zeros(10, dtype="float32"))) # correct input should pass
+
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched type on argument #1 when calling:\n"
+ " `func(a0: bool, a1: Tensor([10], float32))`,\n"
+ " expected Tensor"
+ ),
+ ):
+ lib(1, 1)
+
+ with pytest.raises(
+ TypeError,
+ match=re.escape(
+ "Mismatched a1.dtype on argument #1 when calling:\n"
+ " `func(a0: bool, a1: Tensor([10], float32))`,\n"
+ " expected float32"
+ ),
+ ):
+ lib(1, tvm.runtime.empty([10], "int32"))
+
+ with pytest.raises(
+ ValueError,
+ match=re.escape(
+ "Invalid a1.shape[0] on argument #1 when calling:\n"
+ " `func(a0: bool, a1: Tensor([10], float32))`,\n"
+ " expected 10"
+ ),
+ ):
+ lib(False, tvm.runtime.empty([11], "float32"))
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index c91b5f5bf2..57544ccef1 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -1004,31 +1004,28 @@ def test_llvm_target_attributes():
llvm_ir_lines = llvm_ir.split("\n")
attribute_definitions = dict()
- attributes_with_target = dict()
- functions_with_target = []
+ function_attr_map = dict()
for line in llvm_ir_lines:
func_def = re.match(
"define.* @(?P<func_name>[^(]*)[(].* #(?P<attr_num>[0-9]+) (!.* |){$", line
)
if func_def:
- functions_with_target.append(func_def.group("func_name"))
- attributes_with_target[func_def.group("attr_num")] = True
+ function_attr_map[func_def.group("func_name")] = func_def.group("attr_num")
continue
attr_def = re.match("attributes #(?P<attr_num>[0-9]+) = {(?P<attr_list>.*)}", line)
if attr_def:
attribute_definitions[attr_def.group("attr_num")] = attr_def.group("attr_list")
- for k in list(attributes_with_target.keys()):
- assert re.match('.*"target-cpu"="skylake".*', attribute_definitions[k])
- assert re.match('.*"target-features"=".*[+]avx512f.*".*', attribute_definitions[k])
-
expected_functions = [
"__tvm_ffi_test_func",
"__tvm_parallel_lambda",
]
for n in expected_functions:
- assert n in functions_with_target
+ assert n in function_attr_map, f"Expected function {n} not found in LLVM IR"
+ attr_num = function_attr_map[n]
+ assert re.match('.*"target-cpu"="skylake".*', attribute_definitions[attr_num])
+ assert re.match('.*"target-features"=".*[+]avx512f.*".*', attribute_definitions[attr_num])
@tvm.testing.requires_llvm
@@ -1262,26 +1259,5 @@ def test_bool_return_value():
assert not built(15)
-def test_invalid_arguments():
- """Integers may be passed to functions accepting bool"""
-
- @I.ir_module
- class Module:
- @T.prim_func
- def main(a0: T.bool, a1: T.Buffer([10], "float32")) -> T.int32:
- T.func_attr({"target": T.target("llvm")})
- return 0
-
- built = tvm.compile(Module)
- with pytest.raises(RuntimeError):
- built(1, 1)
-
- with pytest.raises(RuntimeError):
- built(1, tvm.runtime.empty([10], "int32"))
-
- with pytest.raises(RuntimeError):
- built(False, tvm.runtime.empty([11], "float32"))
-
-
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_analysis_well_formed.py b/tests/python/relax/test_analysis_well_formed.py
index 9ce65efe48..00ee1f250a 100644
--- a/tests/python/relax/test_analysis_well_formed.py
+++ b/tests/python/relax/test_analysis_well_formed.py
@@ -685,16 +685,16 @@ def test_pass_dltensor_arg_to_tir():
T.func_attr({"tir.is_scheduled": True, "tir.is_host_func": True})
# From #include <tvm/tir/builtin.h>
- kArrTypeCode = T.meta_var(5)
- kArrTypeBits = T.meta_var(6)
- kArrTypeLanes = T.meta_var(7)
+ kDLTensorTypeCode = T.meta_var(5)
+ kDLTensorTypeBits = T.meta_var(6)
+ kDLTensorTypeLanes = T.meta_var(7)
# From #include <dlpack/dlpack.h>
kDLBfloat = T.meta_var(4)
- type_code = T.tvm_struct_get(tensor, 0, kArrTypeCode, dtype="uint8")
- type_bits = T.tvm_struct_get(tensor, 0, kArrTypeBits, dtype="uint8")
- type_lanes = T.tvm_struct_get(tensor, 0, kArrTypeLanes, dtype="uint16")
+ type_code = T.tvm_struct_get(tensor, 0, kDLTensorTypeCode, dtype="uint8")
+ type_bits = T.tvm_struct_get(tensor, 0, kDLTensorTypeBits, dtype="uint8")
+ type_lanes = T.tvm_struct_get(tensor, 0, kDLTensorTypeLanes, dtype="uint16")
is_bfloat16: T.bool = (
(type_code == kDLBfloat) and (type_bits == 16) and (type_lanes == 1)
diff --git a/tests/python/tir-base/test_tir_constructor.py b/tests/python/tir-base/test_tir_constructor.py
index d05201a143..fa3881f23a 100644
--- a/tests/python/tir-base/test_tir_constructor.py
+++ b/tests/python/tir-base/test_tir_constructor.py
@@ -141,8 +141,15 @@ def test_stmt_constructor():
assert isinstance(x, tvm.tir.AttrStmt)
assert x.value.value == 1
- x = tvm.tir.AssertStmt(tvm.tir.const(1, "bool"), tvm.runtime.convert("hellow"))
+ x = tvm.tir.AssertStmt(
+ tvm.tir.const(1, "bool"),
+ tvm.tir.StringImm("RuntimeError"),
+ [tvm.tir.StringImm("hellow")],
+ )
assert isinstance(x, tvm.tir.AssertStmt)
+ assert x.error_kind.value == "RuntimeError"
+ assert len(x.message_parts) == 1
+ assert x.message_parts[0].value == "hellow"
x = tvm.tir.For(tvm.tir.Var("x", "int32"), 0, 10, tvm.tir.ForKind.SERIAL, nop)
assert isinstance(x, tvm.tir.For)
diff --git a/tests/python/tir-transform/test_tir_transform_make_packed_api.py b/tests/python/tir-transform/test_tir_transform_make_packed_api.py
index 12951a69fb..0d39da0d9e 100644
--- a/tests/python/tir-transform/test_tir_transform_make_packed_api.py
+++ b/tests/python/tir-transform/test_tir_transform_make_packed_api.py
@@ -14,8 +14,12 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+"""Tests for tir.transform.MakePackedAPI TIR transform.
+
+Tests verify the transform output using TVMScript before/after patterns.
+Runtime error tests are in tests/python/codegen/test_codegen_error_handling.py.
+"""
-import numpy as np
import pytest
import tvm
@@ -25,31 +29,6 @@ from tvm.script import ir as I
from tvm.script import tir as T
-def _find_assignment(stmt, var_name):
- while not isinstance(stmt, tvm.tir.LetStmt):
- stmt = stmt.body
-
- if stmt.var.name != var_name:
- return _find_assignment(stmt.body, var_name)
-
- return stmt
-
-
-def _find_next(stmt, type):
- search_stack = [stmt]
-
- while search_stack:
- stmt = search_stack.pop()
- if isinstance(stmt, type):
- return stmt
- elif isinstance(stmt, tvm.tir.SeqStmt):
- search_stack.extend(reversed(stmt))
- else:
- search_stack.append(stmt.body)
-
- return None
-
-
def _find_compute_scope(func):
result = None
@@ -132,7 +111,7 @@ def test_internal_subroutine_call():
compute_scope = _find_compute_scope(after["main"])
subroutine_call_op = compute_scope.body.value.op
assert isinstance(subroutine_call_op, tvm.ir.GlobalVar), (
- f"The main function's CallNode should use the subroutine's GLobalVar as the operation, "
+ f"The main function's CallNode should use the subroutine's GlobalVar as the operation, "
f"but instead has an operation of type {subroutine_call_op}"
)
@@ -175,73 +154,8 @@ def test_subroutine_call_to_externally_visible_subroutine():
)
-def test_function_call_with_wrong_argument_count():
- """Argument counts must be checked before accessing the type codes"""
-
- @T.prim_func
- def func(
- A: T.Buffer([16, 16], "int32"),
- B: T.Buffer([16, 16], "int32"),
- C: T.Buffer([16, 16], "int32"),
- D: T.Buffer([16, 16], "int32"),
- ):
- pass
-
- built = tvm.compile(func, target="llvm")
-
- with pytest.raises(tvm.TVMError):
- built()
-
-
-def test_function_call_with_wrong_type_code():
- """Type codes must be checked before accessing the arguments"""
-
- @T.prim_func
- def func(A: T.Buffer([16, 16], "int32")):
- pass
-
- built = tvm.compile(func, target="llvm")
-
- with pytest.raises(tvm.TVMError):
- built(0)
-
-
-def test_function_call_with_null_data_pointer():
- """The data pointer must be checked before accessing the array"""
-
- @T.prim_func
- def func(A: T.Buffer([16, 16], "int32"), B: T.Buffer([16, 16], "int32")):
- for i, j in T.grid(16, 16):
- B[i, j] = A[i, j]
-
- built = tvm.compile(func, target="llvm")
-
- A = tvm.runtime.tensor(np.zeros([16], dtype="int32"))
- B = tvm.runtime.empty([16, 16], "int32", tvm.cpu())
-
- with pytest.raises(tvm.TVMError):
- built(A, B)
-
-
-def test_function_call_with_wrong_dimensionality():
- """The dimensionality must be checked before validating the shape"""
-
- @T.prim_func
- def func(A: T.Buffer([16, 16], "int32"), B: T.Buffer([16, 16], "int32")):
- for i, j in T.grid(16, 16):
- B[i, j] = A[i, j]
-
- built = tvm.compile(func, target="llvm")
-
- A = tvm.runtime.tensor(np.zeros([16], dtype="int32"))
- B = tvm.runtime.empty([16], "int32", tvm.cpu())
-
- with pytest.raises(tvm.TVMError):
- built(A, B)
-
-
def test_zero_arg_function():
- """Only check non-null args when num_args>0"""
+ """Zero-arg function emits num_args check but no null-pointer check."""
@I.ir_module
class Before:
@@ -254,19 +168,22 @@ def test_zero_arg_function():
class Expected:
@T.prim_func
def func_without_arg(
- self: T.handle,
+ self_handle: T.handle,
args: T.handle,
num_args: T.int32,
- result: T.handle("void"),
+ result: T.handle("void", "global"),
) -> T.int32:
T.func_attr(
{
"calling_conv": 1,
- "target": T.target("llvm"),
"global_symbol": "__tvm_ffi_func_without_arg",
+ "target": T.target("llvm"),
}
)
- assert num_args == 0, "func_without_arg: num_args should be 0"
+ assert num_args == 0, (
+ "TypeError",
+ ["Expected ", "0", " arguments", " when calling:\n `", "func_without_arg()", "`"],
+ )
with T.attr(0, "compute_scope", "func_without_arg_compute_"):
T.tvm_struct_set(result, 0, 13, 1)
T.tvm_struct_set(result, 0, 14, 0)
@@ -275,25 +192,11 @@ def test_zero_arg_function():
return 0
After = tvm.tir.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(Expected, After)
+ tvm.ir.assert_structural_equal(After, Expected)
def test_int_parameter():
- """Boolean may be passed to functions accepting int
-
- A PackedFunc produced by compiling an IRModule should support the
- same type conversions as the C++ implementation. When a function
- accepts an integer argument, the caller may call it with a boolean
- value.
-
- This also provides backwards compatibility for functions that were
- defined as accepting an integer, but are called with a boolean
- argument. Prior to PackedFunc interface supporting boolean
- arguments directly, the argument would be converted from boolean
- to integer to be stored in a TVMValue. After adding support for
- boolean arguments, this usage should not cause an error.
-
- """
+ """Int parameter emits type check accepting int or bool."""
@I.ir_module
class Before:
@@ -309,22 +212,38 @@ def test_int_parameter():
class Expected:
@T.prim_func
def main(
- self: T.handle,
+ self_handle: T.handle,
args: T.handle,
num_args: T.int32,
- result: T.handle("void"),
+ result: T.handle("void", "global"),
) -> T.int32:
T.func_attr(
{
"calling_conv": 1,
- "target": T.target("llvm"),
"global_symbol": "__tvm_ffi_main",
+ "target": T.target("llvm"),
}
)
- assert num_args == 1, "main: num_args should be 1"
- assert not T.isnullptr(args), "main: args pointer is NULL"
+ assert num_args == 1, (
+ "TypeError",
+ ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: int32)", "`"],
+ )
+ assert not T.isnullptr(args), (
+ "TypeError",
+ ["args pointer is NULL", " when calling:\n `", "main(arg: int32)", "`"],
+ )
arg_type_index: T.int32 = T.tvm_struct_get(args, 0, 13, "int32")
- assert arg_type_index == 1 or arg_type_index == 2, "main: Expect arg[0] to be int"
+ assert arg_type_index == 1 or arg_type_index == 2, (
+ "TypeError",
+ [
+ "Mismatched type on argument #",
+ "0",
+ " when calling:\n `",
+ "main(arg: int32)",
+ "`,\n expected ",
+ "int",
+ ],
+ )
arg: T.int32 = T.Cast("int32", T.tvm_struct_get(args, 0, 15, "int64"))
with T.attr(0, "compute_scope", "main_compute_"):
if arg > 0:
@@ -340,18 +259,11 @@ def test_int_parameter():
return 0
After = tvm.tir.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(Expected, After)
+ tvm.ir.assert_structural_equal(After, Expected)
def test_bool_parameter():
- """An integer may be passed to a function acccepting Boolean
-
- A PackedFunc produced by compiling an IRModule should support the
- same type conversions as the C++ implementation. When a function
- accepts a boolean argument, the caller may call it with an integer
- value.
-
- """
+ """Bool parameter emits type check accepting bool or int."""
@I.ir_module
class Before:
@@ -367,22 +279,38 @@ def test_bool_parameter():
class Expected:
@T.prim_func
def main(
- self: T.handle,
+ self_handle: T.handle,
args: T.handle,
num_args: T.int32,
- result: T.handle("void"),
+ result: T.handle("void", "global"),
) -> T.int32:
T.func_attr(
{
"calling_conv": 1,
- "target": T.target("llvm"),
"global_symbol": "__tvm_ffi_main",
+ "target": T.target("llvm"),
}
)
- assert num_args == 1, "main: num_args should be 1"
- assert not T.isnullptr(args), "main: args pointer is NULL"
+ assert num_args == 1, (
+ "TypeError",
+ ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: bool)", "`"],
+ )
+ assert not T.isnullptr(args), (
+ "TypeError",
+ ["args pointer is NULL", " when calling:\n `", "main(arg: bool)", "`"],
+ )
arg_type_index: T.int32 = T.tvm_struct_get(args, 0, 13, "int32")
- assert arg_type_index == 2 or arg_type_index == 1, "main: Expect arg[0] to be boolean"
+ assert arg_type_index == 2 or arg_type_index == 1, (
+ "TypeError",
+ [
+ "Mismatched type on argument #",
+ "0",
+ " when calling:\n `",
+ "main(arg: bool)",
+ "`,\n expected ",
+ "boolean",
+ ],
+ )
arg: T.bool = T.Cast("bool", T.tvm_struct_get(args, 0, 15, "int64"))
with T.attr(0, "compute_scope", "main_compute_"):
if arg:
@@ -398,7 +326,103 @@ def test_bool_parameter():
return 0
After = tvm.tir.transform.MakePackedAPI()(Before)
- tvm.ir.assert_structural_equal(Expected, After)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_float_parameter():
+ """Float parameter emits type check accepting float, int, or bool."""
+
+ @I.ir_module
+ class Before:
+ @T.prim_func
+ def main(arg: T.float32) -> T.int32:
+ T.func_attr({"target": T.target("llvm", host="llvm")})
+ if arg > T.float32(0):
+ return 10
+ else:
+ return 20
+
+ @I.ir_module
+ class Expected:
+ @T.prim_func
+ def main(
+ self_handle: T.handle,
+ args: T.handle,
+ num_args: T.int32,
+ result: T.handle("void", "global"),
+ ) -> T.int32:
+ T.func_attr(
+ {
+ "calling_conv": 1,
+ "global_symbol": "__tvm_ffi_main",
+ "target": T.target("llvm"),
+ }
+ )
+ assert num_args == 1, (
+ "TypeError",
+ ["Expected ", "1", " arguments", " when calling:\n `", "main(arg: float32)", "`"],
+ )
+ assert not T.isnullptr(args), (
+ "TypeError",
+ ["args pointer is NULL", " when calling:\n `", "main(arg: float32)", "`"],
+ )
+ arg_type_index: T.int32 = T.tvm_struct_get(args, 0, 13, "int32")
+ assert arg_type_index == 3 or arg_type_index == 1 or arg_type_index == 2, (
+ "TypeError",
+ [
+ "Mismatched type on argument #",
+ "0",
+ " when calling:\n `",
+ "main(arg: float32)",
+ "`,\n expected ",
+ "float",
+ ],
+ )
+ arg: T.float32 = T.Select(
+ arg_type_index == 3,
+ T.Cast("float32", T.tvm_struct_get(args, 0, 15, "float64")),
+ T.Cast("float32", T.tvm_struct_get(args, 0, 15, "int64")),
+ )
+ with T.attr(0, "compute_scope", "main_compute_"):
+ if arg > T.float32(0.0):
+ T.tvm_struct_set(result, 0, 13, 1)
+ T.tvm_struct_set(result, 0, 14, 0)
+ T.tvm_struct_set(result, 0, 15, T.Cast("int64", 10))
+ return 0
+ else:
+ T.tvm_struct_set(result, 0, 13, 1)
+ T.tvm_struct_set(result, 0, 14, 0)
+ T.tvm_struct_set(result, 0, 15, T.Cast("int64", 20))
+ return 0
+ return 0
+
+ After = tvm.tir.transform.MakePackedAPI()(Before)
+ tvm.ir.assert_structural_equal(After, Expected)
+
+
+def test_forward_reference_symbolic_variable():
+ """MakePackedAPI succeeds when a symbolic variable is used before it is defined.
+
+ When buffer A has shape (batch_size+1,) and buffer B has shape (batch_size,),
+ batch_size is referenced (in A's shape check) before it is defined (from B's
+ shape). The three-sequence separation (init_nest, asserts, decl_buffers)
+ ensures all variable definitions precede all assertions.
+ """
+
+ @I.ir_module
+ class Before:
+ @T.prim_func
+ def main(a: T.handle, b: T.handle):
+ T.func_attr({"target": T.target("llvm", host="llvm")})
+ batch_size = T.int64()
+ A = T.match_buffer(a, (batch_size + 1,), "int32")
+ B = T.match_buffer(b, (batch_size,), "int32")
+ for i in range(batch_size):
+ B[i] = A[i] + A[i + 1]
+
+ # Should not raise "variable batch_size has been used before definition"
+ After = tvm.tir.transform.MakePackedAPI()(Before)
+ assert len(After["main"].params) == 4
if __name__ == "__main__":
diff --git a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
index 71e8f37803..3a5d41415a 100644
--- a/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_ir_builder_tir.py
@@ -301,7 +301,11 @@ def test_ir_builder_tir_assert():
# AssertStmt is a leaf. The frame emits the assert and then the body stmts as siblings.
assert_expected = tir.SeqStmt(
[
- tir.AssertStmt(T.int32() == 0, tir.StringImm("a is 0")),
+ tir.AssertStmt(
+ T.int32() == 0,
+ tir.StringImm("RuntimeError"),
+ [tir.StringImm("a is 0")],
+ ),
tir.Evaluate(0),
]
)
diff --git a/tests/python/tvmscript/test_tvmscript_printer_tir.py b/tests/python/tvmscript/test_tvmscript_printer_tir.py
index 93ddff0fd4..b2deb93155 100644
--- a/tests/python/tvmscript/test_tvmscript_printer_tir.py
+++ b/tests/python/tvmscript/test_tvmscript_printer_tir.py
@@ -289,7 +289,7 @@ def test_assert_stmt():
_assert_print(
obj,
"""
-assert T.bool(True), "assertion"
+assert T.bool(True), ("RuntimeError", ["assertion"])
T.evaluate(0)
""",
)
diff --git a/tests/python/tvmscript/test_tvmscript_roundtrip.py b/tests/python/tvmscript/test_tvmscript_roundtrip.py
index 39ad3e978f..a0b77a4a98 100644
--- a/tests/python/tvmscript/test_tvmscript_roundtrip.py
+++ b/tests/python/tvmscript/test_tvmscript_roundtrip.py
@@ -3728,5 +3728,53 @@ def test_address_of_buffer():
assert "T.address_of(A[0, 0])" in func.script()
+def test_assert_stmt_roundtrip_runtime_error():
+ """RuntimeError assert roundtrips through print->parse."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("RuntimeError", ["x must be positive"])
+
+ script = func.script(show_meta=True)
+ roundtrip = tvm.script.from_source(script, check_well_formed=False)
+ tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
+
+
+def test_assert_stmt_roundtrip_value_error():
+ """ValueError assert roundtrips through print->parse."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("ValueError", ["Shape mismatch"])
+
+ script = func.script(show_meta=True)
+ roundtrip = tvm.script.from_source(script, check_well_formed=False)
+ tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
+
+
+def test_assert_stmt_roundtrip_type_error():
+ """TypeError assert roundtrips through print->parse."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("TypeError", ["Expected Tensor but got int"])
+
+ script = func.script(show_meta=True)
+ roundtrip = tvm.script.from_source(script, check_well_formed=False)
+ tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
+
+
+def test_assert_stmt_roundtrip_multi_parts():
+ """Multi-part message assert roundtrips with structural equality."""
+
+ @T.prim_func
+ def func(x: T.int32):
+ assert x > 0, ("TypeError", ["Expected ", "Tensor", " but got ", "int"])
+
+ script = func.script(show_meta=True)
+ roundtrip = tvm.script.from_source(script, check_well_formed=False)
+ tvm.ir.assert_structural_equal(func, roundtrip, map_free_vars=True)
+
+
if __name__ == "__main__":
tvm.testing.main()