[mlir][linalg][nfc] Fix `linalg.matmul_transpose_a` def. #97690

JerryShih · 2024-07-04T08:15:46Z

The matmul_transpose_a input data format should be KxM * KxN.

github-actions · 2024-07-04T08:16:06Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be
notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write
permissions for the repository. In which case you can instead tag reviewers by
name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review
by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate
is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-07-04T08:16:35Z

@llvm/pr-subscribers-mlir-linalg

@llvm/pr-subscribers-mlir

Author: Jerry Shih (JerryShih)

Changes

The matmul_transpose_a input data format should be KxM * KxN.

Full diff: https://github.com/llvm/llvm-project/pull/97690.diff

2 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml (+1-1)
(modified) mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py (+2-2)

diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml b/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
index fad234a9dcae9..abb79278eddd4 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 43410aaa6af1b..59b3ba914eaab 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),
 ):

banach-space · 2024-07-05T08:01:27Z

The matmul_transpose_a input data format should be KxM * KxN.

That's already the case, isn't it? IIUC, this PR is merely changing the variable names in the YAML spec file. Making all definitions use consistent naming is definitely desirable, but the suggested change looks also inconsistent with other example. See my suggestion inline.

Also, this looks like an NFC change, so please update the name and the summary.

JerryShih · 2024-07-08T02:41:56Z

The matmul_transpose_a input data format should be KxM * KxN.

That's already the case, isn't it? IIUC, this PR is merely changing the variable names in the YAML spec file. Making all definitions use consistent naming is definitely desirable, but the suggested change looks also inconsistent with other example.

Here is the original matmul_transpose_a def:

def matmul_transpose_a(
    A=TensorDef(T1, S.K, S.N),
    B=TensorDef(T2, S.K, S.M),
    C=TensorDef(U, S.M, S.N, output=True),
    cast=TypeFnAttrDef(default=TypeFn.cast_signed),
):

It uses KxN for matrix A and KxM for matrix B.

And please check the batch_matmul_transpose_a and matmul def:

@linalg_structured_op
def matmul(
    A=TensorDef(T1, S.M, S.K),
    B=TensorDef(T2, S.K, S.N),
    C=TensorDef(U, S.M, S.N, output=True),
    cast=TypeFnAttrDef(default=TypeFn.cast_signed),
):

@linalg_structured_op
def batch_matmul_transpose_a(
    A=TensorDef(T1, Batch, S.K, S.M),
    B=TensorDef(T2, Batch, S.K, S.N),
    C=TensorDef(U, Batch, S.M, S.N, output=True),
):

The current matmul_transpose_a def is inconsistent with other ops def.
And the matmul_transpose_a's output is MxN.
From the current matmul_transpose_a def, the output shape should be:
A^T*B = NxK * KxM = NxM

See my suggestion inline.

I don't see the inlined suggestion at github? Do we have another review system for llvm project?

Also, this looks like an NFC change, so please update the name and the summary.

Yes, the input/output's affine_map doesn't used in llvm code-gen tools now. I will update the patch name/summary.

The `matmul_transpose_a` input data format should be `KxM * KxN` instead of current `KxN * KxM` format. It's a NFC fix.

banach-space · 2024-07-05T07:58:04Z

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)>


To keep the spec of linalg.matmul_transpose_a consistent with the other variants, this should be (s0, s2):

llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

Line 1412 in e975ff0

shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>

llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

Line 1094 in e975ff0

shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>

Now, on line 1329, you'd want to change:

shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>

to:

shape_map: affine_map<()[s0, s1, s2] -> (s1, s0)>

As in, IIUC:

s0 -> M dim

s1 -> K dim

s2 -> N dim

Why would K be in the middle?

I'm realising that shape_map != indexing_maps and that all my thinking was wrong. Sorry about the confusion.

What's the intended use of shape_map? Is that required specifically for https://mlir.llvm.org/docs/Dialects/Linalg/OpDSL/?

The YAML is generated from OpDSL and then processed by tablegen to generate op definitions. We don't go from YAML to OpDSL so nothing there is required for that. Not clear about tablegen.

Thanks!

So this PR doesn't change the indexing_maps - these seem to be actually correct. And shape_map seem to be dropped when converting to tablegen. I guess this is only fixing OpDSL?

@JerryShih How can we test/check/verify this change?

So this PR doesn't change the indexing_maps - these seem to be actually correct. And shape_map seem to be dropped when converting to tablegen. I guess this is only fixing OpDSL?

