From ad89a279e15c535ae465b8bfbf2b555b1e82fe7e Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Tue, 11 Aug 2026 07:23:38 -0400 Subject: [PATCH 1/3] [FIX][RELAX] Preserve tensor-derived symbols during fusion --- src/relax/transform/fuse_ops.cc | 32 ++- tests/python/relax/test_transform_fuse_ops.py | 194 ++++++++++++++++++ 2 files changed, 224 insertions(+), 2 deletions(-) diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index d2f9870a3857..302b4d194fd2 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -679,10 +679,38 @@ class FunctionCreator : public ExprMutator { if (const auto* tuple = expr.as()) { return std::all_of(tuple->fields.begin(), tuple->fields.end(), [this](const Expr& e) { return IsInlinableConstants(e); }); + } else if (auto prim_value = expr.as()) { + ffi::Array undefined_vars = tvm::tirx::UndefinedVars(prim_value.value()); + if (undefined_vars.empty()) { + return true; + } + + // A symbolic value directly defined by an existing tensor/shape parameter is already part + // of the grouped function's shape environment. A variable that only occurs inside a + // derived dimension is not directly definable, but leaving it symbolic lets CreateFunction + // add the explicit Shape parameter required to define it. In either case, lifting the value + // as an unrelated scalar parameter would sever the relation to the parameter types. + auto parameter_types = + TupleType(params_.Map([](const Var& param) { return GetType(param); })); + std::unordered_set definable_shape_vars; + for (const tirx::Var& var : DefinableTIRVarsInType(parameter_types)) { + definable_shape_vars.insert(var); + } + if (std::all_of(undefined_vars.begin(), undefined_vars.end(), + [&definable_shape_vars](const tirx::Var& var) { + return definable_shape_vars.count(var); + })) { + return true; + } + + std::unordered_set referenced_shape_vars; + for (const tirx::Var& var : TIRVarsInType(parameter_types)) { + referenced_shape_vars.insert(var); + } + return std::all_of(undefined_vars.begin(), undefined_vars.end(), + [&](const tirx::Var& var) { return referenced_shape_vars.count(var); }); } else if (expr.as() || expr.as()) { return false; - } else if (auto prim_value = expr.as()) { - return tvm::tirx::UndefinedVars(prim_value.value()).empty(); } else if (const auto* shape_expr = expr.as()) { return std::all_of(shape_expr->values.begin(), shape_expr->values.end(), [](const PrimExpr& e) { return tvm::tirx::UndefinedVars(e).empty(); }); diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py index ed3b02d9117d..715c2d3a65e8 100644 --- a/tests/python/relax/test_transform_fuse_ops.py +++ b/tests/python/relax/test_transform_fuse_ops.py @@ -1352,6 +1352,200 @@ def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"): _check(Before, Expected) +def test_symbolic_prim_arg_bound_by_tensor_shape(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n), "float32") + for i in range(n): + with T.sblock("add_one"): + vi = T.axis.spatial(n, i) + out[0, vi] = x[0, vi] + T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n), "float32") + for i in range(n): + with T.sblock("exp"): + vi = T.axis.spatial(n, i) + out[0, vi] = T.exp(x[0, vi]) + + @R.function + def main( + x: R.Tensor((1, "n"), dtype="float32"), + ) -> R.Tensor((1, "n"), dtype="float32"): + n = T.int64() + cls = Before + with R.dataflow(): + lv = R.call_tir( + cls.add_one, + (x, n), + out_ty=R.Tensor((1, n), dtype="float32"), + ) + gv = R.call_tir( + cls.exp, + (lv, n), + out_ty=R.Tensor((1, n), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 1 + assert fused.ret_ty.shape is not None + + +def test_symbolic_prim_arg_reused_from_derived_tensor_shape(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer( + x_handle, + (T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)), + "float32", + ) + out = T.match_buffer( + out_handle, + (T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)), + "float32", + ) + for i in range((n - T.int64(1)) // T.int64(4) + T.int64(1)): + with T.sblock("add_one"): + vi = T.axis.spatial((n - T.int64(1)) // T.int64(4) + T.int64(1), i) + out[0, vi] = x[0, vi] + T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer( + x_handle, + (T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)), + "float32", + ) + out = T.match_buffer( + out_handle, + (T.int64(1), (n - T.int64(1)) // T.int64(4) + T.int64(1)), + "float32", + ) + for i in range((n - T.int64(1)) // T.int64(4) + T.int64(1)): + with T.sblock("exp"): + vi = T.axis.spatial((n - T.int64(1)) // T.int64(4) + T.int64(1), i) + out[0, vi] = T.exp(x[0, vi]) + + @R.function + def main( + source: R.Tensor(("n",), dtype="float32"), + x: R.Tensor((1, "(n - 1) // 4 + 1"), dtype="float32"), + ) -> R.Tensor((1, "(n - 1) // 4 + 1"), dtype="float32"): + n = T.int64() + cls = Before + with R.dataflow(): + lv = R.call_tir( + cls.add_one, + (x, n), + out_ty=R.Tensor((1, (n - 1) // 4 + 1), dtype="float32"), + ) + gv = R.call_tir( + cls.exp, + (lv, n), + out_ty=R.Tensor((1, (n - 1) // 4 + 1), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 2 + assert isinstance(fused.params[1].ty, relax.ShapeType) + assert all(not isinstance(param.ty, tvm.ir.PrimType) for param in fused.params) + + mod = relax.transform.FuseTIR()(mod) + fused_tir = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert tvm.tirx.analysis.verify_well_formed(fused_tir) + + +def test_symbolic_prim_arg_not_bound_by_derived_tensor_shape(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def add_one(x_handle: T.handle, n: T.int64, m: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n + T.int64(1)), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n + T.int64(1)), "float32") + for i in range(n + T.int64(1)): + with T.sblock("add_one"): + vi = T.axis.spatial(n + T.int64(1), i) + out[0, vi] = x[0, vi] + T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def exp(x_handle: T.handle, n: T.int64, m: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n + T.int64(1)), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n + T.int64(1)), "float32") + for i in range(n + T.int64(1)): + with T.sblock("exp"): + vi = T.axis.spatial(n + T.int64(1), i) + out[0, vi] = T.exp(x[0, vi]) + + @R.function + def main( + shape: R.Shape(["n", "m"]), + x: R.Tensor((1, "n + 1"), dtype="float32"), + ) -> R.Tensor((1, "n + 1"), dtype="float32"): + n = T.int64() + m = T.int64() + cls = Before + with R.dataflow(): + lv = R.call_tir( + cls.add_one, + (x, n, m), + out_ty=R.Tensor((1, n + 1), dtype="float32"), + ) + gv = R.call_tir( + cls.exp, + (lv, n, m), + out_ty=R.Tensor((1, n + 1), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 3 + assert sum(isinstance(param.ty, relax.ShapeType) for param in fused.params) == 1 + assert sum(isinstance(param.ty, tvm.ir.PrimType) for param in fused.params) == 1 + + def test_shape_expr_arg(): @I.ir_module(s_tir=True) class Before: From d00e421bfa373a704510815a31576bf7ec7d9b91 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Wed, 26 Aug 2026 02:54:52 -0400 Subject: [PATCH 2/3] [FIX][RELAX] Defer symbolic fusion argument classification --- src/relax/transform/fuse_ops.cc | 117 +++++---- tests/python/relax/test_transform_fuse_ops.py | 222 +++++++++++++++++- 2 files changed, 294 insertions(+), 45 deletions(-) diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index 302b4d194fd2..5ebe36373b71 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -522,7 +522,11 @@ class FunctionCreator : public ExprMutator { params_.insert(params_.begin() + param_idx, item_params.begin(), item_params.end()); } - // Step 3. Visit each binding and collect outputs one by one. + // Step 3. Now that the complete function boundary is known, inline primitive arguments whose + // symbols participate in the boundary's tensor/shape types. + InlineShapeDependentPrimArgs(); + + // Step 4. Visit each binding and collect outputs one by one. ffi::Array outputs(output_vars_.size(), Expr()); for (const Binding& binding : bindings_) { // Special handing for TupleGetItem. @@ -554,7 +558,7 @@ class FunctionCreator : public ExprMutator { } } - // Step 4. Finish constructing the new block. + // Step 5. Finish constructing the new block. BindingBlock new_block = builder_->EndBlock(); if (outputs.empty()) { // If the result is not used outside @@ -606,6 +610,53 @@ class FunctionCreator : public ExprMutator { return std::nullopt; } + Expr ResolveOuterBinding(const Expr& expr) { + Expr bound_value = expr; + std::unordered_set visited; + while (const auto* current_var = bound_value.as()) { + if (!visited.insert(current_var).second) break; + auto it = outer_bindings_.find(ffi::GetRef(current_var)); + if (it == outer_bindings_.end()) break; + bound_value = (*it).second; + } + return bound_value; + } + + void InlineShapeDependentPrimArgs() { + ffi::Array boundary_types = params_.Map([](const Var& param) { return GetType(param); }); + for (const VarNode* output_var : output_vars_) { + boundary_types.push_back(GetType(ffi::GetRef(output_var))); + } + + std::unordered_set boundary_shape_vars; + for (const tirx::Var& var : TIRVarsInType(TupleType(boundary_types))) { + boundary_shape_vars.insert(var); + } + + for (const Expr& argument : deferred_prim_args_) { + auto it = std::find_if(arguments_.begin(), arguments_.end(), + [&](const Expr& candidate) { return candidate.same_as(argument); }); + TVM_FFI_ICHECK(it != arguments_.end()); + + Expr bound_value = ResolveOuterBinding(argument); + bool inline_argument = + bound_value.same_as(argument) && IsShapeDependentPrimExpr(argument, boundary_shape_vars); + bool inline_bound_value = !bound_value.same_as(argument) && + IsShapeDependentPrimExpr(bound_value, boundary_shape_vars); + if (!inline_argument && !inline_bound_value) continue; + + if (inline_bound_value) { + const auto* argument_var = argument.as(); + TVM_FFI_ICHECK(argument_var); + inlined_bindings_[argument_var] = bound_value; + } + + size_t param_idx = it - arguments_.begin(); + arguments_.erase(arguments_.begin() + param_idx); + params_.erase(params_.begin() + param_idx); + } + } + /*! * \brief Check whether the input expression is defined within this function. If not, create a new * parameter for the expression. @@ -624,14 +675,7 @@ class FunctionCreator : public ExprMutator { const auto* var = expr.as(); if (var != nullptr && defined_vars_.count(var) == 0) { Var bound_var = ffi::GetRef(var); - Expr bound_value = bound_var; - std::unordered_set visited; - while (const auto* current_var = bound_value.as()) { - if (!visited.insert(current_var).second) break; - auto it = outer_bindings_.find(ffi::GetRef(current_var)); - if (it == outer_bindings_.end()) break; - bound_value = (*it).second; - } + Expr bound_value = ResolveOuterBinding(bound_var); if (!bound_value.same_as(bound_var) && IsInlinableConstants(bound_value)) { inlined_bindings_[var] = bound_value; return; @@ -646,6 +690,9 @@ class FunctionCreator : public ExprMutator { Var param(std::move(name), GetType(expr)); arguments_.push_back(expr); params_.push_back(param); + if (IsSymbolicPrimExpr(expr)) { + deferred_prim_args_.push_back(expr); + } } // Mark the tuple parameter is partially referenced in the beginning. @@ -675,42 +722,30 @@ class FunctionCreator : public ExprMutator { // Check if the expression is constant PrimExpr or ShapeExpr or tuple of them that can be // inlined in the composite functions and excluded from args/params. + bool IsSymbolicPrimExpr(const Expr& expr) { + if (expr.as()) return false; + if (auto prim_value = expr.as()) { + return !tvm::tirx::UndefinedVars(prim_value.value()).empty(); + } + return false; + } + + bool IsShapeDependentPrimExpr(const Expr& expr, + const std::unordered_set& referenced_shape_vars) { + if (!IsSymbolicPrimExpr(expr)) return false; + ffi::Array undefined_vars = tvm::tirx::UndefinedVars(expr.as_or_throw()); + return std::all_of(undefined_vars.begin(), undefined_vars.end(), + [&](const tirx::Var& var) { return referenced_shape_vars.count(var); }); + } + bool IsInlinableConstants(const Expr& expr) { if (const auto* tuple = expr.as()) { return std::all_of(tuple->fields.begin(), tuple->fields.end(), [this](const Expr& e) { return IsInlinableConstants(e); }); - } else if (auto prim_value = expr.as()) { - ffi::Array undefined_vars = tvm::tirx::UndefinedVars(prim_value.value()); - if (undefined_vars.empty()) { - return true; - } - - // A symbolic value directly defined by an existing tensor/shape parameter is already part - // of the grouped function's shape environment. A variable that only occurs inside a - // derived dimension is not directly definable, but leaving it symbolic lets CreateFunction - // add the explicit Shape parameter required to define it. In either case, lifting the value - // as an unrelated scalar parameter would sever the relation to the parameter types. - auto parameter_types = - TupleType(params_.Map([](const Var& param) { return GetType(param); })); - std::unordered_set definable_shape_vars; - for (const tirx::Var& var : DefinableTIRVarsInType(parameter_types)) { - definable_shape_vars.insert(var); - } - if (std::all_of(undefined_vars.begin(), undefined_vars.end(), - [&definable_shape_vars](const tirx::Var& var) { - return definable_shape_vars.count(var); - })) { - return true; - } - - std::unordered_set referenced_shape_vars; - for (const tirx::Var& var : TIRVarsInType(parameter_types)) { - referenced_shape_vars.insert(var); - } - return std::all_of(undefined_vars.begin(), undefined_vars.end(), - [&](const tirx::Var& var) { return referenced_shape_vars.count(var); }); } else if (expr.as() || expr.as()) { return false; + } else if (auto prim_value = expr.as()) { + return tvm::tirx::UndefinedVars(prim_value.value()).empty(); } else if (const auto* shape_expr = expr.as()) { return std::all_of(shape_expr->values.begin(), shape_expr->values.end(), [](const PrimExpr& e) { return tvm::tirx::UndefinedVars(e).empty(); }); @@ -723,6 +758,8 @@ class FunctionCreator : public ExprMutator { std::unordered_set defined_vars_; /*! \brief Caller variables replaced by statically inlinable bound values. */ std::unordered_map inlined_bindings_; + /*! \brief Symbolic primitive arguments classified after the complete boundary is known. */ + ffi::Array deferred_prim_args_; /*! \brief The number of parameters reserved for constants */ int n_param_for_const_ = 0; /*! \brief The output vars */ diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py index 715c2d3a65e8..850f34888d91 100644 --- a/tests/python/relax/test_transform_fuse_ops.py +++ b/tests/python/relax/test_transform_fuse_ops.py @@ -1352,11 +1352,11 @@ def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"): _check(Before, Expected) -def test_symbolic_prim_arg_bound_by_tensor_shape(): +def test_symbolic_prim_arg_before_tensor_arg(): @I.ir_module(s_tir=True) class Before: @T.prim_func(private=True, s_tir=True) - def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle): + def add_one(n: T.int64, x_handle: T.handle, out_handle: T.handle): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) x = T.match_buffer(x_handle, (T.int64(1), n), "float32") out = T.match_buffer(out_handle, (T.int64(1), n), "float32") @@ -1366,7 +1366,7 @@ def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle): out[0, vi] = x[0, vi] + T.float32(1) @T.prim_func(private=True, s_tir=True) - def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle): + def exp(n: T.int64, x_handle: T.handle, out_handle: T.handle): T.func_attr({"op_pattern": 0, "tirx.noalias": True}) x = T.match_buffer(x_handle, (T.int64(1), n), "float32") out = T.match_buffer(out_handle, (T.int64(1), n), "float32") @@ -1384,12 +1384,12 @@ def main( with R.dataflow(): lv = R.call_tir( cls.add_one, - (x, n), + (n, x), out_ty=R.Tensor((1, n), dtype="float32"), ) gv = R.call_tir( cls.exp, - (lv, n), + (n, lv), out_ty=R.Tensor((1, n), dtype="float32"), ) R.output(gv) @@ -1397,6 +1397,7 @@ def main( mod = relax.transform.AnnotateTIROpPattern()(Before) mod = relax.transform.FuseOps()(mod) + assert relax.analysis.check_well_formed(mod) fused = next( mod[global_var] @@ -1404,8 +1405,18 @@ def main( if global_var.name_hint.startswith("fused_") ) assert len(fused.params) == 1 + assert not isinstance(fused.params[0].ty, tvm.ir.PrimType) assert fused.ret_ty.shape is not None + mod = relax.transform.FuseTIR()(mod) + assert relax.analysis.check_well_formed(mod) + fused_tir = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert tvm.tirx.analysis.verify_well_formed(fused_tir) + def test_symbolic_prim_arg_reused_from_derived_tensor_shape(): @I.ir_module(s_tir=True) @@ -1546,6 +1557,207 @@ def main( assert sum(isinstance(param.ty, tvm.ir.PrimType) for param in fused.params) == 1 +def test_primitive_call_arg_not_inlined(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def add_scalar(x_handle: T.handle, value: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(4),), "int64") + out = T.match_buffer(out_handle, (T.int64(4),), "int64") + for i in range(4): + with T.sblock("add_scalar"): + vi = T.axis.spatial(4, i) + out[vi] = x[vi] + value + + @T.prim_func(private=True, s_tir=True) + def double(x_handle: T.handle, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(4),), "int64") + out = T.match_buffer(out_handle, (T.int64(4),), "int64") + for i in range(4): + with T.sblock("double"): + vi = T.axis.spatial(4, i) + out[vi] = x[vi] * T.int64(2) + + @R.function + def main(x: R.Tensor((4,), dtype="int64")): + cls = Before + with R.dataflow(): + value: R.Prim("int64") = R.call_pure_packed("get_scalar", ty_args=R.Prim("int64")) + lv = R.call_tir( + cls.add_scalar, + (x, value), + out_ty=R.Tensor((4,), dtype="int64"), + ) + gv = R.call_tir( + cls.double, + (lv,), + out_ty=R.Tensor((4,), dtype="int64"), + ) + R.output(gv, value) + return gv, value + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + assert relax.analysis.check_well_formed(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 2 + assert sum(isinstance(param.ty, tvm.ir.PrimType) for param in fused.params) == 1 + + mod = relax.transform.FuseTIR()(mod) + assert relax.analysis.check_well_formed(mod) + fused_tir = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert tvm.tirx.analysis.verify_well_formed(fused_tir) + + +def test_primitive_call_arg_used_by_output_shape_not_inlined(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def make(n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + out = T.match_buffer(out_handle, (n,), "float32") + for i in range(n): + with T.sblock("make"): + vi = T.axis.spatial(n, i) + out[vi] = T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def double(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (n,), "float32") + out = T.match_buffer(out_handle, (n,), "float32") + for i in range(n): + with T.sblock("double"): + vi = T.axis.spatial(n, i) + out[vi] = x[vi] * T.float32(2) + + @R.function(pure=False) + def main(): + cls = Before + n: R.Prim("int64") = R.call_packed("get_extent", ty_args=R.Prim("int64")) + with R.dataflow(): + lv = R.call_tir( + cls.make, + (n,), + out_ty=R.Tensor((n,), dtype="float32"), + ) + gv = R.call_tir( + cls.double, + (lv, n), + out_ty=R.Tensor((n,), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + assert relax.analysis.check_well_formed(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 1 + assert isinstance(fused.params[0].ty, tvm.ir.PrimType) + + packed_calls = [] + + def collect_packed_calls(expr): + if ( + isinstance(expr, relax.Call) + and isinstance(expr.op, relax.ExternFunc) + and expr.op.global_symbol == "get_extent" + ): + packed_calls.append(expr) + + relax.analysis.post_order_visit(mod["main"], collect_packed_calls) + assert len(packed_calls) == 1 + + packed_calls.clear() + relax.analysis.post_order_visit( + fused, + collect_packed_calls, + ) + assert not packed_calls + + +def test_symbolic_prim_arg_used_only_by_output_shape(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def make(n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + out = T.match_buffer(out_handle, (n,), "float32") + for i in range(n): + with T.sblock("make"): + vi = T.axis.spatial(n, i) + out[vi] = T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def double(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (n,), "float32") + out = T.match_buffer(out_handle, (n,), "float32") + for i in range(n): + with T.sblock("double"): + vi = T.axis.spatial(n, i) + out[vi] = x[vi] * T.float32(2) + + @R.function + def main( + source: R.Tensor(("n",), dtype="float32"), + ) -> R.Tensor(("n",), dtype="float32"): + n = T.int64() + cls = Before + with R.dataflow(): + lv = R.call_tir( + cls.make, + (n,), + out_ty=R.Tensor((n,), dtype="float32"), + ) + gv = R.call_tir( + cls.double, + (lv, n), + out_ty=R.Tensor((n,), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + assert relax.analysis.check_well_formed(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 1 + assert isinstance(fused.params[0].ty, relax.ShapeType) + assert fused.ret_ty.shape is not None + + mod = relax.transform.FuseTIR()(mod) + assert relax.analysis.check_well_formed(mod) + fused_tir = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert tvm.tirx.analysis.verify_well_formed(fused_tir) + + def test_shape_expr_arg(): @I.ir_module(s_tir=True) class Before: From a9137fc7ae8866875f80c511d3b32b3a152b3e91 Mon Sep 17 00:00:00 2001 From: Akaash Parthasarathy Date: Wed, 26 Aug 2026 09:56:38 -0400 Subject: [PATCH 3/3] [TEST][RELAX] Cover both symbolic fusion argument orders --- tests/python/relax/test_transform_fuse_ops.py | 66 +++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py index 850f34888d91..bee21e0686f6 100644 --- a/tests/python/relax/test_transform_fuse_ops.py +++ b/tests/python/relax/test_transform_fuse_ops.py @@ -1352,6 +1352,72 @@ def main(s: R.Shape(["n"])) -> R.Tensor((1, 1, "n", "n"), dtype="float32"): _check(Before, Expected) +def test_symbolic_prim_arg_after_tensor_arg(): + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def add_one(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n), "float32") + for i in range(n): + with T.sblock("add_one"): + vi = T.axis.spatial(n, i) + out[0, vi] = x[0, vi] + T.float32(1) + + @T.prim_func(private=True, s_tir=True) + def exp(x_handle: T.handle, n: T.int64, out_handle: T.handle): + T.func_attr({"op_pattern": 0, "tirx.noalias": True}) + x = T.match_buffer(x_handle, (T.int64(1), n), "float32") + out = T.match_buffer(out_handle, (T.int64(1), n), "float32") + for i in range(n): + with T.sblock("exp"): + vi = T.axis.spatial(n, i) + out[0, vi] = T.exp(x[0, vi]) + + @R.function + def main( + x: R.Tensor((1, "n"), dtype="float32"), + ) -> R.Tensor((1, "n"), dtype="float32"): + n = T.int64() + cls = Before + with R.dataflow(): + lv = R.call_tir( + cls.add_one, + (x, n), + out_ty=R.Tensor((1, n), dtype="float32"), + ) + gv = R.call_tir( + cls.exp, + (lv, n), + out_ty=R.Tensor((1, n), dtype="float32"), + ) + R.output(gv) + return gv + + mod = relax.transform.AnnotateTIROpPattern()(Before) + mod = relax.transform.FuseOps()(mod) + assert relax.analysis.check_well_formed(mod) + + fused = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert len(fused.params) == 1 + assert not isinstance(fused.params[0].ty, tvm.ir.PrimType) + assert fused.ret_ty.shape is not None + + mod = relax.transform.FuseTIR()(mod) + assert relax.analysis.check_well_formed(mod) + fused_tir = next( + mod[global_var] + for global_var in mod.get_global_vars() + if global_var.name_hint.startswith("fused_") + ) + assert tvm.tirx.analysis.verify_well_formed(fused_tir) + + def test_symbolic_prim_arg_before_tensor_arg(): @I.ir_module(s_tir=True) class Before: