[Mlir-commits] [mlir] [mlir][vectorize] Support affine.apply in SuperVectorize (PR #77968)

[Hsiangkai Wang]

Hsiangkai wrote:

> I barely remember this code but IIRC we were able to deal with affine.apply ops. @sergei-grechanik ?

In the current implementation, the vectorizer will fail. The debug log is

```
******************************************
******************************************
[early-vect] new pattern on parent op
func.func @vec_affine_apply(%arg0: memref<8x12x16xf32>, %arg1: memref<8x24x48xf32>) {
  affine.for %arg2 = 0 to 8 {
    affine.for %arg3 = 0 to 24 {
      affine.for %arg4 = 0 to 48 {
        %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
        %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
        %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xf32>
        affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
      }
    }
  }
  return
}
[early-vect]+++++ Vectorizing: affine.for %arg4 = 0 to 48 {
  %0 = affine.apply affine_map<(d0) -> (d0 mod 12)>(%arg3)
  %1 = affine.apply affine_map<(d0) -> (d0 mod 16)>(%arg4)
  %2 = affine.load %arg0[%arg2, %0, %1] : memref<8x12x16xf32>
  affine.store %2, %arg1[%arg2, %arg3, %arg4] : memref<8x24x48xf32>
}
[early-vect]+++++ commit vectorized op:
"affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index, upperBoundMap = affine_map<() -> (48)>}> ({
^bb0(%arg4: index):
  %0 = "affine.apply"(%arg3) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (index) -> index
  %1 = "affine.apply"(%arg4) <{map = affine_map<(d0) -> (d0 mod 16)>}> : (index) -> index
  %2 = "affine.load"(%arg0, %arg2, %0, %1) <{map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>}> : (memref<8x12x16xf32>, index, index, index) -> f32
  "affine.store"(%2, %arg1, %arg2, %arg3, %arg4) <{map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>}> : (f32, memref<8x24x48xf32>, index, index, index) -> ()
  "affine.yield"() : () -> ()
}) : () -> ()
into
"affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 8 : index, upperBoundMap = affine_map<() -> (48)>}> ({
^bb0(%arg4: index):
}) : () -> ()
[early-vect]+++++ Vectorizing: %0 = "affine.apply"(%arg3) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (index) -> index
[early-vect]+++++ vectorize operand: <block argument> of type 'index' at index: 0-> uniform: %0 = "vector.broadcast"(%arg3) : (index) -> vector<8xindex>
[early-vect]+++++ commit vectorized op:
%1 = "affine.apply"(%arg3) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (index) -> index
into
%1 = "affine.apply"(%0) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (vector<8xindex>) -> vector<8xindex>
[early-vect]+++++ Vectorizing: %2 = "affine.apply"(%arg4) <{map = affine_map<(d0) -> (d0 mod 16)>}> : (index) -> index
[early-vect]+++++ vectorize operand: <block argument> of type 'index' at index: 0-> unsupported block argument

[early-vect]+++++ an operand failed vectorize
[early-vect]+++++ failed vectorizing the operation: %2 = "affine.apply"(%arg4) <{map = affine_map<(d0) -> (d0 mod 16)>}> : (index) -> index
[early-vect]+++++ failed vectorization for: "affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 1 : index, upperBoundMap = affine_map<() -> (48)>}> ({
^bb0(%arg4: index):
  %1 = "affine.apply"(%arg3) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (index) -> index
  %2 = "affine.apply"(%arg4) <{map = affine_map<(d0) -> (d0 mod 16)>}> : (index) -> index
  %3 = "affine.load"(%arg0, %arg2, %1, %2) <{map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>}> : (memref<8x12x16xf32>, index, index, index) -> f32
  "affine.store"(%3, %arg1, %arg2, %arg3, %arg4) <{map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>}> : (f32, memref<8x24x48xf32>, index, index, index) -> ()
  "affine.yield"() : () -> ()
}) : () -> ()
[early-vect]+++++ erasing:
"affine.for"() <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array<i32: 0, 0, 0>, step = 8 : index, upperBoundMap = affine_map<() -> (48)>}> ({
^bb0(%arg4: index):
  %1 = "affine.apply"(%0) <{map = affine_map<(d0) -> (d0 mod 12)>}> : (vector<8xindex>) -> vector<8xindex>
}) : () -> ()
```

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