Detecting reduction loops in affine dialect

Hi,
I need to detect affine loops representing reductions, i.e. that look similar to this:

affine.for %j = 0 to 64 {
  %x = affine.load %0[%i] : memref<64xf32>
  %y = affine.load %1[%i, %j] : memref<64x64xf32>
  %z = addf %x, %y : f32
  affine.store %z, %0[%i] : memref<64xf32>
}

Are there utilities in MLIR that do something like this? So far the closest thing I found is the isLoopParallel function, so maybe I have to implement something along the lines, but may be I missed something more relevant for this task?

Kindly pinging @bondhugula and @ftynse

There aren’t currently any utilities to detect reductions. However, once detected, there is support to represent such reductions via iter_args and return values on the affine.for as well as the lowering from there to scf.for and then to the standard dialect.

The simple case of detection sounds relatively straightforward, assuming no aliasing like the rest of affine transformations. Find an affine.store, follow use-def chains and check if the value stored transitively originates from an affine.load of the same memory location, and that there are no other accesses to that location.

Something like mem2reg on affine loops could also help exposing iter_args and going from there.

Thanks! Running mem2reg first would definitely make things easier (however, if I’m not mistaken, there isn’t such a pass for affine yet).