[Mlir-commits] [mlir] [mlir][linalg] Fix `linalg.matmul_transpose_a` def. (PR #97690)

[Jerry Shih]

https://github.com/JerryShih created https://github.com/llvm/llvm-project/pull/97690

The `matmul_transpose_a` input data format should be `KxM * KxN`.

>From 4c654ed5af456667505a42799bb420b825849e40 Mon Sep 17 00:00:00 2001
From: Jerry Shih <bignose1007 at gmail.com>
Date: Mon, 6 May 2024 01:55:50 -0700
Subject: [PATCH] [mlir][linalg] Fix `linalg.matmul_transpose_a` def.

The `matmul_transpose_a` input data format should be `KxM * KxN`.
---
 .../mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml      | 2 +-
 mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py  | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index fad234a9dcae9c..abb79278eddd4c 100644
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
@@ -1336,7 +1336,7 @@ structured_op: !LinalgStructuredOpConfig
     name: C
     kind: output_tensor
     type_var: U
-    shape_map: affine_map<()[s0, s1, s2] -> (s2, s1)>
+    shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
   - !LinalgOperandDefConfig
     name: cast
     kind: type_fn_attr
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
index 43410aaa6af1be..59b3ba914eaabe 100644
--- a/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
+++ b/mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
@@ -429,8 +429,8 @@ def quantized_matmul(
 
 @linalg_structured_op
 def matmul_transpose_a(
-    A=TensorDef(T1, S.K, S.N),
-    B=TensorDef(T2, S.K, S.M),
+    A=TensorDef(T1, S.K, S.M),
+    B=TensorDef(T2, S.K, S.N),
     C=TensorDef(U, S.M, S.N, output=True),
     cast=TypeFnAttrDef(default=TypeFn.cast_signed),
 ):