From 5c96c6efba255094871aabf1872e14ce706f9917 Mon Sep 17 00:00:00 2001 From: Joren Dumoulin Date: Mon, 20 Jan 2025 17:32:38 +0100 Subject: [PATCH] dialects: (linalg) quantized matmul op generic prints 5 affinemaps (#3772) The two zero points given also need affinemaps (even though they are integers) --- tests/filecheck/dialects/linalg/linalg_ops.mlir | 2 +- xdsl/dialects/linalg.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/filecheck/dialects/linalg/linalg_ops.mlir b/tests/filecheck/dialects/linalg/linalg_ops.mlir index dd9a9d3485..95f24552bf 100644 --- 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}> : (i32, i32) -> i32 // CHECK-GENERIC-NEXT: %46 = "arith.addi"(%40, %45) <{overflowFlags = #arith.overflow}> : (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 index 3e06938a92..065280c3be 100644 --- 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))), ] )