[Mlir-commits] [mlir] [mlir][vector] Deal with special patterns when emulating masked load/store (PR #75587)

[Hsiangkai Wang]

Hsiangkai wrote:

I use the following mlir func as input (without this PR)
```
func.func @vector_maskedstore_c2(%arg0: memref<4x5xf32>, %arg1: vector<4xf32>) {
  %c7 = arith.constant 7 : index
  %c6 = arith.constant 6 : index
  %c5 = arith.constant 5 : index
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  %0 = vector.constant_mask [2] : vector<4xi1>
  %1 = vector.extract %0[0] : i1 from vector<4xi1>
  scf.if %1 {
    %5 = vector.extract %arg1[0] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c4] : memref<4x5xf32>
  }
  %2 = vector.extract %0[1] : i1 from vector<4xi1>
  scf.if %2 {
    %5 = vector.extract %arg1[1] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c5] : memref<4x5xf32>
  }
  %3 = vector.extract %0[2] : i1 from vector<4xi1>
  scf.if %3 {
    %5 = vector.extract %arg1[2] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c6] : memref<4x5xf32>
  }
  %4 = vector.extract %0[3] : i1 from vector<4xi1>
  scf.if %4 {
    %5 = vector.extract %arg1[3] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c7] : memref<4x5xf32>
  }
  return
}
```
After applying `--canonicalize --cse`, there is no change at all.
The output is
```
func.func @vector_maskedstore_c2(%arg0: memref<4x5xf32>, %arg1: vector<4xf32>) {
  %c7 = arith.constant 7 : index
  %c6 = arith.constant 6 : index
  %c5 = arith.constant 5 : index
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  %0 = vector.constant_mask [2] : vector<4xi1>
  %1 = vector.extract %0[0] : i1 from vector<4xi1>
  scf.if %1 {
    %5 = vector.extract %arg1[0] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c4] : memref<4x5xf32>
  }
  %2 = vector.extract %0[1] : i1 from vector<4xi1>
  scf.if %2 {
    %5 = vector.extract %arg1[1] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c5] : memref<4x5xf32>
  }
  %3 = vector.extract %0[2] : i1 from vector<4xi1>
  scf.if %3 {
    %5 = vector.extract %arg1[2] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c6] : memref<4x5xf32>
  }
  %4 = vector.extract %0[3] : i1 from vector<4xi1>
  scf.if %4 {
    %5 = vector.extract %arg1[3] : f32 from vector<4xf32>
    memref.store %5, %arg0[%c0, %c7] : memref<4x5xf32>
  }
  return
}
```

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