Yes, the input/output's affine_map doesn't used in mlir-linalg-ods-gen tool now. I just try to fix the weird declarion in core_named_ops.py.

@JerryShih How can we test/check/verify this change?

Sorry, I don't know how to test/verify this update. It's just try to fix the matrix dim name.

This can be tested by manually writing an instance of linalg.matmul_tranpose_a taking tensors of certain fixed size (says M=10,N=20,K=30) and making sure it passes the verifier. Without this change, it wouldn't because the dimension sizes wouldn't match.

@ftynse

This can be tested by manually writing an instance of linalg.matmul_tranpose_a taking tensors of certain fixed size (says M=10,N=20,K=30) and making sure it passes the verifier. Without this change, it wouldn't because the dimension sizes wouldn't match.

There is already a linalg.matmul_tranpose_a test inside.

llvm-project/mlir/test/Dialect/Linalg/named-ops.mlir

Lines 1193 to 1200 in ed4e75d

// CHECK-LABEL: func @matmul_transpose_a

// CHECK: linalg.matmul_transpose_a

// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<5x3xf32>, memref<5x7xf32>)

// CHECK-SAME: outs(%{{.+}} : memref<3x7xf32>)

func.func @matmul_transpose_a(%arg0: memref<5x3xf32>, %arg1: memref<5x7xf32>, %arg2: memref<3x7xf32>) {

linalg.matmul_transpose_a ins(%arg0, %arg1 : memref<5x3xf32>, memref<5x7xf32>) outs(%arg2: memref<3x7xf32>)

return

}

And this pr is "nfc" patch. It doesn't change the correctness of linalg.matmul_tranpose_a op.

Here is the affine_map used in yaml op generator.

llvm-project/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

Line 834 in ed4e75d

// Generating the getIndexingMaps() method.

It looks like only uses the affine_map in indexing_maps

llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

Lines 1099 to 1103 in 5fc7342

indexing_maps: !LinalgIndexingMapsConfig

static_indexing_maps:

- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>

- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>

- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>

Use linalg.matmul as the example:

vim ./llvm-project/tools/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yamlgen.cpp.inc .... ArrayAttr MatmulOp::getIndexingMaps() { static const char memoizeAttr[] = "linalg.memoized_indexing_maps"; ArrayAttr cached = getOperation()->getAttrOfType<ArrayAttr>(memoizeAttr); if (cached) return cached; MLIRContext *context = getContext(); auto symbolBindings = getSymbolBindings(*this); SmallVector<AffineMap> maps; maps.push_back(llvm::cast<AffineMapAttr>(mlir::parseAttribute("affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>", context)).getValue()); maps.back() = simplifyAffineMap(maps.back().replaceDimsAndSymbols({}, symbolBindings, 3, 0)); maps.push_back(llvm::cast<AffineMapAttr>(mlir::parseAttribute("affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>", context)).getValue()); maps.back() = simplifyAffineMap(maps.back().replaceDimsAndSymbols({}, symbolBindings, 3, 0)); maps.push_back(llvm::cast<AffineMapAttr>(mlir::parseAttribute("affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>", context)).getValue()); maps.back() = simplifyAffineMap(maps.back().replaceDimsAndSymbols({}, symbolBindings, 3, 0)); cached = Builder(context).getAffineMapArrayAttr(maps); getOperation()->setAttr(memoizeAttr, cached); return cached; }

I don't see the op def which is related the affine_map in args parts.

This pr only update the args parts which is not used in yaml generator.

llvm-project/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

Lines 1079 to 1098 in 5fc7342

args:

- !LinalgOperandDefConfig

name: A

kind: input_tensor

type_var: T1

shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>

- !LinalgOperandDefConfig

name: B

kind: input_tensor

type_var: T2

shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>

- !LinalgOperandDefConfig

name: C

kind: output_tensor

type_var: U

shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>

- !LinalgOperandDefConfig

name: cast

kind: type_fn_attr

default_fn: cast_signed

This pr only update the args parts which is not used in yaml generator.

Where is it used then? We need to find that place and add a test there. If there is no such place, we should drop this altogether.

Note, we (Intel Labs) are working on a rewrite of matmul in table-gen, so that these variants are not relevant any more. They should continue here in case there's usage downstream, but hopefully they'll become unused soon, so we can remove them.

