[all-commits] [llvm/llvm-project] db115b: [mlir][Linalg] Fix non-matmul linalg structured op...

[Kunwar Grover via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: db115ba3efee9c940539667842a1092d8d956850
      https://github.com/llvm/llvm-project/commit/db115ba3efee9c940539667842a1092d8d956850
  Author: Kunwar Grover <groverkss at gmail.com>
  Date:   2024-11-16 (Sat, 16 Nov 2024)

  Changed paths:
    M mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
    M mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    M mlir/test/Dialect/Linalg/rank-reduce-contraction-ops.mlir
    M mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

  Log Message:
  -----------
  [mlir][Linalg] Fix non-matmul linalg structured ops (#116412)

https://github.com/llvm/llvm-project/commit/3ad0148020ca91cc288bffd8ad36e25f7555a3bb
broke linalg structured ops other than MatmulOp.

The patch:

- Changes the printer to hide additional attributes, which weren't
hidden before: "indexing_maps".
- Changes the build of every linalg structured op to have an indexing
map for matmul.

These changes combined, hide the problem until you print the operation
in it's generic form.

Reproducer:

```mlir
func.func public @bug(%arg0 : tensor<5x10x20xf32>, %arg1 : tensor<5x20x40xf32>, %arg3 : tensor<5x10x40xf32>) -> tensor<5x10x40xf32> {
  %out = linalg.batch_matmul ins(%arg0, %arg1 : tensor<5x10x20xf32>, tensor<5x20x40xf32>)
      outs(%arg3 : tensor<5x10x40xf32>) -> tensor<5x10x40xf32>
  func.return %out : tensor<5x10x40xf32>
}
```

Prints fine, with `mlir-opt <file>`, but if you do `mlir-opt
--mlir-print-op-generic <file>`:

```
#map = affine_map<(d0, d1, d2) -> (d0, d2)>
#map1 = affine_map<(d0, d1, d2) -> (d2, d1)>
#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map4 = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>
#map5 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
"builtin.module"() ({
  "func.func"() <{function_type = (tensor<5x10x20xf32>, tensor<5x20x40xf32>, tensor<5x10x40xf32>) -> tensor<5x10x40xf32>, sym_name = "bug", sym_visibility = "public"}> ({
  ^bb0(%arg0: tensor<5x10x20xf32>, %arg1: tensor<5x20x40xf32>, %arg2: tensor<5x10x40xf32>):
    %0 = "linalg.batch_matmul"(%arg0, %arg1, %arg2) <{operandSegmentSizes = array<i32: 2, 1>}> ({
    ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):
      %1 = "arith.mulf"(%arg3, %arg4) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
      %2 = "arith.addf"(%arg5, %1) <{fastmath = #arith.fastmath<none>}> : (f32, f32) -> f32
      "linalg.yield"(%2) : (f32) -> ()
    }) {indexing_maps = [#map, #map1, #map2], linalg.memoized_indexing_maps = [#map3, #map4, #map5]} : (tensor<5x10x20xf32>, tensor<5x20x40xf32>, tensor<5x10x40xf32>) -> tensor<5x10x40xf32>
    "func.return"(%0) : (tensor<5x10x40xf32>) -> ()
  }) : () -> ()
}) : () -> ()
```

The batch_matmul operation's builder now always inserts a indexing_map
which is unrelated to the operation itself. This was caught when a
transformation from one LinalgStructuredOp to another, tried to pass
it's attributes to the other ops builder and there were multiple
indexing_map attributes in the result.

This patch fixes this by specializing the builders for MatmulOp with
indexing map information.



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications