How to efficiently vectorize linalg.generic reduction op?

MLIR
1

Hi all,

I have a question on vectorize reduction op.

Our current lowering pipeline for loop vectorization is -linalg-bufferize -convert-linalg-to-affine-loops -affine-super-vectorize.

This works well with element wise operation but not for reduction, for example in max operation:

linalg input:

   linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%alloc_2 : memref<128x
f32>) outs(%alloc_4 : memref<f32>) {                                                                                                              
    ^bb0(%in: f32, %out: f32):                                                                                                                    
      %25 = arith.maxf %in, %out : f32                                                                                                            
      linalg.yield %25 : f32                                                                                                                      
    }

After -linalg-bufferize -convert-linalg-to-affine-loops:

    affine.for %arg13 = 0 to 128 {
      %20 = affine.load %alloc_2[%arg13] : memref<128xf32>
      %21 = affine.load %alloc_4[] : memref<f32>
      %22 = arith.maxf %20, %21 : f32
      affine.store %22, %alloc_4[] : memref<f32>
    }

Aftert -affine-super-vectorize:

    "affine.for"() ({
      ^bb0(%arg13: index):
        %36 = "arith.constant"() <{value = 0.000000e+00 : f32}> : () -> f32
        %37 = "vector.transfer_read"(%21, %arg13, %36) {operand_segment_sizes = array<i32: 1, 1, 1, 0>, permutation_map = #map} : (memref<128xf32>
, index, f32) -> vector<128xf32>
        %38 = "arith.constant"() <{value = 0.000000e+00 : f32}> : () -> f32
        %39 = "vector.transfer_read"(%23, %38) {operand_segment_sizes = array<i32: 1, 0, 1, 0>, permutation_map = #map1} : (memref<f32>, f32) -> v
ector<128xf32>
        %40 = "arith.maxf"(%37, %39) <{fastmath = #arith.fastmath<none>}> : (vector<128xf32>, vector<128xf32>) -> vector<128xf32>
        "vector.transfer_write"(%40, %23) {operand_segment_sizes = array<i32: 1, 1, 0, 0>, permutation_map = #map1} : (vector<128xf32>, memref<f32
>) -> ()
        "affine.yield"() : () -> ()
      }) {lower_bound = #map1, step = 128 : index, upper_bound = #map2} : () -> ()
It generated broadcast/reduce read/write, vectorized max operation and didn’t get rid of the loop.

Ideally we want to go from linalg reduction to one single reduction op like: vector.reduction.

%1 = vector.reduction <maxf>, %0 : vector<128xf16> into f16

Any thoughts? Do you know any standard pass that can achieve above transform?

2

Have you tried Linalg vectorization. There are several examples here . They are using transform dialect for testing, but they are calling linalg::vectorize under the hood. I dont know if there is a pass that you can just call using mlir-opt

3

Thanks for the pointing. This works well for me.