[Unity] Support pattern-based rewriting

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()

python/tvm/relax/dpl/pattern.py

@@ -20,7 +20,7 @@
 # pylint: disable=pointless-statement
 
 import typing
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union, Callable
 
 import tvm
 import tvm._ffi as tvm_ffi
@@ -31,7 +31,7 @@
 from ...ir import make_node
 from ...ir.base import Node
 from ...runtime import Object
-from ..expr import Expr, Var
+from ..expr import Expr, Var, Function
 from . import _ffi as ffi
 
 
@@ -1115,3 +1115,39 @@ def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None
  return is_op(activation)(out)
 
  return out
+
+
+def rewrite(
+ pattern: DFPattern, rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr], func: Function
+) -> Function:
+ """
+ Rewrite a function with the given pattern and the rewriter function.
+
+ Parameters
+ ----------
+ pattern: DFPattern
+ The pattern to match.
+
+ rewriter: Callable[[Expr, Dict[DFPattern, Expr]], Expr]
+ The function to be called on a successful matching for rewriting. Given the matched
+ call node and the map of patterns and matched expressions, it should return a new call node
+ to replace the original one or the original matched call node as is.
+
+ For example, to replace x + x with 2 * x, we can write the rewriter as follows:
+ ```
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(orig, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+ ```
+
+ func: Function
+ The function to rewrite.
+
+ Returns
+ -------
+ rewritten_func: Function
+ The rewritten or the input function, depending on the pattern matching result.
+ """
+ return ffi.rewrite(pattern, rewriter, func)

src/relax/ir/dataflow_matcher.cc

@@ -766,5 +766,53 @@ Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& d
 
 TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
 
+/*!
+ * \brief Apply pattern matching to each call node and replace matching ones with the output of
+ * a user-provided rewriter function.
+ */
+class PatternRewriter : ExprMutator {
+ public:
+ using ExprMutator::VisitExpr_;
+
+ PatternRewriter(DFPattern pat, PackedFunc rewriter_func)
+ : pattern_(pat), rewriter_func_(rewriter_func) {}
+
+ static Expr Run(DFPattern pat, PackedFunc rewriter_func, Function f) {
+ PatternRewriter rewriter(pat, rewriter_func);
+ return RemoveAllUnused(Downcast<Function>(rewriter.VisitExpr(f)));
+ }
+
+ void VisitBinding_(const VarBindingNode* binding) final {
+ bindings_.Set(binding->var, binding->value);
+ ExprMutator::VisitBinding_(binding);
+ if (auto it = memo_.find(binding->value.get()); it != memo_.end()) {
+ // We need to update the binding to pass to ExtractMatchedExpr, so that the rewritten
+ // expression can be subject to further pattern matchings.
+ bindings_.Set(binding->var, it->second);
+ }
+ }
+
+ Expr VisitExpr_(const CallNode* call_node) final {
+ auto call = ExprMutator::VisitExpr_(call_node);
+ if (auto matches_opt = ExtractMatchedExpr(pattern_, call, bindings_)) {
+ auto rewriten_expr = rewriter_func_(call, matches_opt.value());
+ memo_[call_node] = rewriten_expr;
+ return rewriten_expr;
+ }
+ return call;
+ }
+
+ private:
+ DFPattern pattern_;
+ PackedFunc rewriter_func_;
+ Map<Var, Expr> bindings_;
+ std::unordered_map<const Object*, Expr> memo_;
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.rewrite")
+ .set_body_typed([](DFPattern pat, PackedFunc rewriter, Function f) {
+ return PatternRewriter::Run(pat, rewriter, f);
+ });
+
 } // namespace relax
 } // namespace tvm

tests/python/relax/test_dataflow_pattern.py

@@ -888,5 +888,123 @@ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
  assert not ctx1.match_dfb(simple_chain.body.blocks[0])
 
 
