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 9094f0e9cf [TIRx] Infer variables from linear packed shapes (#20782)
9094f0e9cf is described below
commit 9094f0e9cfd3bf460cdb80ab782c52bc3d0725f0
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Oct 17 17:50:12 2026 -0400
[TIRx] Infer variables from linear packed shapes (#20782)
Prefer bare-variable dimensions across packed tensor arguments, then
infer unresolved shape variables from linear dimensions such as `4 * n +
3`. Require exact divisibility before binding and check remaining
constraints after the variables are defined, preserving the inferred
relation in argument diagnostics.
---
tirx/src/tirx/transform/tvm_ffi_binder.cc | 92 +++++++-
tirx/src/tirx/transform/tvm_ffi_binder.h | 26 ++-
tirx/tests/python/codegen/llvm/test_basic.py | 21 ++
.../python/tirx/transform/test_make_packed_api.py | 259 +++++++++++++++++++--
4 files changed, 369 insertions(+), 29 deletions(-)
diff --git a/tirx/src/tirx/transform/tvm_ffi_binder.cc
b/tirx/src/tirx/transform/tvm_ffi_binder.cc
index 3ead929e6f..bad702479a 100644
--- a/tirx/src/tirx/transform/tvm_ffi_binder.cc
+++ b/tirx/src/tirx/transform/tvm_ffi_binder.cc
@@ -231,6 +231,70 @@ bool TVMFFIABIBuilder::BindScalar(const PrimExpr& arg,
const PrimExpr& value,
return false;
}
+bool TVMFFIABIBuilder::TryBindLinearIntExpr(const PrimExpr& arg, const
PrimExpr& value,
+ const ffi::reflection::AccessPath&
path) {
+ // Recognize factor * var + offset, allowing either operand order for Add
and Mul.
+ PrimExpr var = arg;
+ PrimExpr factor = IntImm(arg.ty(), 1);
+ PrimExpr offset = IntImm(arg.ty(), 0);
+
+ // Remove the constant offset, if present.
+ if (const auto* add = var.as<AddNode>()) {
+ if (add->b.as<IntImmNode>()) {
+ var = add->a;
+ offset = add->b;
+ } else if (add->a.as<IntImmNode>()) {
+ var = add->b;
+ offset = add->a;
+ } else {
+ return false;
+ }
+ } else if (const auto* sub = var.as<SubNode>()) {
+ if (!sub->b.as<IntImmNode>()) return false;
+ var = sub->a;
+ offset = -sub->b;
+ }
+
+ // Remove the constant factor; the remainder must be a variable.
+ if (const auto* mul = var.as<MulNode>()) {
+ if (mul->b.as<IntImmNode>()) {
+ var = mul->a;
+ factor = mul->b;
+ } else if (mul->a.as<IntImmNode>()) {
+ var = mul->b;
+ factor = mul->a;
+ } else {
+ return false;
+ }
+ }
+ if (!var.as<VarNode>() || arg.same_as(var) || IsZero(factor) ||
+ var_defs_.count(var.as<VarNode>())) {
+ return false;
+ }
+
+ // Check divisibility before binding the inferred variable to the quotient.
+ PrimExpr delta = value - offset;
+ std::string shape_path = RenderAccessPath(path);
+ std::string decoded_expr = shape_path;
+ if (!IsZero(offset)) {
+ std::ostringstream os;
+ os << "(" << shape_path << " - " << offset.as_or_throw<IntImm>()->value <<
")";
+ decoded_expr = os.str();
+ }
+ if (!IsOne(factor)) {
+ std::ostringstream os;
+ os << factor.as_or_throw<IntImm>()->value;
+ std::string factor_text = os.str();
+ EmitAssert(truncmod(delta, factor) == 0, "ValueError", "Invalid ",
shape_path, " on argument #",
+ std::to_string(GetParamIndex(path)), when_calling_imm_,
sig_imm_, "`,\n expected ",
+ decoded_expr, " to be divisible by ", factor_text);
+ decoded_expr += " // " + factor_text;
+ }
+ BindScalar(var, truncdiv(delta, factor), path, true);
+ shape_var_exprs_.emplace(var.as<VarNode>(), std::move(decoded_expr));
+ return true;
+}
+
bool TVMFFIABIBuilder::BindPointer(const Var& arg, const Expr& value,
const ffi::reflection::AccessPath& path,
bool with_lets) {
TVM_FFI_ICHECK(ffi::StructuralEqual()(arg->ty, value->ty));
@@ -335,6 +399,9 @@ void TVMFFIABIBuilder::RenderPendingAsserts() {
if (pending_const_asserts_.empty()) return;
ExprPathRenderer renderer([this](const VarNode* v) -> std::string {
+ if (auto inferred = shape_var_exprs_.find(v); inferred !=
shape_var_exprs_.end()) {
+ return "(" + inferred->second + ")";
+ }
auto it = var_defs_.find(v);
if (it != var_defs_.end()) {
ffi::String path = RenderAccessPath(it->second.first_def_path);
@@ -603,12 +670,13 @@ void TVMFFIABIBuilder::DecodeParam(int param_index) {
void TVMFFIABIBuilder::DecodeAllParams() {
int num_args = static_cast<int>(params_.size());
- // Phase 1: Decode each packed argument (type-check, value load, scalar bind)
+ // Decode scalar parameters before any tensor shape can reference them.
for (int i = 0; i < num_args; ++i) {
DecodeParam(i);
}
- // Phase 2: Bind DLTensor tensors (shape, strides, dtype, device checks)
+ // Shape pass 1: Prefer bare-variable dimensions across all arguments.
+ // Per-tensor safety checks still precede the corresponding field loads.
for (int i = 0; i < num_args; ++i) {
Var param = params_[i];
if (auto tensor = param.as<TensorVar>()) {
@@ -640,6 +708,17 @@ void TVMFFIABIBuilder::DecodeAllParams() {
}
}
}
+ // Shape pass 2: Infer only variables not defined by a bare dimension.
+ for (auto& shape : pending_int_exprs_) {
+ shape.needs_check = !TryBindLinearIntExpr(shape.arg, shape.value,
shape.path);
+ }
+ // Shape pass 3: Check the remaining expressions after all bindings are
known.
+ for (const auto& shape : pending_int_exprs_) {
+ if (shape.needs_check) {
+ BindScalar(shape.arg, shape.value, shape.path, true);
+ }
+ }
+ pending_int_exprs_.clear();
}
// ============================================================
@@ -772,8 +851,13 @@ Expr TVMFFIABIBuilder::DecodeParamDLTensor(const
TensorVar& tensor, const PrimEx
break;
}
ffi::reflection::AccessPath shape_k_path =
param_path->Attr(ffi::String("shape"))->ArrayItem(k);
- BindScalar(tensor.shape()[k], cast(LoadInt64ArrayElem(shape_ptr, k),
tensor.shape()[k].ty()),
- shape_k_path, true);
+ PrimExpr shape = tensor.shape()[k];
+ PrimExpr value = cast(LoadInt64ArrayElem(shape_ptr, k), shape.ty());
+ if (shape.as<VarNode>()) {
+ BindScalar(shape, value, shape_k_path, true);
+ } else {
+ pending_int_exprs_.push_back({shape, value, shape_k_path});
+ }
}
// ── Section: strides ─────────────────────────────────────────
diff --git a/tirx/src/tirx/transform/tvm_ffi_binder.h
b/tirx/src/tirx/transform/tvm_ffi_binder.h
index 46f323f8b1..832a2c3450 100644
--- a/tirx/src/tirx/transform/tvm_ffi_binder.h
+++ b/tirx/src/tirx/transform/tvm_ffi_binder.h
@@ -127,7 +127,8 @@ class TVMFFIABIBuilder {
* 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 tensors to handles (shape, strides, dtype, device
checks)
+ * 2. Binds DLTensor tensors to handles, preferring bare-variable dimensions
+ * 3. Infers unresolved linear shapes, then adds remaining shape checks
*
* Statements are routed to init_nest_, asserts_, or decl_tensors_.
*/
@@ -226,6 +227,13 @@ class TVMFFIABIBuilder {
// ── Private binding submethods (all take ffi::reflection::AccessPath)
───────────
+ /*!
+ * \brief Infer a variable from a linear shape when no bare dimension
defines it.
+ * \return Whether this shape defined a variable and needs no further
equality check.
+ */
+ bool TryBindLinearIntExpr(const PrimExpr& arg, const PrimExpr& value,
+ const ffi::reflection::AccessPath& path);
+
/*!
* \brief Internal scalar bind with ffi::reflection::AccessPath tracking and
rich error messages.
*
@@ -383,8 +391,24 @@ class TVMFFIABIBuilder {
PrimExpr expected_expr;
};
+ /*! \brief An integer expression deferred until bare dimensions are bound. */
+ struct PendingIntExpr {
+ /*! \brief The expected dimension expression from the tensor signature. */
+ PrimExpr arg;
+ /*! \brief The dimension value loaded from the runtime argument. */
+ PrimExpr value;
+ /*! \brief The argument access path used in diagnostics. */
+ ffi::reflection::AccessPath path;
+ /*! \brief Whether equality still needs checking after linear inference. */
+ bool needs_check{true};
+ };
+
+ /*! \brief Non-variable shapes, processed after all bare dimensions have
been bound. */
+ std::vector<PendingIntExpr> pending_int_exprs_;
/*! \brief The definition map: VarNode* -> VarDefInfo (value +
first_def_path). */
std::unordered_map<const VarNode*, VarDefInfo> var_defs_;
+ /*! \brief Decoded shape expressions for diagnostics of inferred variables.
*/
+ std::unordered_map<const VarNode*, std::string> shape_var_exprs_;
/*! \brief Variable definitions: Binds, shape/strides DeclTensors. */
std::vector<Stmt> init_nest_;
/*! \brief Validation checks: all AssertStmts. */
diff --git a/tirx/tests/python/codegen/llvm/test_basic.py
b/tirx/tests/python/codegen/llvm/test_basic.py
index 8c71c1e004..7802879ff1 100644
--- a/tirx/tests/python/codegen/llvm/test_basic.py
+++ b/tirx/tests/python/codegen/llvm/test_basic.py
@@ -236,6 +236,27 @@ def test_llvm_vadd_pipeline():
tvm_tirx.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy())
[email protected](not env.has_llvm(), reason="need llvm")
+def test_llvm_add_linear_shape():
+ n = T.dynamic("n", "int64")
+
+ @T.function
+ def add(
+ A: T.Tensor((4 * n + 3,), "float32"),
+ B: T.Tensor((4 * n + 3,), "float32"),
+ C: T.Tensor((4 * n + 3,), "float32"),
+ ):
+ for i in range(4 * n + 3):
+ C[i] = A[i] + B[i]
+
+ lib = tvm_tirx.tirx.compile(add, target="llvm")
+ a = tvm_tirx.runtime.tensor(np.arange(19, dtype="float32"))
+ b = tvm_tirx.runtime.tensor(np.ones(19, dtype="float32"))
+ c = tvm_tirx.runtime.tensor(np.zeros(19, dtype="float32"))
+ lib(a, b, c)
+ np.testing.assert_equal(c.numpy(), a.numpy() + b.numpy())
+
+
@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_llvm_madd_pipeline():
def check_llvm(nn, base, stride):
diff --git a/tirx/tests/python/tirx/transform/test_make_packed_api.py
b/tirx/tests/python/tirx/transform/test_make_packed_api.py
index 367a5fae34..98d96a5fdb 100644
--- a/tirx/tests/python/tirx/transform/test_make_packed_api.py
+++ b/tirx/tests/python/tirx/transform/test_make_packed_api.py
@@ -247,30 +247,30 @@ def test_target_host_removed():
" bytes",
],
)
- assert T.Cast("int32", T.abi_field_get(main_A_shape, 0, 17,
ty="int64")) == 1, (
+ assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
"ValueError",
[
"Invalid ",
- "A.shape[0]",
+ "A.byte_offset",
" on argument #",
"0",
" when calling:\n `",
"main(A: Tensor([1], float32))",
"`,\n expected ",
- "1",
+ "0",
],
)
- assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
+ assert T.Cast("int32", T.abi_field_get(main_A_shape, 0, 17,
ty="int64")) == 1, (
"ValueError",
[
"Invalid ",
- "A.byte_offset",
+ "A.shape[0]",
" on argument #",
"0",
" when calling:\n `",
"main(A: Tensor([1], float32))",
"`,\n expected ",
- "0",
+ "1",
],
)
A = T.decl_tensor(
@@ -901,6 +901,217 @@ def test_float_parameter():
tvm_tirx.ir.assert_structural_equal(After, Expected)
+def test_linear_shape_binding():
+ n = T.dynamic("n", "int64")
+ signature = "main(A: Tensor([T.int64(4) * n + T.int64(3)], float32))"
+
+ @I.ir_module
+ class Before:
+ @T.function
+ def main(A: T.Tensor((4 * n + 3,), "float32")):
+ T.func_attr({"target": T.target("llvm", host="llvm")})
+ T.evaluate(n)
+
+ @I.ir_module
+ class Expected:
+ @T.function
+ def main(
+ self_handle: T.Ptr(T.void),
+ args: T.Ptr(T.void),
+ num_args: T.int32,
+ result: T.Ptr(T.void),
+ ) -> T.int32:
+ T.func_attr({"global_symbol": "__tvm_ffi_main"})
+ T.func_attr({"calling_conv": 1, "target": T.target("llvm")})
+ assert num_args == 1, (
+ "TypeError",
+ ["Expected ", "1", " arguments", " when calling:\n `",
signature, "`"],
+ )
+ assert not T.isnullptr(args), (
+ "TypeError",
+ ["args pointer is NULL", " when calling:\n `", signature,
"`"],
+ )
+ A_type_index: T.let[T.int32] = T.abi_field_get(args, 0, 13,
ty=T.int32)
+ assert (
+ A_type_index == 0 or A_type_index == 4 or A_type_index == 7 or
A_type_index >= 64
+ ), (
+ "TypeError",
+ [
+ "Mismatched type on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "Tensor",
+ ],
+ )
+ A_handle: T.let[T.Ptr(T.void)] = T.if_then_else(
+ A_type_index == 70,
+ T.handle_add_byte_offset(T.abi_field_get(args, 0, 15,
ty=T.Ptr(T.void)), 24),
+ T.abi_field_get(args, 0, 15, ty=T.Ptr(T.void)),
+ )
+ assert not T.isnullptr(A_handle), (
+ "TypeError",
+ [
+ "Mismatched type on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "Tensor",
+ ],
+ )
+ assert 1 == T.abi_field_get(A_handle, 0, 4, ty=T.int32), (
+ "ValueError",
+ [
+ "Mismatched ",
+ "A",
+ ".ndim on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "1",
+ ],
+ )
+ assert (
+ T.abi_field_get(A_handle, 0, 5, ty=T.uint8) == T.uint8(2)
+ and T.abi_field_get(A_handle, 0, 6, ty=T.uint8) == T.uint8(32)
+ and T.abi_field_get(A_handle, 0, 7, ty=T.uint16) == T.uint16(1)
+ ), (
+ "TypeError",
+ [
+ "Mismatched ",
+ "A",
+ ".dtype on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "float32",
+ ],
+ )
+ main_A_shape: T.let[T.Ptr(T.int64)] = T.abi_field_get(A_handle, 0,
2, ty=T.Ptr(T.int64))
+ main_A_strides: T.let[T.Ptr(T.int64)] = T.abi_field_get(
+ A_handle, 0, 3, ty=T.Ptr(T.int64)
+ )
+ assert T.abi_field_get(A_handle, 0, 10, ty=T.int32) == 1, (
+ "ValueError",
+ [
+ "Mismatched ",
+ "A",
+ ".device_type on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "cpu",
+ ],
+ )
+ dev_id: T.let[T.int32] = T.abi_field_get(A_handle, 0, 9,
ty=T.int32)
+ assert T.truncmod(
+ T.abi_field_get(main_A_shape, 0, 17, ty=T.int64) - T.int64(3),
T.int64(4)
+ ) == T.int64(0), (
+ "ValueError",
+ [
+ "Invalid ",
+ "A.shape[0]",
+ " on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "(A.shape[0] - 3)",
+ " to be divisible by ",
+ "4",
+ ],
+ )
+ n: T.let[T.int64] = T.Div(
+ T.abi_field_get(main_A_shape, 0, 17, ty=T.int64) - T.int64(3),
T.int64(4)
+ )
+ with T.device_context(1, dev_id):
+ if not T.isnullptr(main_A_strides):
+ assert T.int64(4) * n + T.int64(3) == T.int64(1) or
T.int64(
+ 1
+ ) == T.abi_field_get(main_A_strides, 0, 17, ty=T.int64), (
+ "ValueError",
+ [
+ "Mismatched ",
+ "A",
+ ".strides on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected to be compact array",
+ ],
+ )
+ assert T.int64(4) * n + T.int64(3) == T.int64(0) or not
T.isnullptr(
+ T.reinterpret(
+ T.abi_field_get(A_handle, 0, 1, ty=T.Ptr(T.void)),
ty=T.Ptr(T.float32)
+ )
+ ), (
+ "ValueError",
+ [
+ "A",
+ " data pointer is NULL on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected non-NULL data pointer",
+ ],
+ )
+ assert T.int64(4) * n + T.int64(3) == T.int64(0) or T.truncmod(
+ T.reinterpret(
+ T.reinterpret(
+ T.reinterpret(
+ T.abi_field_get(A_handle, 0, 1,
ty=T.Ptr(T.void)),
+ ty=T.Ptr(T.float32),
+ ),
+ ty=T.Ptr(T.void),
+ ),
+ ty=T.uint64,
+ ),
+ T.uint64(64),
+ ) == T.uint64(0), (
+ "ValueError",
+ [
+ "Misaligned Tensor data on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected data alignment=",
+ "64",
+ " bytes",
+ ],
+ )
+ assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty=T.uint64), (
+ "ValueError",
+ [
+ "Invalid ",
+ "A.byte_offset",
+ " on argument #",
+ "0",
+ " when calling:\n `",
+ signature,
+ "`,\n expected ",
+ "0",
+ ],
+ )
+ A = T.decl_tensor(
+ T.reinterpret(
+ T.abi_field_get(A_handle, 0, 1, ty=T.Ptr(T.void)),
ty=T.Ptr(T.float32)
+ ),
+ ty_args=[T.Tensor((T.int64(4) * n + T.int64(3),),
T.float32)],
+ )
+ T.assume_aligned(A, 64)
+ with T.compute_scope("main_compute_"):
+ T.evaluate(n)
+ return 0
+
+ After = tvm_tirx.tirx.transform.MakePackedAPI()(Before)
+ tvm_tirx.ir.assert_structural_equal(After, Expected)
+
+
def test_forward_reference_symbolic_variable():
"""MakePackedAPI succeeds when a symbolic variable is used before it is
defined.
@@ -1239,45 +1450,45 @@ def test_forward_reference_symbolic_variable():
" bytes",
],
)
- assert batch_size + T.int64(1) == T.abi_field_get(
- main_A_shape, 0, 17, ty="int64"
- ), (
+ assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
"ValueError",
[
"Invalid ",
- "A.shape[0]",
+ "A.byte_offset",
" on argument #",
"0",
" when calling:\n `",
signature,
"`,\n expected ",
- "B.shape[0] + 1",
+ "0",
],
)
- assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
+ assert T.uint64(0) == T.abi_field_get(B_handle, 0, 8,
ty="uint64"), (
"ValueError",
[
"Invalid ",
- "A.byte_offset",
+ "B.byte_offset",
" on argument #",
- "0",
+ "1",
" when calling:\n `",
signature,
"`,\n expected ",
"0",
],
)
- assert T.uint64(0) == T.abi_field_get(B_handle, 0, 8,
ty="uint64"), (
+ assert batch_size + T.int64(1) == T.abi_field_get(
+ main_A_shape, 0, 17, ty="int64"
+ ), (
"ValueError",
[
"Invalid ",
- "B.byte_offset",
+ "A.shape[0]",
" on argument #",
- "1",
+ "0",
" when calling:\n `",
signature,
"`,\n expected ",
- "0",
+ "B.shape[0] + 1",
],
)
A = T.decl_tensor(
@@ -1485,30 +1696,30 @@ def test_tensor_alignment_attached_to_tensor_var():
" bytes",
],
)
- assert T.Cast("int32", T.abi_field_get(main_A_shape, 0, 17,
ty="int64")) == 16, (
+ assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
"ValueError",
[
"Invalid ",
- "A.shape[0]",
+ "A.byte_offset",
" on argument #",
"0",
" when calling:\n `",
"main(A: Tensor([16], float32))",
"`,\n expected ",
- "16",
+ "0",
],
)
- assert T.uint64(0) == T.abi_field_get(A_handle, 0, 8,
ty="uint64"), (
+ assert T.Cast("int32", T.abi_field_get(main_A_shape, 0, 17,
ty="int64")) == 16, (
"ValueError",
[
"Invalid ",
- "A.byte_offset",
+ "A.shape[0]",
" on argument #",
"0",
" when calling:\n `",
"main(A: Tensor([16], float32))",
"`,\n expected ",
- "0",
+ "16",
],
)
A = T.decl_tensor(
