wip

Author: masahi

tests/python/relax/test_codegen_cutlass.py

@@ -1961,8 +1961,8 @@ def rewrite_attention(f):
     K_BNSH_T = is_op("relax.permute_dims")(K_BNSH)
 
     matmul1 = is_op("relax.matmul")(Q_BNSH, K_BNSH_T)
-    multiply = is_op("relax.multiply")(matmul1, is_const())
-    softmax = is_op("relax.nn.softmax")(multiply)
+    multiply = is_op("relax.divide")(matmul1, is_const())
+    softmax = is_op("relax.astype")(is_op("relax.nn.softmax")(is_op("relax.astype")(multiply)))
     matmul2 = is_op("relax.matmul")(softmax, V_BNSH)
     pattern = is_op("relax.permute_dims")(is_op("relax.reshape")(matmul2, is_shape([4, 32, 16, 8])))
 
@@ -1996,9 +1996,9 @@ def main(
 
                 lv6 = R.permute_dims(lv3, axes=[0, 2, 1])
                 lv7 = R.matmul(lv1, lv6, out_dtype="float16")
-                lv3_1 = R.const(0.5, "float16")
-                lv8 = R.multiply(lv7, lv3_1)
-                lv11 = R.nn.softmax(lv8, axis=2)
+                lv3_1 = R.const(2.0, "float16")
+                lv8 = R.divide(lv7, lv3_1)
+                lv11 = R.astype(R.nn.softmax(R.astype(lv8, "float32"), axis=2), "float16")
                 lv12 = R.matmul(lv11, lv5, out_dtype="float16")
                 lv13 = R.reshape(lv12, R.shape([4, 32, 16, 8]))
                 lv6_1 = R.permute_dims(lv13, axes=[0, 2, 1, 3])