MLIR Affine Dialect Loop Fusion pass appears to not work

[rohany]

I’m playing around with the Affine dialect of MLIR, in particular the Loop Fusion pass. I haven’t been able to get it to work with mlir-opt on some simple examples, including the ones from the documentation page here: https://mlir.llvm.org/docs/Passes/#-affine-loop-fusion-fuse-affine-loop-nests.

At a high level, I’m doing the following:

• Copying some MLIR source with fusable loops into a file “testing.mlir”
• Running bin/mlir-opt testing.mlir —affine-loop-fusion —dump-pass-pipeline

Here is concrete input and output:

➜  build git:(main) ✗ cat ../testing.mlir                                                                                                                                                                                      
func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
  %0 = memref.alloc() : memref<10xf32>
  %1 = memref.alloc() : memref<10xf32>
  %cst = arith.constant 0.000000e+00 : f32
  affine.for %arg2 = 0 to 10 {
    affine.store %cst, %0[%arg2] : memref<10xf32>
    affine.store %cst, %1[%arg2] : memref<10xf32>
  }
  affine.for %arg2 = 0 to 10 {
    %2 = affine.load %0[%arg2] : memref<10xf32>
    %3 = arith.addf %2, %2 : f32
    affine.store %3, %arg0[%arg2] : memref<10xf32>
  }
  affine.for %arg2 = 0 to 10 {
    %2 = affine.load %1[%arg2] : memref<10xf32>
    %3 = arith.mulf %2, %2 : f32
    affine.store %3, %arg1[%arg2] : memref<10xf32>
  }
  return
}

func.func @sibling_fusion(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>,
                     %arg2: memref<10x10xf32>, %arg3: memref<10x10xf32>,
                     %arg4: memref<10x10xf32>) {
  affine.for %arg5 = 0 to 3 {
    affine.for %arg6 = 0 to 3 {
      %0 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
      %1 = affine.load %arg1[%arg5, %arg6] : memref<10x10xf32>
      %2 = arith.mulf %0, %1 : f32
      affine.store %2, %arg3[%arg5, %arg6] : memref<10x10xf32>
    }
  }
  affine.for %arg5 = 0 to 3 {
    affine.for %arg6 = 0 to 3 {
      %0 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
      %1 = affine.load %arg2[%arg5, %arg6] : memref<10x10xf32>
      %2 = arith.addf %0, %1 : f32
      affine.store %2, %arg4[%arg5, %arg6] : memref<10x10xf32>
    }
  }
  return
}

➜  build git:(main) ✗ bin/mlir-opt ../testing.mlir --affine-loop-fusion --dump-pass-pipeline
Pass Manager with 1 passes:
builtin.module(affine-loop-fusion{fusion-compute-tolerance=3.000000e-01 fusion-fast-mem-space=0 fusion-local-buf-threshold=0 fusion-maximal=false mode=producer})

module {
  func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
    %alloc = memref.alloc() : memref<10xf32>
    %alloc_0 = memref.alloc() : memref<10xf32>
    %cst = arith.constant 0.000000e+00 : f32
    affine.for %arg2 = 0 to 10 {
      affine.store %cst, %alloc[%arg2] : memref<10xf32>
      affine.store %cst, %alloc_0[%arg2] : memref<10xf32>
    }
    affine.for %arg2 = 0 to 10 {
      %0 = affine.load %alloc[%arg2] : memref<10xf32>
      %1 = arith.addf %0, %0 : f32
      affine.store %1, %arg0[%arg2] : memref<10xf32>
    }
    affine.for %arg2 = 0 to 10 {
      %0 = affine.load %alloc_0[%arg2] : memref<10xf32>
      %1 = arith.mulf %0, %0 : f32
      affine.store %1, %arg1[%arg2] : memref<10xf32>
    }
    return
  }
  func.func @sibling_fusion(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>, %arg2: memref<10x10xf32>, %arg3: memref<10x10xf32>, %arg4: memref<10x10xf32>) {
    affine.for %arg5 = 0 to 3 {
      affine.for %arg6 = 0 to 3 {
        %0 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
        %1 = affine.load %arg1[%arg5, %arg6] : memref<10x10xf32>
        %2 = arith.mulf %0, %1 : f32
        affine.store %2, %arg3[%arg5, %arg6] : memref<10x10xf32>
      }
    }
    affine.for %arg5 = 0 to 3 {
      affine.for %arg6 = 0 to 3 {
        %0 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
        %1 = affine.load %arg2[%arg5, %arg6] : memref<10x10xf32>
        %2 = arith.addf %0, %1 : f32
        affine.store %2, %arg4[%arg5, %arg6] : memref<10x10xf32>
      }
    }
    return
  }
}

