dialects: (linalg) quantized matmul op generic prints 5 affinemaps (#…

…3772) The two zero points given also need affinemaps (even though they are integers)
xdslproject · Jan 20, 2025 · 5c96c6e · 5c96c6e
1 parent 40aaea2
commit 5c96c6e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tests/filecheck/dialects/linalg/linalg_ops.mlir b/tests/filecheck/dialects/linalg/linalg_ops.mlir
@@ -150,4 +150,4 @@ linalg.quantized_matmul ins(%5, %6, %7, %8 : tensor<64x9216xi8>, tensor<9216x409
 // CHECK-GENERIC-NEXT:         %45 = "arith.muli"(%42, %44) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
 // CHECK-GENERIC-NEXT:         %46 = "arith.addi"(%40, %45) <{overflowFlags = #arith.overflow<none>}> : (i32, i32) -> i32
 //  CHECK-GENERIC-NEXT:        "linalg.yield"(%46) : (i32) -> ()
-// CHECK-GENERIC-NEXT:       }) {linalg.memoized_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>]} : (tensor<64x9216xi8>, tensor<9216x4096xi8>, i32, i32, tensor<64x4096xi32>) -> tensor<64x4096xi32>
+// CHECK-GENERIC-NEXT:       }) {linalg.memoized_indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>, affine_map<(d0, d1, d2) -> (d2, d1)>, affine_map<(d0, d1, d2) -> ()>, affine_map<(d0, d1, d2) -> ()>, affine_map<(d0, d1, d2) -> (d0, d1)>]} : (tensor<64x9216xi8>, tensor<9216x4096xi8>, i32, i32, tensor<64x4096xi32>) -> tensor<64x4096xi32>
diff --git a/xdsl/dialects/linalg.py b/xdsl/dialects/linalg.py
@@ -935,6 +935,8 @@ def hidden_region(args: tuple[BlockArgument, ...]) -> None:
                 [
                     AffineMapAttr(AffineMap.from_callable(lambda i, _, k: (i, k))),
                     AffineMapAttr(AffineMap.from_callable(lambda _, j, k: (k, j))),
+                    AffineMapAttr(AffineMap(3, 0, ())),
+                    AffineMapAttr(AffineMap(3, 0, ())),
                     AffineMapAttr(AffineMap.from_callable(lambda i, j, _: (i, j))),
                 ]
             )