diff --git a/src/relax/transform/fuse_ops.cc b/src/relax/transform/fuse_ops.cc index 5ebe36373b71..a6ba798dcd38 100644 --- a/src/relax/transform/fuse_ops.cc +++ b/src/relax/transform/fuse_ops.cc @@ -28,6 +28,7 @@ */ #include +#include #include #include #include @@ -106,6 +107,110 @@ constexpr uint32_t kMaxFusedOps = 256; TVM_REGISTER_PASS_CONFIG_OPTION("relax.FuseOps.max_depth", int64_t); +// Shape symbols are not ordinary dataflow variables. Record the match_cast that +// first defines each symbol, so both partitioning and scheduling see its uses. +class SymbolicDependencyCollector { + public: + using DependencyMap = std::unordered_map>; + using VisitResult = ffi::Expected>; + + static DependencyMap Collect(const Function& func) { + SymbolicDependencyCollector collector; + ffi::StructuralVisit( + func, + [&](const Function& function, ffi::StructuralVisitorObj* visitor) -> VisitResult { + auto saved_producers = collector.producers_; + auto saved_binding = collector.current_binding_; + collector.current_binding_ = nullptr; + for (const Var& param : function->params) { + // Scalar parameters can themselves be used as shape symbols. + collector.producers_.emplace(param.get(), nullptr); + collector.DefineSymbols(GetType(param)); + } + auto result = visitor->VisitExpected(function->body); + collector.current_binding_ = saved_binding; + collector.producers_ = std::move(saved_producers); + return result; + }, + [&](const If& if_expr, ffi::StructuralVisitorObj* visitor) -> VisitResult { + visitor->Visit(if_expr->cond); + auto saved_producers = collector.producers_; + auto saved_binding = collector.current_binding_; + // Branch-local definitions are invisible to the other branch and to + // the enclosing binding, including uses in the branch result type. + collector.current_binding_ = nullptr; + visitor->Visit(if_expr->true_branch); + collector.producers_ = saved_producers; + visitor->Visit(if_expr->false_branch); + collector.producers_ = std::move(saved_producers); + collector.current_binding_ = saved_binding; + return visitor->VisitExpected(GetType(if_expr)); + }, + [&](const Binding& binding, ffi::StructuralVisitorObj* visitor) -> VisitResult { + auto saved_binding = collector.current_binding_; + collector.current_binding_ = binding->var.get(); + visitor->Visit(GetBoundValue(binding)); + visitor->Visit(GetType(binding->var)); + if (const auto* match_cast = binding.as()) { + // Existing symbols in the value or target type are uses, not redefinitions. + visitor->Visit(match_cast->ty); + collector.DefineSymbols(match_cast->ty); + } + collector.current_binding_ = saved_binding; + return std::nullopt; + }, + [&](const Var& var, ffi::StructuralVisitorObj* visitor) -> VisitResult { + collector.UseVar(var); + return visitor->VisitExpected(GetType(var)); + }, + [](const FuncType&, ffi::StructuralVisitorObj*) -> VisitResult { + // Function-type symbols are locally bound, not caller dependencies. + return std::nullopt; + }); + return collector.dependencies_; + } + + private: + void UseVar(const Var& var) { + auto it = producers_.find(var.get()); + if (current_binding_ && it != producers_.end() && it->second && + it->second != current_binding_) { + auto& deps = dependencies_[current_binding_]; + Var producer = ffi::GetRef(it->second); + if (std::none_of(deps.begin(), deps.end(), + [&](const Var& dep) { return dep.same_as(producer); })) { + deps.push_back(producer); + } + } + } + + void DefineSymbols(const Type& ty) { + auto define_shape = [&](const ffi::Array& values) { + for (const PrimExpr& value : values) { + // Only a bare symbol is definable; compound expressions are constraints. + if (const auto* var = value.as()) { + producers_.emplace(var, current_binding_); + } + } + }; + ffi::StructuralWalk( + ty, + [](const FuncType&) -> ffi::Expected { return ffi::WalkResult::Skip(); }, + [&](const ShapeType& shape) -> ffi::Expected { + if (shape->values.has_value()) define_shape(shape->values.value()); + return ffi::WalkResult::Skip(); + }, + [&](const ShapeExpr& shape) -> ffi::Expected { + define_shape(shape->values); + return ffi::WalkResult::Skip(); + }); + } + + const VarNode* current_binding_{nullptr}; + std::unordered_map producers_; + DependencyMap dependencies_; +}; + class GraphCreator : public ExprVisitor { public: /*! @@ -130,7 +235,9 @@ class GraphCreator : public ExprVisitor { func->GetAttr(attr::kCodegen).has_value()) { continue; } - creator(ffi::GetRef(func)); + auto function = ffi::GetRef(func); + creator.symbolic_deps_ = SymbolicDependencyCollector::Collect(function); + creator(function); } // The algorithm of the graph creator ensures that each created node will be added to the @@ -167,7 +274,8 @@ class GraphCreator : public ExprVisitor { void VisitBinding_(const MatchCastNode* binding) final { IndexedForwardGraph::Node* node = CreateNode(binding->var.get()); - SetNodePattern(node, OpPatternKind::kOpaque); + VisitUnsupportedNode(binding->value, node); + AddSymbolicDependencies(binding->var, node); AddToPostDFSOrder(node, binding->var.get()); } @@ -190,11 +298,18 @@ class GraphCreator : public ExprVisitor { // Case 3. The type of the expression is not fusion-supported. // In this case, we skip adding edges, adding an empty node into graph. } + AddSymbolicDependencies(binding->var, node); AddToPostDFSOrder(node, binding->var.get()); } /********** Non-Leaf Expression Nodes **********/ + void AddSymbolicDependencies(const Var& var, IndexedForwardGraph::Node* node) { + for (const Var& producer : symbolic_deps_[var.get()]) { + AddEdge(graph_.node_map.at(producer.get()), node, OpPatternKind::kOpaque); + } + } + void VisitCall(const CallNode* call, IndexedForwardGraph::Node* binding_var_node) { TVM_FFI_ICHECK_NOTNULL(binding_var_node); @@ -381,6 +496,8 @@ class GraphCreator : public ExprVisitor { std::unordered_set initialized_nodes_; /*! \brief The model params in the function input */ std::unordered_set input_params_; + /*! \brief Dependencies on bindings that define shape symbols. */ + SymbolicDependencyCollector::DependencyMap symbolic_deps_; }; /*! @@ -839,6 +956,7 @@ class OperatorFusor : public ExprMutator { if (func->IsInstance() && !func->HasNonzeroAttr(attr::kPrimitive) && !func->GetAttr(attr::kCodegen).has_value()) { outer_bindings_ = AnalyzeVar2Value(func); + symbolic_deps_ = SymbolicDependencyCollector::Collect(func.as_or_throw()); auto updated_func = VisitExpr(func).as_or_throw(); builder_->UpdateFunction(gv, updated_func); outer_bindings_ = {}; @@ -862,6 +980,7 @@ class OperatorFusor : public ExprMutator { BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) final { group2func_.clear(); + group_deps_.clear(); // Step 1. Collect the bindings for each grouped function. CollectFuncBindings(block->bindings); @@ -1011,20 +1130,11 @@ class OperatorFusor : public ExprMutator { // - If the var's group is same as the binding's, the var is defined in the same group // - If the var's group is different with the binding's, the var must be the output from // another group. Mark it to be the group output. - auto update_boundary = [this, binding, &cur_group](const Expr& e) { + auto update_boundary = [this, &cur_group](const Expr& e) { if (e->IsInstance() && obj2group_.count(e.get())) { const Var& used_var = e.as_or_throw(); Group* producer_group = GetGroupFromVar(used_var); - // Only check those group defined before. - // Skip the vars from input or groups with single binding. - if (producer_group != cur_group) { - for (Group* depgroup : group_deps_[producer_group]) { - TVM_FFI_ICHECK(depgroup != cur_group) - << "A cyclic dependency detected between the groups " << binding->var->name - << " and " << used_var->name << " are in."; - } - group_deps_[cur_group].push_back(producer_group); - } + AddGroupDependency(cur_group, producer_group); if (auto producer = group2func_.find(producer_group); producer_group != cur_group && producer != group2func_.end()) { @@ -1040,6 +1150,21 @@ class OperatorFusor : public ExprMutator { TVM_FFI_ICHECK_NOTNULL(match_cast); PostOrderVisit(match_cast->value, update_boundary); } + + // Shape dependencies constrain scheduling without adding tensor outputs or + // parameters to the grouped functions. + for (const Var& producer : symbolic_deps_[binding->var.get()]) { + if (obj2group_.count(producer.get())) { + AddGroupDependency(cur_group, GetGroupFromVar(producer)); + } + } + } + } + + void AddGroupDependency(Group* consumer, Group* producer) { + auto& deps = group_deps_[consumer]; + if (consumer != producer && std::find(deps.begin(), deps.end(), producer) == deps.end()) { + deps.push_back(producer); } } @@ -1094,14 +1219,18 @@ class OperatorFusor : public ExprMutator { } std::unordered_set visited; + std::unordered_set visiting; std::function)> dfs_visit; - dfs_visit = [this, &visited, &dfs_visit](Group* g, auto leaf_fun) { + dfs_visit = [this, &visited, &visiting, &dfs_visit](Group* g, auto leaf_fun) { if (!visited.count(g)) { - visited.insert(g); + TVM_FFI_ICHECK(visiting.insert(g).second) + << "A cyclic dependency detected between fusion groups."; for (auto dep : group_deps_[g]) { dfs_visit(dep, leaf_fun); } + visiting.erase(g); + visited.insert(g); leaf_fun(g); } }; @@ -1129,6 +1258,8 @@ class OperatorFusor : public ExprMutator { std::unordered_map group2func_; /*! \brief Bindings visible while rewriting the current Relax function. */ ffi::Map outer_bindings_; + /*! \brief Dependencies on bindings that define shape symbols. */ + SymbolicDependencyCollector::DependencyMap symbolic_deps_; /*! * \brief A map from a group to its dependent groups, used to detect cyclic dependencies. * \note Use vector so we can be deterministic, there won't be a lot of dep groups so diff --git a/tests/python/relax/test_frontend_onnx.py b/tests/python/relax/test_frontend_onnx.py index 92f62d32f66d..004ae17a623a 100644 --- a/tests/python/relax/test_frontend_onnx.py +++ b/tests/python/relax/test_frontend_onnx.py @@ -7636,6 +7636,80 @@ def main( # ) +@pytest.mark.parametrize("input_shape", [(1, 4, 5, 5), (2, 3, 4, 4)]) +def test_size_slice_reshape_with_fusion(input_shape): + """Regression for #20177: preserve the runtime Slice shape before allocation.""" + x = helper.make_tensor_value_info("x", TensorProto.FLOAT, input_shape) + y = helper.make_tensor_value_info("y", TensorProto.FLOAT, input_shape) + initializers = [ + helper.make_tensor("flat_shape", TensorProto.INT64, [1], [-1]), + helper.make_tensor("size_shape", TensorProto.INT64, [1], [1]), + helper.make_tensor("slice_starts", TensorProto.INT64, [1], [0]), + helper.make_tensor("slice_axes", TensorProto.INT64, [1], [0]), + helper.make_tensor("slice_steps", TensorProto.INT64, [1], [1]), + ] + nodes = [ + helper.make_node("Relu", ["x"], ["h0"]), + helper.make_node("Relu", ["h0"], ["positive"]), + helper.make_node("Reshape", ["x", "flat_shape"], ["flat"]), + helper.make_node("Size", ["flat"], ["size"]), + helper.make_node("Reshape", ["size", "size_shape"], ["size_1d"]), + helper.make_node( + "Slice", + ["flat", "slice_starts", "size_1d", "slice_axes", "slice_steps"], + ["projected"], + ), + helper.make_node("Shape", ["x"], ["x_shape"]), + helper.make_node("Reshape", ["projected", "x_shape"], ["splice"]), + helper.make_node("Add", ["positive", "splice"], ["y"]), + ] + graph = helper.make_graph( + nodes, + "vm_shape_lower_size_slice", + [x], + [y], + initializers, + ) + # Keep the model compatible with the ONNX Runtime versions used in CI. + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)], ir_version=8) + + onnx.checker.check_model(model) + session = onnxruntime.InferenceSession( + model.SerializeToString(), providers=["CPUExecutionProvider"] + ) + mod = from_onnx(model, opset=18, keep_params_in_input=True) + pipeline = tvm.transform.Sequential( + [ + relax.backend.DispatchSampling(), + relax.backend.DispatchSortScan(), + relax.transform.LegalizeOps(), + relax.transform.AnnotateTIROpPattern(), + relax.transform.FoldConstant(), + relax.transform.FuseOps(fuse_opt_level=2), + relax.transform.FuseTIR(), + relax.transform.RewriteDataflowReshape(), + relax.transform.ToNonDataflow(), + relax.transform.RemovePurityChecking(), + relax.transform.CallTIRRewrite(), + relax.transform.StaticPlanBlockMemory(), + relax.transform.LowerAllocTensor(), + relax.transform.KillAfterLastUse(), + relax.transform.LowerRuntimeBuiltin(), + relax.transform.ComputePrimValue(), + relax.transform.VMShapeLower(emit_err_ctx=True), + relax.transform.AttachGlobalSymbol(), + ] + ) + mod, params = relax.frontend.detach_params(mod) + executable = tvm.compile(mod, target="llvm", relax_pipeline=pipeline) + vm = relax.VirtualMachine(executable, tvm.cpu()) + data = (np.arange(np.prod(input_shape), dtype="float32") % 7 - 3).reshape(input_shape) + actual = vm["main"](tvm.runtime.tensor(data), *params.get("main", [])).numpy() + expected = session.run(None, {"x": data})[0] + tvm.testing.assert_allclose(actual, expected) + tvm.testing.assert_allclose(actual, np.maximum(data, 0) + data) + + def test_slice_dynamic_inputs_ir(): slice_node = helper.make_node("Slice", ["x", "starts", "ends", "axes", "steps"], ["y"]) diff --git a/tests/python/relax/test_transform_fuse_ops.py b/tests/python/relax/test_transform_fuse_ops.py index bee21e0686f6..884c7b8d1f92 100644 --- a/tests/python/relax/test_transform_fuse_ops.py +++ b/tests/python/relax/test_transform_fuse_ops.py @@ -17,6 +17,8 @@ # ruff: noqa: E501, F841 +import pytest + import tvm import tvm.testing from tvm import relax, topi @@ -1878,6 +1880,151 @@ def main(s: R.Shape(["n"]), kv_cache: R.Any): _check(Before, Expected) +@pytest.mark.parametrize("match_tensor", [False, True]) +def test_match_cast_shape_dependency(match_tensor): + """A symbol used only in call_tir's output type still depends on match_cast.""" + + @I.ir_module(s_tir=True) + class Before: + @T.prim_func(private=True, s_tir=True) + def copy(x: T.Buffer((8,), "float32"), out_handle: T.handle): + T.func_attr({"op_pattern": 8, "tirx.noalias": True}) + n = T.int64() + out = T.match_buffer(out_handle, (n,), "float32") + for i in range(n): + with T.sblock("copy"): + vi = T.axis.spatial(n, i) + out[vi] = x[vi] + + @R.function + def main(x: R.Tensor((8,), "float32"), shape: R.Shape(ndim=1)): + s = T.int64() + with R.dataflow(): + positive = R.nn.relu(x) + positive2 = R.nn.relu(positive) + bound = R.match_cast( + positive if match_tensor else shape, + R.Tensor((s,), "float32") if match_tensor else R.Shape([s]), + ) + copied = R.call_tir(Before.copy, (x,), out_ty=R.Tensor((s,), "float32")) + reshaped = R.reshape(copied, (8,)) + out = R.add(positive2, reshaped) + R.output(out) + return out + + mod = relax.transform.LegalizeOps()(Before) + mod = relax.transform.AnnotateTIROpPattern()(mod) + mod = relax.transform.FuseOps(fuse_opt_level=2)(mod) + assert relax.analysis.check_well_formed(mod) + bindings = mod["main"].body.blocks[0].bindings + match_index = next(i for i, b in enumerate(bindings) if isinstance(b, relax.MatchCast)) + copy_index = next( + i + for i, b in enumerate(bindings) + if isinstance(b.value, relax.Call) + and b.value.op == tvm.ir.Op.get("relax.call_tir") + and b.value.args[0].name_hint == "copy" + ) + assert match_index < copy_index + # If match_cast consumes the first Relu, fusing it with Add would create a + # cycle through the copy's shape dependency. The second Relu can still fuse. + fused_name = "fused_relu_add" if match_tensor else "fused_relu_relu_add" + assert any(gv.name_hint.startswith(fused_name) for gv in mod.get_global_vars()) + mod = relax.transform.FuseTIR()(mod) + assert relax.analysis.check_well_formed(mod) + + +@pytest.mark.parametrize( + "symbol_source", ["match_cast", "parameter", "tensor_parameter", "function_type", "outer_block"] +) +@pytest.mark.parametrize("use_shape", [False, True]) +def test_match_cast_symbol_scope(symbol_source, use_shape): + """Keep the original definition across repeated casts, blocks, and functions.""" + bb = relax.BlockBuilder() + for name in ["main", "other"]: + # Symbols with the same name must be tracked by identity. + m = tvm.tirx.Var("s", "int64") + n = tvm.tirx.Var("s", "int64") + x = relax.Var("x", R.Tensor((8,), "float32")) + shape = relax.Var( + "shape", R.Shape([m, n]) if symbol_source == "parameter" else R.Shape(ndim=2) + ) + params = [x, shape] + if symbol_source == "tensor_parameter": + params.append(relax.Var("known_shape", R.Tensor((m, n), "float32"))) + if symbol_source == "function_type": + # These symbols are bound inside the callable's signature and must + # not hide the caller's match_cast definitions, even if shared. + tensor_ty = R.Tensor((m, n), "float32") + params.append(relax.Var("callback", relax.FuncType([tensor_ty], tensor_ty))) + with bb.function(name, params): + if symbol_source == "outer_block": + with bb.dataflow(): + bound = bb.match_cast(shape, R.Shape([m, n])) + bb.emit_output(bound) + bb.emit(relax.op.shape_of(x)) + with bb.dataflow(): + positive = bb.emit(relax.op.nn.relu(x)) + positive2 = bb.emit(relax.op.nn.relu(positive)) + if symbol_source in ["match_cast", "function_type"]: + bb.match_cast(shape, R.Shape([m, n])) + arg = relax.ShapeExpr([m * n]) if use_shape else m + n + value = bb.emit( + relax.call_pure_packed( + "test.symbolic_arg", arg, ty_args=R.Tensor((8,), "float32") + ) + ) + # This cast only checks existing symbols. Treating it as a new + # producer would create a cycle with the fused Relu/Add group. + bb.match_cast(positive, R.Tensor((m * n,), "float32")) + bb.match_cast(shape, R.Shape([m, n])) + out = bb.emit_output(relax.op.add(positive2, value)) + bb.emit_func_output(out) + + mod = relax.transform.LegalizeOps()(bb.get()) + mod = relax.transform.AnnotateTIROpPattern()(mod) + mod = relax.transform.FuseOps(fuse_opt_level=2)(mod) + assert relax.analysis.check_well_formed(mod) + for name in ["main", "other"]: + if symbol_source == "outer_block": + assert len(mod[name].body.blocks) == 3 + bindings = [binding for block in mod[name].body.blocks for binding in block.bindings] + calls = [b for b in bindings if isinstance(b.value, relax.Call)] + symbolic_call = next( + b for b in calls if b.value.op == tvm.ir.Op.get("relax.call_pure_packed") + ) + if symbol_source in ["match_cast", "function_type", "outer_block"]: + assert any( + isinstance(b, relax.MatchCast) for b in bindings[: bindings.index(symbolic_call)] + ) + assert sum(isinstance(b, relax.MatchCast) for b in bindings) == ( + 3 if symbol_source in ["match_cast", "function_type", "outer_block"] else 2 + ) + + +def test_match_cast_symbol_scope_in_if(): + """Branch-local symbols must not become producers for later bindings.""" + + @I.ir_module + class Before: + @R.function + def main(x: R.Tensor((8,), "float32"), shape: R.Shape(ndim=1), cond: R.Tensor((), "bool")): + n = T.int64() + if cond: + true_bound = R.match_cast(shape, R.Shape([n])) + branch_out = x + else: + false_bound = R.match_cast(shape, R.Shape([n])) + branch_out = x + with R.dataflow(): + bound = R.match_cast(shape, R.Shape([n])) + out = R.call_pure_packed("test.symbolic_arg", n, ty_args=R.Tensor((8,), "float32")) + R.output(out) + return out + + _check(Before, Before) + + def test_skipping_match_cast(): @I.ir_module(s_tir=True) class Module: diff --git a/tests/python/relax/test_transform_fuse_ops_by_pattern.py b/tests/python/relax/test_transform_fuse_ops_by_pattern.py index 7f18a4a07595..87a5ec616fee 100644 --- a/tests/python/relax/test_transform_fuse_ops_by_pattern.py +++ b/tests/python/relax/test_transform_fuse_ops_by_pattern.py @@ -1460,6 +1460,58 @@ def main( assert after["main"].body.body.same_as(grouped_result) +def test_match_cast_checks_scalar_parameter(): + """Checking a scalar parameter must not create a cyclic shape dependency.""" + + @I.ir_module + class Before: + @R.function + def main(x: R.Tensor((8,), "float32"), n: T.int64): + with R.dataflow(): + positive = R.nn.relu(x) + positive2 = R.nn.relu(positive) + bound = R.match_cast(positive, R.Tensor((n,), "float32")) + value = R.call_pure_packed( + "test.symbolic_arg", n, ty_args=R.Tensor((8,), "float32") + ) + out = R.add(positive2, value) + R.output(out) + return out + + @I.ir_module + class Expected: + @R.function(private=True) + def fused_relax_nn_relu_relax_nn_relu_relax_add( + x: R.Tensor((8,), "float32"), value: R.Tensor((8,), "float32") + ): + R.func_attr({"Composite": "test.relu_relu_add", "Primitive": True}) + with R.dataflow(): + positive = R.nn.relu(x) + positive2 = R.nn.relu(positive) + out = R.add(positive2, value) + R.output(positive, out) + return (out, positive) + + @R.function + def main(x: R.Tensor((8,), "float32"), n: T.int64): + cls = Expected + with R.dataflow(): + value = R.call_pure_packed( + "test.symbolic_arg", n, ty_args=R.Tensor((8,), "float32") + ) + fused = cls.fused_relax_nn_relu_relax_nn_relu_relax_add(x, value) + positive = fused[1] + out = fused[0] + bound = R.match_cast(positive, R.Tensor((n,), "float32")) + R.output(out) + return out + + pattern = is_op("relax.add")( + is_op("relax.nn.relu")(is_op("relax.nn.relu")(wildcard())), wildcard() + ) + check(Before, [("test.relu_relu_add", pattern)], Expected) + + def test_inline_bound_static_shape_argument(): """A static leaf binding should not become a grouped-function parameter."""