Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)>
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

To keep the spec of linalg.matmul_transpose_a consistent with the other variants, this should be (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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why would K be in the middle?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

Copy link
Author

@JerryShih JerryShih Jul 9, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@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

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

- !LinalgOperandDefConfig
name: cast
kind: type_fn_attr
Expand Down
4 changes: 2 additions & 2 deletions mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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),
):
Expand Down
Loading