fix msc testcase

Author: mshr-h

python/tvm/contrib/msc/core/transform/pattern.py

@@ -330,15 +330,17 @@ def make_relax_attention_pattern() -> (
     q_trans = relax_pattern.is_op("relax.permute_dims")(weight_q)
     k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
     v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
-    out = relax_pattern.is_op("relax.nn.attention")(q_trans, k_trans, v_trans)
+    attention = relax_pattern.is_op("relax.nn.attention")(q_trans, k_trans, v_trans)
+    out = relax_pattern.is_op("relax.permute_dims")(attention)
     annotations = {
         "weight_q": weight_q,
         "weight_k": weight_k,
         "weight_v": weight_v,
         "q_trans": q_trans,
         "k_trans": k_trans,
         "v_trans": v_trans,
-        "attention": out,
+        "attention": attention,
+        "out": out,
     }
     return out, annotations
 
@@ -378,7 +380,8 @@ def make_relax_mask_attention_pattern() -> (
     q_trans = relax_pattern.is_op("relax.permute_dims")(weight_q)
     k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
     v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
-    out = relax_pattern.is_op("relax.nn.attention_bias")(q_trans, k_trans, v_trans, mask)
+    attention = relax_pattern.is_op("relax.nn.attention_bias")(q_trans, k_trans, v_trans, mask)
+    out = relax_pattern.is_op("relax.permute_dims")(attention)
     annotations = {
         "weight_q": weight_q,
         "weight_k": weight_k,
@@ -387,7 +390,8 @@ def make_relax_mask_attention_pattern() -> (
         "q_trans": q_trans,
         "k_trans": k_trans,
         "v_trans": v_trans,
-        "attention": out,
+        "attention": attention,
+        "out": out,
     }
     return out, annotations
 

tests/python/contrib/test_msc/test_graph_build.py

@@ -2363,13 +2363,13 @@ def forward(self, q_data, k_data, v_data):
         ],
         "outputs": [
             {
-                "name": "attention",
-                "shape": [1, seq, 8, 64],
+                "name": "permute_dims_3",
+                "shape": [1, 8, seq, 64],
                 "dtype": "float32",
-                "layout": "ABCD",
+                "layout": "ACBD",
             }
         ],
-        "nodes": {"total": 4, "input": 3, "msc.attention": 1},
+        "nodes": {"total": 5, "input": 3, "msc.attention": 1, "permute_dims": 1},
     }
     if dynamic:
         expected1["prims"] = {"total": 1, "shape": 1}
@@ -2395,13 +2395,13 @@ def forward(self, q_data, k_data, v_data, mask):
         ],
         "outputs": [
             {
-                "name": "attention_bias",
-                "shape": [1, seq, 8, 64],
+                "name": "permute_dims_3",
+                "shape": [1, 8, seq, 64],
                 "dtype": "float32",
-                "layout": "ABCD",
+                "layout": "ACBD",
             }
         ],
-        "nodes": {"total": 5, "input": 4, "msc.attention": 1},
+        "nodes": {"total": 6, "input": 4, "msc.attention": 1, "permute_dims": 1},
     }
     if dynamic:
         expected2["prims"] = {"total": 1, "shape": 1}