[all-commits] [llvm/llvm-project] 26722f: [MLIR] Fix incorrect memref::DimOp canonicalizatio...

[Sayan Saha via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 26722f5b61575fb0e58ff2933e7bea03353ff441
      https://github.com/llvm/llvm-project/commit/26722f5b61575fb0e58ff2933e7bea03353ff441
  Author: Sayan Saha <sayan.jubiee at gmail.com>
  Date:   2024-03-11 (Mon, 11 Mar 2024)

  Changed paths:
    M mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    M mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
    M mlir/test/Dialect/MemRef/canonicalize.mlir
    M mlir/test/Dialect/Tensor/canonicalize.mlir

  Log Message:
  -----------
  [MLIR] Fix incorrect memref::DimOp canonicalization, add tensor::DimOp canonicalization (#84225)

The current canonicalization of `memref.dim` operating on the result of
`memref.reshape` into `memref.load` is incorrect as it doesn't check
whether the `index` operand of `memref.dim` dominates the source
`memref.reshape` op. It always introduces `memref.load` right after
`memref.reshape` to ensure the `memref` is not mutated before the
`memref.load` call. As a result, the following error is observed:

```
$> mlir-opt --canonicalize input.mlir

func.func @reshape_dim(%arg0: memref<*xf32>, %arg1: memref<?xindex>, %arg2: index) -> index {
    %c4 = arith.constant 4 : index
    %reshape = memref.reshape %arg0(%arg1) : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
    %0 = arith.muli %arg2, %c4 : index
    %dim = memref.dim %reshape, %0 : memref<*xf32>
    return %dim : index
  }
```

results in:

```
dominator.mlir:22:12: error: operand #1 does not dominate this use
    %dim = memref.dim %reshape, %0 : memref<*xf32>
           ^
dominator.mlir:22:12: note: see current operation: %1 = "memref.load"(%arg1, %2) <{nontemporal = false}> : (memref<?xindex>, index) -> index
dominator.mlir:21:10: note: operand defined here (op in the same block)
    %0 = arith.muli %arg2, %c4 : index
```

Properly fixing this issue requires a dominator analysis which is
expensive to run within a canonicalization pattern. So, this patch fixes
the canonicalization pattern by being more strict/conservative about the
legality condition in which we perform this canonicalization.
The more general pattern is also added to `tensor.dim`. Since tensors are
immutable we don't need to worry about where to introduce the
`tensor.extract` call after canonicalization.



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