The new matmul will have the proposed optional affine map to encode transpose and broadcast on the matmul operations themselves. Anyone using them as before will not be impacted, anyone using *_matmul_transposed_* can start using the new *_matmul { #transpose_map } variants.

@shahidact

banach-space · 2024-07-08T08:22:32Z

I don't see the inlined suggestion at github? Do we have another review system for llvm project?

I'm really sorry, I forgot to press the "submit" button :( I've just sent my comment. Looks like we are on the same page?

ftynse

Looks reasonable, could we see the result of converting this to linalg.generic to double-check? The fact that the K dimension is leading in shape_map is rather confusing, though I do think it's just the order in which dimensions appear rather than the order of loop iterators.

ftynse · 2024-07-08T20:18:49Z

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)>


Why would K be in the middle?

JerryShih requested review from ftynse, makslevental, stellaraccident, dcaballe, nicolasvasilache and rengolin as code owners July 4, 2024 08:15

llvmbot added mlir:linalg mlir:python mlir labels Jul 4, 2024

[mlir][linalg][nfc] Fix linalg.matmul_transpose_a def.

5fc7342

The `matmul_transpose_a` input data format should be `KxM * KxN` instead of current `KxN * KxM` format. It's a NFC fix.

JerryShih force-pushed the dev/jerrys/matmul_transpose_a-def branch from 4c654ed to 5fc7342 Compare July 8, 2024 02:46

JerryShih changed the title ~~[mlir][linalg] Fix linalg.matmul_transpose_a def.~~ [mlir][linalg][nfc] Fix linalg.matmul_transpose_a def. Jul 8, 2024

banach-space reviewed Jul 8, 2024

View reviewed changes

ftynse reviewed Jul 8, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[mlir][linalg][nfc] Fix `linalg.matmul_transpose_a` def. #97690

[mlir][linalg][nfc] Fix `linalg.matmul_transpose_a` def. #97690

JerryShih commented Jul 4, 2024

github-actions bot commented Jul 4, 2024

llvmbot commented Jul 4, 2024 •

edited

Loading

banach-space commented Jul 5, 2024

JerryShih commented Jul 8, 2024

banach-space Jul 5, 2024

ftynse Jul 8, 2024

banach-space Jul 8, 2024

ftynse Jul 9, 2024

banach-space Jul 9, 2024

JerryShih Jul 9, 2024 •

edited

Loading

ftynse Jul 11, 2024

JerryShih Jul 26, 2024

ftynse Aug 13, 2024

rengolin Aug 13, 2024

banach-space commented Jul 8, 2024

ftynse left a comment

ftynse Jul 8, 2024

	// CHECK-LABEL: func @matmul_transpose_a
	// CHECK: linalg.matmul_transpose_a
	// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<5x3xf32>, memref<5x7xf32>)
	// CHECK-SAME: outs(%{{.+}} : memref<3x7xf32>)
	func.func @matmul_transpose_a(%arg0: memref<5x3xf32>, %arg1: memref<5x7xf32>, %arg2: memref<3x7xf32>) {
	linalg.matmul_transpose_a ins(%arg0, %arg1 : memref<5x3xf32>, memref<5x7xf32>) outs(%arg2: memref<3x7xf32>)
	return
	}

	indexing_maps: !LinalgIndexingMapsConfig
	static_indexing_maps:
	- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
	- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
	- affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>

	args:
	- !LinalgOperandDefConfig
	name: A
	kind: input_tensor
	type_var: T1
	shape_map: affine_map<()[s0, s1, s2] -> (s0, s1)>
	- !LinalgOperandDefConfig
	name: B
	kind: input_tensor
	type_var: T2
	shape_map: affine_map<()[s0, s1, s2] -> (s1, s2)>
	- !LinalgOperandDefConfig
	name: C
	kind: output_tensor
	type_var: U
	shape_map: affine_map<()[s0, s1, s2] -> (s0, s2)>
	- !LinalgOperandDefConfig
	name: cast
	kind: type_fn_attr
	default_fn: cast_signed

[mlir][linalg][nfc] Fix linalg.matmul_transpose_a def. #97690

Are you sure you want to change the base?

[mlir][linalg][nfc] Fix linalg.matmul_transpose_a def. #97690

Conversation

JerryShih commented Jul 4, 2024

github-actions bot commented Jul 4, 2024

llvmbot commented Jul 4, 2024 • edited Loading

banach-space commented Jul 5, 2024

JerryShih commented Jul 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

JerryShih Jul 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

banach-space commented Jul 8, 2024

ftynse left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

[mlir][linalg][nfc] Fix `linalg.matmul_transpose_a` def. #97690

[mlir][linalg][nfc] Fix `linalg.matmul_transpose_a` def. #97690

llvmbot commented Jul 4, 2024 •

edited

Loading

JerryShih Jul 9, 2024 •

edited

Loading