This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx-standalone-root in repository https://gitbox.apache.org/repos/asf/tvm.git
commit ca2d38491dadd607c9f737a3b9355e4a529f17c2 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(
