(tvm) branch main updated: [Unity][Transform] Check for permute_dims in ExpandMatmulOfSum (#16590)

[lu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

lunderberg pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new 8f4259710a [Unity][Transform] Check for permute_dims in ExpandMatmulOfSum (#16590)
8f4259710a is described below

commit 8f4259710a298249416bdea2ba06380e06a6ec18
Author: Eric Lunderberg <Lu...@users.noreply.github.com>
AuthorDate: Thu Feb 22 11:19:39 2024 -0600

    [Unity][Transform] Check for permute_dims in ExpandMatmulOfSum (#16590)
    
    This pattern occurs whenever `relax.op.linear` is used.
---
 src/relax/transform/expand_matmul_of_sum.cc        | 18 ++++++++++++-
 .../relax/test_transform_expand_matmul_of_sum.py   | 30 ++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/src/relax/transform/expand_matmul_of_sum.cc b/src/relax/transform/expand_matmul_of_sum.cc
index 76ebae9498..9066205634 100644
--- a/src/relax/transform/expand_matmul_of_sum.cc
+++ b/src/relax/transform/expand_matmul_of_sum.cc
@@ -34,6 +34,7 @@
 
 #include "../op/tensor/binary.h"
 #include "../op/tensor/linear_algebra.h"
+#include "../op/tensor/manipulate.h"
 
 namespace tvm {
 namespace relax {
@@ -49,7 +50,11 @@ std::tuple<DFPattern, TypedPackedFunc<Expr(Expr, Map<DFPattern, Expr>)>> CreateP
 
   auto pat_rhs_a = WildcardPattern();
   auto pat_rhs_b = WildcardPattern();
-  auto pat_rhs = IsOp("relax.add")(pat_rhs_a, pat_rhs_b);
+  auto pat_rhs_sum = IsOp("relax.add")(pat_rhs_a, pat_rhs_b);
+
+  auto pat_rhs_permute_dims = IsOp("relax.permute_dims")(pat_rhs_sum);
+
+  auto pat_rhs = pat_rhs_sum | pat_rhs_permute_dims;
 
   auto pat_matmul = IsOp("relax.matmul")(pat_lhs, pat_rhs);
 
@@ -72,6 +77,17 @@ std::tuple<DFPattern, TypedPackedFunc<Expr(Expr, Map<DFPattern, Expr>)>> CreateP
       return expr;
     }
 
+    if (matches.count(pat_rhs_permute_dims)) {
+      auto call_permute = Downcast<Call>(matches[pat_rhs_permute_dims]);
+      auto attrs = call_permute->attrs.as<PermuteDimsAttrs>();
+      ICHECK(attrs) << "Operator permute_dims should have PermuteDimsAttrs, "
+                    << "but " << call_permute << " has attributes " << call_permute->attrs;
+      auto axes = attrs->axes;
+
+      rhs_a = permute_dims(rhs_a, axes);
+      rhs_b = permute_dims(rhs_b, axes);
+    }
+
     return add(matmul(lhs, rhs_a, DataType::Void()), matmul(lhs, rhs_b, DataType::Void()));
   };
 
diff --git a/tests/python/relax/test_transform_expand_matmul_of_sum.py b/tests/python/relax/test_transform_expand_matmul_of_sum.py
index 67e59225c5..b380d15842 100644
--- a/tests/python/relax/test_transform_expand_matmul_of_sum.py
+++ b/tests/python/relax/test_transform_expand_matmul_of_sum.py
@@ -123,5 +123,35 @@ class TestExpansionOfRuntimeAddition(Base):
             return out
 
 
+class TestRHSPermuteDims(Base):
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor([16], "float32"),
+            A: R.Tensor([32, 16], "float32"),
+            B: R.Tensor([32, 16], "float32"),
+        ) -> R.Tensor([32], "float32"):
+            linear_weight = R.add(A, B)
+            matmul_weight = R.permute_dims(linear_weight)
+            out = R.matmul(x, matmul_weight)
+            return out
+
+    @I.ir_module
+    class Expected:
+        @R.function
+        def main(
+            x: R.Tensor([16], "float32"),
+            A: R.Tensor([32, 16], "float32"),
+            B: R.Tensor([32, 16], "float32"),
+        ) -> R.Tensor([32], "float32"):
+            A_transpose = R.permute_dims(A)
+            lhs = R.matmul(x, A_transpose)
+            B_transpose = R.permute_dims(B)
+            rhs = R.matmul(x, B_transpose)
+            out = R.add(lhs, rhs)
+            return out
+
+
 if __name__ == "__main__":
     tvm.testing.main()