+def test_rewrite_simple():
+ @R.function
+ def main(x: R.Tensor((16, 16), "float32")) -> R.Tensor((16, 16), "float32"):
+ with R.dataflow():
+ x2 = R.add(x, x)
+ x4 = R.add(x2, x2)
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected1(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(2, "float32"))
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(lv, R.const(2, "float32"))
+ R.output(x4)
+ return x4
+
+ @R.function
+ def expected2(x: R.Tensor((16, 16), dtype="float32")) -> R.Tensor((16, 16), dtype="float32"):
+ with R.dataflow():
+ x4: R.Tensor((16, 16), dtype="float32") = R.multiply(x, R.const(4, "float32"))
+ R.output(x4)
+ return x4
+
+ x = wildcard()
+ pattern = is_op("relax.add")(x, x)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(2, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected1)
+
+ add1 = is_op("relax.add")(x, x)
+ pattern = is_op("relax.add")(add1, add1)
+
+ def rewriter(_, matchings):
+ return R.multiply(matchings[x], R.const(4, "float32"))
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected2)
+
+ # No rewriting, return the original call node as is
+ def rewriter(orig, _):
+ return orig
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, main)
+
+
+def test_rewrite_attention():
+ @R.function
+ def main(
+ Q: R.Tensor((2, 4096, 8, 40), "float32"),
+ K: R.Tensor((2, 4096, 8, 40), "float32"),
+ V: R.Tensor((2, 4096, 8, 40), "float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), "float32"):
+ with R.dataflow():
+ lv58 = R.permute_dims(Q, axes=[0, 2, 1, 3])
+ lv59 = R.reshape(lv58, R.shape([16, 4096, 40]))
+
+ lv61 = R.permute_dims(K, axes=[0, 2, 1, 3])
+ lv62 = R.reshape(lv61, R.shape([16, 4096, 40]))
+
+ lv64 = R.permute_dims(V, axes=[0, 2, 1, 3])
+ lv65 = R.reshape(lv64, R.shape([16, 4096, 40]))
+
+ lv62_transposed = R.permute_dims(lv62, axes=[0, 2, 1])
+ lv3_1 = R.matmul(lv59, lv62_transposed)
+ lv68 = R.multiply(lv3_1, R.const(0.15811388194561005, "float32"))
+ lv69 = R.nn.softmax(lv68, axis=-1)
+ lv_3 = R.matmul(lv69, lv65)
+
+ lv71 = R.reshape(lv_3, R.shape([2, 8, 4096, 40]))
+ lv72 = R.permute_dims(lv71, axes=[0, 2, 1, 3])
+ R.output(lv72)
+
+ return lv72
+
+ @R.function
+ def expected(
+ Q: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ K: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ V: R.Tensor((2, 4096, 8, 40), dtype="float32"),
+ ) -> R.Tensor((2, 4096, 8, 40), dtype="float32"):
+ with R.dataflow():
+ lv72: R.Tensor((2, 4096, 8, 40), dtype="float32") = R.nn.attention(Q, V, K)
+ R.output(lv72)
+ return lv72
+
+ def BSNH_to_BSH(tensor):
+ return is_op("relax.reshape")(is_op("relax.permute_dims")(tensor), wildcard())
+
+ def BSH_to_BSNH(tensor):
+ return is_op("relax.permute_dims")(is_op("relax.reshape")(tensor, wildcard()))
+
+ Q = wildcard()
+ K = wildcard()
+ V = wildcard()
+
+ Q_3D = BSNH_to_BSH(Q)
+ V_3D = BSNH_to_BSH(V)
+ K_3D = BSNH_to_BSH(K)
+
+ matmul1 = is_op("relax.matmul")(Q_3D, is_op("relax.permute_dims")(V_3D))
+ multiply = is_op("relax.multiply")(matmul1, is_const())
+ softmax = is_op("relax.nn.softmax")(multiply)
+ matmul2 = is_op("relax.matmul")(softmax, K_3D)
+
+ pattern = BSH_to_BSNH(matmul2)
+
+ def rewriter(_, matchings):
+ return R.nn.attention(matchings[Q], matchings[K], matchings[V])
+
+ rewritten = rewrite(pattern, rewriter, main)
+ tvm.ir.assert_structural_equal(rewritten, expected)
+
+
 if __name__ == "__main__":
  tvm.testing.main()