[mlir][vector] Fix a `target-rank=0` unrolling (#73365)
Fixes https://github.com/llvm/llvm-project/issues/64269. With this patch, calling `mlir-opt "-convert-vector-to-scf=full-unroll target-rank=0"` on ```mlir func.func @main(%vec : vector<2xi32>) { %alloc = memref.alloc() : memref<4xi32> %c0 = arith.constant 0 : index vector.transfer_write %vec, %alloc[%c0] : vector<2xi32>, memref<4xi32> return } ``` will result in ```mlir module { func.func @main(%arg0: vector<2xi32>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %alloc = memref.alloc() : memref<4xi32> %0 = vector.extract %arg0[0] : i32 from vector<2xi32> %1 = vector.broadcast %0 : i32 to vector<i32> vector.transfer_write %1, %alloc[%c0] : vector<i32>, memref<4xi32> %2 = vector.extract %arg0[1] : i32 from vector<2xi32> %3 = vector.broadcast %2 : i32 to vector<i32> vector.transfer_write %3, %alloc[%c1] : vector<i32>, memref<4xi32> return } } ``` I've also tried to proactively find other `target-rank=0` bugs, but couldn't find any. `options.targetRank` is only used 8 times throughout the `mlir` folder, all inside `VectorToSCF.cpp`. None of the other uses look like they could cause a crash. I've also tried ```mlir func.func @main(%vec : vector<2xi32>) -> vector<2xi32> { %alloc = memref.alloc() : memref<4xindex> %c0 = arith.constant 0 : index %out = vector.transfer_read %alloc[%c0], %c0 : memref<4xindex>, vector<2xi32> return %out : vector<2xi32> } ``` with `"--convert-vector-to-scf=full-unroll target-rank=0"` and that also didn't crash. (Maybe obvious. I have to admit that I'm not very familiar with these ops.)
Loading
Please sign in to comment