[Mlir-commits] [mlir] [mlir][Transforms] Dialect conversion: Simplify handling of dropped arguments (PR #97213)

[llvmlistbot at llvm.org]

MaheshRavishankar wrote:

I am  posting here (but I can make it an issue as well), cause I dont fully follow if this is a downstream error or an error in this PR.  To repro you can use this example

```
func.func @alloc_transfer_read_write_vector4_vector8(%arg0: memref<4096x4096xf32>, %x: index, %y: index) {
  %cst = arith.constant 0.000000e+00 : f32
  %0 = memref.alloc() : memref<128x32xf32, 3>
  %v = vector.transfer_read %arg0[%x, %y], %cst : memref<4096x4096xf32>, vector<4xf32>
  vector.transfer_write %v, %0[%x, %y] : vector<4xf32>, memref<128x32xf32, 3>
  %mat = vector.transfer_read %arg0[%x, %y], %cst : memref<4096x4096xf32>, vector<8xf32>
  vector.transfer_write %mat, %0[%x, %y] : vector<8xf32>, memref<128x32xf32, 3>
  memref.dealloc %0 : memref<128x32xf32, 3>
  return
}
```

and the following command from IREE

```
iree-opt --pass-pipeline="builtin.module(func.func(iree-spirv-vectorize-load-store))" repro.mlir
```
With this change, I am hitting an assertion [here](https://github.com/iree-org/iree/blob/488cfc9688032d61cd47c63ff795d7f3e4642a16/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVVectorizeLoadStore.cpp) cause the pass expects the element type to be vector, but it isnt. Debugging this and adding a break point + `--debug` flag I see this being printed before the assertion

```
mlir-asm-printer: 'func.func' failed to verify and will be printed in generic form
"func.func"() <{function_type = (memref<4096x1024xvector<4xf32>>, index, index) -> (), sym_name = "alloc_transfer_read_write_vector4_vector8"}> ({
^bb0(%arg0: memref<4096x1024xvector<4xf32>>, %arg1: index, %arg2: index):
  %0 = "builtin.unrealized_conversion_cast"(%arg0) : (memref<4096x1024xvector<4xf32>>) -> memref<4096x4096xf32>
  %1 = "arith.constant"() <{value = 0.000000e+00 : f32}> : () -> f32
  %2 = "memref.alloc"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<128x8xvector<4xf32>, 3>
  %3 = "memref.alloc"() <{operandSegmentSizes = array<i32: 0, 0>}> : () -> memref<128x32xf32, 3>
  %4 = "vector.transfer_read"(<<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, %1) <{in_bounds = [false], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d1)>}> : (memref<4096x4096xf32>, index, index, f32) -> vector<4xf32>
  "vector.transfer_write"(%4, %3, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>) <{in_bounds = [false], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0, d1) -> (d1)>}> : (vector<4xf32>, memref<128x32xf32, 3>, index, index) -> ()
  %5 = "vector.transfer_read"(<<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, %1) <{in_bounds = [false], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d1)>}> : (memref<4096x4096xf32>, index, index, f32) -> vector<8xf32>
  "vector.transfer_write"(%5, %3, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>) <{in_bounds = [false], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0, d1) -> (d1)>}> : (vector<8xf32>, memref<128x32xf32, 3>, index, index) -> ()
  "memref.dealloc"(%3) : (memref<128x32xf32, 3>) -> ()
  "func.return"() : () -> ()
}) : () -> ()


} -> SUCCESS
//===-------------------------------------------===//

//===-------------------------------------------===//
Legalizing operation : 'vector.transfer_read'(0x555555665f40) {
  %4 = "vector.transfer_read"(<<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, %1) <{in_bounds = [false], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d1)>}> : (memref<4096x4096xf32>, index, index, f32) -> vector<4xf32>

  * Fold {
  } -> FAILURE : unable to fold

  * Pattern : 'vector.transfer_read -> ()' {
Trying to match "mlir::iree_compiler::(anonymous namespace)::ProcessTransferRead"
```

Looking at the IR (and stepping through dialect conversion as much as I understood it, there might be couple of things happening
1) At [this](https://github.com/iree-org/iree/blob/488cfc9688032d61cd47c63ff795d7f3e4642a16/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVVectorizeLoadStore.cpp) point, the `adaptor.geSource()` is supposed to be of type `memref<...<vector<4xf32>>` but it is instead the original type. So somewhere the value tracking went off the tracks.
2) The ```vector.transfer_read"(<<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, <<UNKNOWN SSA VALUE>>, %1) ``` is also suspect. I think the `signatureConversion` that happens when this pattern runs on the `func.func` was supposed to instead have this 
```vector.transfer_read"(%0, %arg1, %arg2, %1) ``` which means something is probably suspect in the signature conversion here (https://github.com/iree-org/iree/blob/488cfc9688032d61cd47c63ff795d7f3e4642a16/compiler/src/iree/compiler/Codegen/SPIRV/SPIRVVectorizeLoadStore.cpp#L312) . I am debugging this, but posting here for some input.

cc @matthias-springer @antiagainst and @bjacob 


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