[mlir][linalg] fix linalg.batch_reduce_matmul auto cast (#102585)

zhiczhong · web-flow · commit 86347f6a477c · 2024-08-12T14:37:57.000+08:00
Fix the auto-cast of `linalg.batch_reduce_matmul` from `cast_to_T(A *
cast_to_T(B)) + C` to `cast_to_T(A) * cast_to_T(B) + C`
diff --git a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -592,8 +592,8 @@ def batch_reduce_matmul(
     """
     domain(D.b, D.m, D.n, D.k)
     implements(ContractionOpInterface)
-    C[D.m, D.n] += TypeFn.cast_signed(
-        U, A[D.b, D.m, D.k] * TypeFn.cast_signed(U, B[D.b, D.k, D.n])
+    C[D.m, D.n] += TypeFn.cast_signed(U, A[D.b, D.m, D.k]) * TypeFn.cast_signed(
+        U, B[D.b, D.k, D.n]
     )
 
 

Original file line number	Diff line number	Diff line change
`@@ -592,8 +592,8 @@ def batch_reduce_matmul(`
`592`	`592`	`"""`
`593`	`593`	`domain(D.b, D.m, D.n, D.k)`
`594`	`594`	`implements(ContractionOpInterface)`
`595`		`- C[D.m, D.n] += TypeFn.cast_signed(`
`596`		`- U, A[D.b, D.m, D.k] * TypeFn.cast_signed(U, B[D.b, D.k, D.n])`
	`595`	`+ C[D.m, D.n] += TypeFn.cast_signed(U, A[D.b, D.m, D.k]) * TypeFn.cast_signed(`
	`596`	`+ U, B[D.b, D.k, D.n]`
`597`	`597`	`)`
`598`	`598`
`599`	`599`