Fix the legalization func of matmul to avoid zero index (#48)

This PR fixes a known issue of the TE legalization func of matmul, which
is used here in the generation of NT matmul. With this PR, there will be
no index 0 whenever the matmul on the particular dimension is not
broadcasting.

Author: MasterJH5574

web_llm/transform/dispatch_tir_operator.py

@@ -266,11 +266,11 @@ def matmul1_before(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, m
     for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), T.int64(1), T.int64(128), n):
         with T.block("matmul"):
             v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
-            T.reads(rxplaceholder[T.int64(0), v_i1, v_i2, v_k], rxplaceholder_1[T.int64(0), v_i1, v_k, v_i3])
+            T.reads(rxplaceholder[v_i0, v_i1, v_i2, v_k], rxplaceholder_1[v_i0, v_i1, v_k, v_i3])
             T.writes(matmul[v_i0, v_i1, v_i2, v_i3])
             with T.init():
                 matmul[v_i0, v_i1, v_i2, v_i3] = T.float32(0)
-            matmul[v_i0, v_i1, v_i2, v_i3] = matmul[v_i0, v_i1, v_i2, v_i3] + rxplaceholder[T.int64(0), v_i1, v_i2, v_k] * rxplaceholder_1[T.int64(0), v_i1, v_k, v_i3]
+            matmul[v_i0, v_i1, v_i2, v_i3] = matmul[v_i0, v_i1, v_i2, v_i3] + rxplaceholder[v_i0, v_i1, v_i2, v_k] * rxplaceholder_1[v_i0, v_i1, v_k, v_i3]
 
 
 @T.prim_func
@@ -448,11 +448,11 @@ def matmul5_before(var_rxplaceholder: T.handle, var_rxplaceholder_1: T.handle, v
     for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), n, T.int64(128), n):
         with T.block("matmul"):
             v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
-            T.reads(rxplaceholder[T.int64(0), v_i1, v_i2, v_k], rxplaceholder_1[T.int64(0), v_i1, v_k, v_i3])
+            T.reads(rxplaceholder[v_i0, v_i1, v_i2, v_k], rxplaceholder_1[v_i0, v_i1, v_k, v_i3])
             T.writes(matmul_1[v_i0, v_i1, v_i2, v_i3])
             with T.init():
                 matmul_1[v_i0, v_i1, v_i2, v_i3] = T.float32(0)
-            matmul_1[v_i0, v_i1, v_i2, v_i3] = matmul_1[v_i0, v_i1, v_i2, v_i3] + rxplaceholder[T.int64(0), v_i1, v_i2, v_k] * rxplaceholder_1[T.int64(0), v_i1, v_k, v_i3]
+            matmul_1[v_i0, v_i1, v_i2, v_i3] = matmul_1[v_i0, v_i1, v_i2, v_i3] + rxplaceholder[v_i0, v_i1, v_i2, v_k] * rxplaceholder_1[v_i0, v_i1, v_k, v_i3]
 
 
 @T.prim_func
@@ -1363,11 +1363,11 @@ def fused_NT_matmul1_divide_add_maximum_before(p_lv28: T.handle, p_lv29: T.handl
     for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), n, n, T.int64(128)):
         with T.block("NT_matmul"):
             v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
-            T.reads(lv28[T.int64(0), v_i1, v_i2, v_k], lv29[T.int64(0), v_i1, v_i3, v_k])
+            T.reads(lv28[v_i0, v_i1, v_i2, v_k], lv29[v_i0, v_i1, v_i3, v_k])
             T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3])
             with T.init():
                 var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = T.float32(0)
-            var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + lv28[T.int64(0), v_i1, v_i2, v_k] * lv29[T.int64(0), v_i1, v_i3, v_k]
+            var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + lv28[v_i0, v_i1, v_i2, v_k] * lv29[v_i0, v_i1, v_i3, v_k]
     for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), n, n):
         with T.block("T_divide"):
             v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])
