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

[Jerry Shih]

================
@@ -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)>
----------------
JerryShih wrote:

@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.
https://github.com/llvm/llvm-project/blob/ed4e75d5e5ada30c37c57df032378a77e6dd598e/mlir/test/Dialect/Linalg/named-ops.mlir#L1193-L1200

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.
https://github.com/llvm/llvm-project/blob/ed4e75d5e5ada30c37c57df032378a77e6dd598e/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp#L834

It looks like only uses the affine_map in `indexing_maps` 

https://github.com/llvm/llvm-project/blob/5fc7342a836cea6409719bdbdb5b69e8e4e7c570/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml#L1099-L1103

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.
https://github.com/llvm/llvm-project/blob/5fc7342a836cea6409719bdbdb5b69e8e4e7c570/mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml#L1079-L1098

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