A related bug is that I seem to be unable to set the value of the argument mode to the loop fusion pass. Even if I do —affine-loop-fusion=“mode=greedy”, the output from the pass manager always reports that mode is producer.

[llvmbot]

@llvm/issue-subscribers-mlir-affine

[bondhugula]

--affine-loop-fusion --dump-pass-pipeline

This won't run it on the func op's loops. Instead, you need:

-pass-pipeline='builtin.module(func.func(affine-loop-fusion))'

You can see test/Transforms/loop-fusion*.mlir for examples of how the loop fusion pass is tested. Here's the output on your example.

module {
  func.func @producer_consumer_fusion(%arg0: memref<10xf32>, %arg1: memref<10xf32>) {
    %alloc = memref.alloc() : memref<1xf32>
    %alloc_0 = memref.alloc() : memref<1xf32>
    %cst = arith.constant 0.000000e+00 : f32
    affine.for %arg2 = 0 to 10 {
      affine.store %cst, %alloc[0] : memref<1xf32>
      affine.store %cst, %alloc_0[0] : memref<1xf32>
      %0 = affine.load %alloc_0[0] : memref<1xf32>
      %1 = arith.mulf %0, %0 : f32
      affine.store %1, %arg1[%arg2] : memref<10xf32>
      %2 = affine.load %alloc[0] : memref<1xf32>
      %3 = arith.addf %2, %2 : f32
      affine.store %3, %arg0[%arg2] : memref<10xf32>
    }
    return
  }
  func.func @sibling_fusion(%arg0: memref<10x10xf32>, %arg1: memref<10x10xf32>, %arg2: memref<10x10xf32>, %arg3: memref<10x10xf32>, %arg4: memref<10x10xf32>) {
    affine.for %arg5 = 0 to 3 {
      affine.for %arg6 = 0 to 3 {
        %0 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
        %1 = affine.load %arg1[%arg5, %arg6] : memref<10x10xf32>
        %2 = arith.mulf %0, %1 : f32
        affine.store %2, %arg3[%arg5, %arg6] : memref<10x10xf32>
        %3 = affine.load %arg0[%arg5, %arg6] : memref<10x10xf32>
        %4 = affine.load %arg2[%arg5, %arg6] : memref<10x10xf32>
        %5 = arith.addf %3, %4 : f32
        affine.store %5, %arg4[%arg5, %arg6] : memref<10x10xf32>
      }
    }
    return
  }
}

Reg. the modes, you'll see different output fused IR as expected when you change the mode -- so the right mode is being passed to the fusion pass. It's perhaps a dump-pass-pipeline issue/behavior that the default option is being printed (greedy is the default and it's what is being printed, not producer).

[joker-eph]

The way the affine pass is defined is a bit weird: it walks only the immediate regions nested under the current operation, why isn't it performing a walk instead?

[joker-eph]

Best to use https://discourse.llvm.org for questions, and also keep things on topic when an issue is started.

[bondhugula]

The way the affine pass is defined is a bit weird: it walks only the immediate regions nested under the current operation, why isn't it performing a walk instead?

The internal method itself can work on any Block; so one can call the affine fusion on, say, each of the Block's of an scf.while loop or inside of a specific affine.for or scf.for loop to fuse only the nests immediately nested at that depth, and for any/all depths. However, for the pass itself, doing a walk of all nested blocks and doing fusion by default at all the depths one after the other isn't what would be expected/the natural outcome in general -- at least as a default. For eg. consider this:

affine.for
  affine.for
      S0

affine.for %i
  affine.for %j
    S1
    affine.for
       S2
    affine.for 
      S3

Depending on the costs and customization, S0 could be pulled right before S1, before S2, right before S3 or anywhere in between (even right inside %i but outside and right before %j). Walking recursively and calling fusion for each block would yield a very different structure/fusion; that's possible to accomplish too easily and more expensive in general. It depends on what a user wants, and one would normally build/customize their pass downstream with the utilities available (LoopFusionUtils is separated out). We do have to expose a few more things though out of LoopFusion.cpp to make the customization easier (e.g. runOnBlock).