@@ -1479,11 +1479,11 @@ def fused_NT_matmul6_divide1_add2_maximum1_before(lv2732: T.Buffer((T.int64(1),
     for i0, i1, i2, i3, k in T.grid(T.int64(1), T.int64(32), T.int64(1), n, T.int64(128)):
         with T.block("NT_matmul"):
             v_i0, v_i1, v_i2, v_i3, v_k = T.axis.remap("SSSSR", [i0, i1, i2, i3, k])
-            T.reads(lv2732[T.int64(0), v_i1, v_i2, v_k], lv2733[T.int64(0), v_i1, v_i3, v_k])
+            T.reads(lv2732[v_i0, v_i1, v_i2, v_k], lv2733[v_i0, v_i1, v_i3, v_k])
             T.writes(var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3])
             with T.init():
                 var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = T.float32(0)
-            var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + lv2732[T.int64(0), v_i1, v_i2, v_k] * lv2733[T.int64(0), v_i1, v_i3, v_k]
+            var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] = var_NT_matmul_intermediate[v_i0, v_i1, v_i2, v_i3] + lv2732[v_i0, v_i1, v_i2, v_k] * lv2733[v_i0, v_i1, v_i3, v_k]
     for ax0, ax1, ax2, ax3 in T.grid(T.int64(1), T.int64(32), T.int64(1), n):
         with T.block("T_divide"):
             v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0, ax1, ax2, ax3])

web_llm/transform/transpose_matmul.py

@@ -44,11 +44,7 @@ def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
                 b_shape.append(1)
 
             is_a_larger = len(a_shape) > len(b_shape)
-            offset = (
-                len(a_shape) - len(b_shape)
-                if is_a_larger
-                else len(b_shape) - len(a_shape)
-            )
+            offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape)
 
             a_relax = relax.Var("a", relax.TensorStructInfo(a.shape))
             bT_shape = list(b.shape)
@@ -70,15 +66,19 @@ def multiply_compute(idx_reduce):
                             a_indices.append(idx_spatial[i])
                         else:
                             b_indices.append(idx_spatial[i])
-                    for i in range(
-                        offset, len(output_shape) - (2 - a_prepended - b_appended)
-                    ):
+                    for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)):
                         a_dim = a_shape[i if is_a_larger else i - offset]
                         b_dim = b_shape[i if not is_a_larger else i - offset]
-                        a_dim_is_one = isinstance(a_dim, tir.IntImm) and a_dim == 1
-                        b_dim_is_one = isinstance(b_dim, tir.IntImm) and b_dim == 1
-                        a_indices.append(0 if a_dim_is_one else idx_spatial[i])
-                        b_indices.append(0 if b_dim_is_one else idx_spatial[i])
+                        dim_equal = a_dim == b_dim
+                        if not isinstance(dim_equal, tir.IntImm) or dim_equal == 0:
+                            a_dim_is_one = isinstance(a_dim, tir.IntImm) and a_dim == 1
+                            b_dim_is_one = isinstance(b_dim, tir.IntImm) and b_dim == 1
+                            a_indices.append(0 if a_dim_is_one else idx_spatial[i])
+                            b_indices.append(0 if b_dim_is_one else idx_spatial[i])
+                        else:
+                            a_indices.append(idx_spatial[i])
+                            b_indices.append(idx_spatial[i])
+
                     if not a_prepended:
                         a_indices.append(idx_spatial[-2 + b_appended])
                     a_indices.append(idx_reduce)
@@ -118,9 +118,7 @@ def multiply_compute(idx_reduce):
 
 @tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
 class FuseTransposeMatmul:
-    def transform_module(
-        self, mod: IRModule, ctx: tvm.transform.PassContext
-    ) -> IRModule:
+    def transform_module(self, mod: IRModule, ctx: tvm.transform.PassContext) -> IRModule:
         mod = relax.transform.FuseOpsByPattern(
             [("transpose_matmul_fuse", *TransposeMatmulCodeGenerator.pattern())]
         )(mod)