Incorrect memref size produced during privatization post-affine-fusion

[bane0101]

Hi,

This issue can almost be considered a clone of #115989, except that the example used here has affine scalar loads and stores, instead of their affine vector counterparts. I see that incorrect memref sizes are generated after fusing the loops below :

// fusion_test.mlir
module {
  func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
    %alloc = memref.alloc() : memref<1x2x4xf32>
    affine.for %i = 0 to 1 {
      %0 = affine.load %arg0[0, 0] : memref<2x4xf32>
      %1 = affine.load %arg0[0, 1] : memref<2x4xf32>
      %2 = affine.load %arg0[0, 2] : memref<2x4xf32>
      %3 = affine.load %arg0[0, 3] : memref<2x4xf32>
      %4 = affine.load %arg0[1, 0] : memref<2x4xf32>
      %5 = affine.load %arg0[1, 1] : memref<2x4xf32>
      %6 = affine.load %arg0[1, 2] : memref<2x4xf32>
      %7 = affine.load %arg0[1, 3] : memref<2x4xf32>

      affine.store %0, %alloc[0, 0, 0] : memref<1x2x4xf32>
      affine.store %1, %alloc[0, 0, 1] : memref<1x2x4xf32>
      affine.store %2, %alloc[0, 0, 2] : memref<1x2x4xf32>
      affine.store %3, %alloc[0, 0, 3] : memref<1x2x4xf32>
      affine.store %4, %alloc[0, 1, 0] : memref<1x2x4xf32>
      affine.store %5, %alloc[0, 1, 1] : memref<1x2x4xf32>
      affine.store %6, %alloc[0, 1, 2] : memref<1x2x4xf32>
      affine.store %7, %alloc[0, 1, 3] : memref<1x2x4xf32>
    }

    affine.for %i = 0 to 2 {
        affine.for %j = 0 to 4 {
            %8 = affine.load %alloc[0, %i, %j] : memref<1x2x4xf32>
            %9 = arith.negf %8 : f32
            affine.store %9, %arg1[0, %i, %j] : memref<1x2x4xf32>
        }
    }
    return
  }
}

$> mlir-opt --affine-loop-fusion fusion_test.mlir
results in the following output:

module {
  func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
    %alloc = memref.alloc() : memref<1x1x1xf32>
    affine.for %arg2 = 0 to 2 {
      affine.for %arg3 = 0 to 4 {
        %0 = affine.load %arg0[0, 0] : memref<2x4xf32>
        %1 = affine.load %arg0[0, 1] : memref<2x4xf32>
        %2 = affine.load %arg0[0, 2] : memref<2x4xf32>
        %3 = affine.load %arg0[0, 3] : memref<2x4xf32>
        %4 = affine.load %arg0[1, 0] : memref<2x4xf32>
        %5 = affine.load %arg0[1, 1] : memref<2x4xf32>
        %6 = affine.load %arg0[1, 2] : memref<2x4xf32>
        %7 = affine.load %arg0[1, 3] : memref<2x4xf32>
        affine.store %0, %alloc[0, 0, 0] : memref<1x1x1xf32>
        affine.store %1, %alloc[0, 0, 1] : memref<1x1x1xf32>
        affine.store %2, %alloc[0, 0, 2] : memref<1x1x1xf32>
        affine.store %3, %alloc[0, 0, 3] : memref<1x1x1xf32>
        affine.store %4, %alloc[0, 1, 0] : memref<1x1x1xf32>
        affine.store %5, %alloc[0, 1, 1] : memref<1x1x1xf32>
        affine.store %6, %alloc[0, 1, 2] : memref<1x1x1xf32>
        affine.store %7, %alloc[0, 1, 3] : memref<1x1x1xf32>
        %8 = affine.load %alloc[0, %arg2, %arg3] : memref<1x1x1xf32>
        %9 = arith.negf %8 : f32
        affine.store %9, %arg1[0, %arg2, %arg3] : memref<1x2x4xf32>
      }
    }
    return
  }
}

Here, the type of %alloc changed from memref<1x2x4xf32> to memref<1x1x1xf32> which is incorrect. The memref privatization code needs fixing.

[bane0101]

That said, I tried to explain to myself what could be a use case where such a privatization is meaningful and legal, and modified the test by rerolling the first loop :

module {
  func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
    %alloc = memref.alloc() : memref<1x2x4xf32>
    affine.for %i = 0 to 2 {
        affine.for %j = 0 to 4 {
            %0 = affine.load %arg0[%i, %j] : memref<2x4xf32>
            affine.store %0, %alloc[0, %i, %j] : memref<1x2x4xf32>
        }
    }

    affine.for %i = 0 to 2 {
        affine.for %j = 0 to 4 {
            %8 = affine.load %alloc[0, %i, %j] : memref<1x2x4xf32>
            %9 = arith.negf %8 : f32
            affine.store %9, %arg1[0, %i, %j] : memref<1x2x4xf32>
        }
    }
    return
  }
}

and found that it made sense with the latest output:

#map = affine_map<(d0, d1) -> (0)>
#map1 = affine_map<(d0, d1) -> (d0)>
#map2 = affine_map<(d0, d1) -> (d1)>
module {
  func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
    %alloc = memref.alloc() : memref<1x1x1xf32>
    affine.for %arg2 = 0 to 2 {
      affine.for %arg3 = 0 to 4 {
        %0 = affine.load %arg0[%arg2, %arg3] : memref<2x4xf32>
        %1 = affine.apply #map(%arg2, %arg3)
        %2 = affine.apply #map1(%arg2, %arg3)
        %3 = affine.apply #map2(%arg2, %arg3)
        affine.store %0, %alloc[0, 0, 0] : memref<1x1x1xf32>
        %4 = affine.apply #map(%arg2, %arg3)
        %5 = affine.apply #map1(%arg2, %arg3)
        %6 = affine.apply #map2(%arg2, %arg3)
        %7 = affine.load %alloc[0, 0, 0] : memref<1x1x1xf32>
        %8 = arith.negf %7 : f32
        affine.store %8, %arg1[0, %arg2, %arg3] : memref<1x2x4xf32>
      }
    }
    return
  }
}

or in simplified form (pre-scalarize) :

module {
  func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
    %alloc = memref.alloc() : memref<1x1x1xf32>
    affine.for %arg2 = 0 to 2 {
      affine.for %arg3 = 0 to 4 {
        %0 = affine.load %arg0[%arg2, %arg3] : memref<2x4xf32>
        affine.store %0, %alloc[0, 0, 0] : memref<1x1x1xf32>
        %7 = affine.load %alloc[0, 0, 0] : memref<1x1x1xf32>
        %8 = arith.negf %7 : f32
        affine.store %8, %arg1[0, %arg2, %arg3] : memref<1x2x4xf32>
      }
    }
    return
  }
}

Here, %alloc actually needs to be memref<1x1x1xf32>.
So it is likely an issue in analysing memref footprint by multiple loads and stores.
Problem somewhere in getConstantBoundingSizeAndShape?

[llvmbot]

@llvm/issue-subscribers-mlir-affine

Author: bane (bane0101)

Hi,

This issue can almost be considered a clone of #115989, except that the example used here has affine scalar loads and stores, instead of their affine vector counterparts. I see that incorrect memref sizes are generated after fusing the loops below :

// fusion_test.mlir
module {
  func.func @<!-- -->main(%arg0: memref&lt;2x4xf32&gt;, %arg1: memref&lt;1x2x4xf32&gt;) {
    %alloc = memref.alloc() : memref&lt;1x2x4xf32&gt;
    affine.for %i = 0 to 1 {
      %0 = affine.load %arg0[0, 0] : memref&lt;2x4xf32&gt;
      %1 = affine.load %arg0[0, 1] : memref&lt;2x4xf32&gt;
      %2 = affine.load %arg0[0, 2] : memref&lt;2x4xf32&gt;
      %3 = affine.load %arg0[0, 3] : memref&lt;2x4xf32&gt;
      %4 = affine.load %arg0[1, 0] : memref&lt;2x4xf32&gt;
      %5 = affine.load %arg0[1, 1] : memref&lt;2x4xf32&gt;
      %6 = affine.load %arg0[1, 2] : memref&lt;2x4xf32&gt;
      %7 = affine.load %arg0[1, 3] : memref&lt;2x4xf32&gt;

      affine.store %0, %alloc[0, 0, 0] : memref&lt;1x2x4xf32&gt;
      affine.store %1, %alloc[0, 0, 1] : memref&lt;1x2x4xf32&gt;
      affine.store %2, %alloc[0, 0, 2] : memref&lt;1x2x4xf32&gt;
      affine.store %3, %alloc[0, 0, 3] : memref&lt;1x2x4xf32&gt;
      affine.store %4, %alloc[0, 1, 0] : memref&lt;1x2x4xf32&gt;
      affine.store %5, %alloc[0, 1, 1] : memref&lt;1x2x4xf32&gt;
      affine.store %6, %alloc[0, 1, 2] : memref&lt;1x2x4xf32&gt;
      affine.store %7, %alloc[0, 1, 3] : memref&lt;1x2x4xf32&gt;
    }

    affine.for %i = 0 to 2 {
        affine.for %j = 0 to 4 {
            %8 = affine.load %alloc[0, %i, %j] : memref&lt;1x2x4xf32&gt;
            %9 = arith.negf %8 : f32
            affine.store %9, %arg1[0, %i, %j] : memref&lt;1x2x4xf32&gt;
        }
    }
    return
  }
}

$> mlir-opt --affine-loop-fusion fusion_test.mlir
results in the following output:

module {
  func.func @<!-- -->main(%arg0: memref&lt;2x4xf32&gt;, %arg1: memref&lt;1x2x4xf32&gt;) {
    %alloc = memref.alloc() : memref&lt;1x1x1xf32&gt;
    affine.for %arg2 = 0 to 2 {
      affine.for %arg3 = 0 to 4 {
        %0 = affine.load %arg0[0, 0] : memref&lt;2x4xf32&gt;
        %1 = affine.load %arg0[0, 1] : memref&lt;2x4xf32&gt;
        %2 = affine.load %arg0[0, 2] : memref&lt;2x4xf32&gt;
        %3 = affine.load %arg0[0, 3] : memref&lt;2x4xf32&gt;
        %4 = affine.load %arg0[1, 0] : memref&lt;2x4xf32&gt;
        %5 = affine.load %arg0[1, 1] : memref&lt;2x4xf32&gt;
        %6 = affine.load %arg0[1, 2] : memref&lt;2x4xf32&gt;
        %7 = affine.load %arg0[1, 3] : memref&lt;2x4xf32&gt;
        affine.store %0, %alloc[0, 0, 0] : memref&lt;1x1x1xf32&gt;
        affine.store %1, %alloc[0, 0, 1] : memref&lt;1x1x1xf32&gt;
        affine.store %2, %alloc[0, 0, 2] : memref&lt;1x1x1xf32&gt;
        affine.store %3, %alloc[0, 0, 3] : memref&lt;1x1x1xf32&gt;
        affine.store %4, %alloc[0, 1, 0] : memref&lt;1x1x1xf32&gt;
        affine.store %5, %alloc[0, 1, 1] : memref&lt;1x1x1xf32&gt;
        affine.store %6, %alloc[0, 1, 2] : memref&lt;1x1x1xf32&gt;
        affine.store %7, %alloc[0, 1, 3] : memref&lt;1x1x1xf32&gt;
        %8 = affine.load %alloc[0, %arg2, %arg3] : memref&lt;1x1x1xf32&gt;
        %9 = arith.negf %8 : f32
        affine.store %9, %arg1[0, %arg2, %arg3] : memref&lt;1x2x4xf32&gt;
      }
    }
    return
  }
}

Here, the type of %alloc changed from memref<1x2x4xf32> to memref<1x1x1xf32> which is incorrect. The memref privatization code needs fixing.

[shrikardongre]

Hi @bane0101 !
I am a first-time contributor and recently have started to get my hands on mlir, very excited about the opportunity to contribute to this project. This issue aligns with my interest in mlir and I would love to work on it. Could you please assign it to me? I am eager to learn and contribute effectively to the LLVM project.

[bane0101]

@shrikardongre sure, please feel free to take this up.

I think I do not have the right permission to assign it to anyone. @EugeneZelenko can you please assign if you have the permission?

[shrikardongre]

@bane0101,Apologies for the delay in addressing the issue; I was unable to work on it due to health reasons, but I am doing well now ,
I have read the full thread of these issues related to loop fusion.
If i am not wrong you are stating that in the first IR with multiple store and load operations happening the memref is not working properly ?

[shrikardongre]

module {
func.func @main(%arg0: memref<2x4xf32>, %arg1: memref<1x2x4xf32>) {
%alloc = memref.alloc() : memref<1x1x1xf32>
affine.for %arg2 = 0 to 2 {
affine.for %arg3 = 0 to 4 {
%0 = affine.load %arg0[%arg2, %arg3] : memref<2x4xf32>
affine.store %0, %alloc[0, 0, 0] : memref<1x1x1xf32>
%7 = affine.load %alloc[0, 0, 0] : memref<1x1x1xf32>
%8 = arith.negf %7 : f32
affine.store %8, %arg1[0, %arg2, %arg3] : memref<1x2x4xf32>
}
}
return
}
}

in this IR why was the %alloc was created at first place cant we just eliminate the need of it and directly use the value stored in %0 .
I might sound a little off topic here but while i was studying the loop fusion pass I came across the feasiblity of fusion and thats when it came to my mind .
Is it something like we will need to clear the memory allocated to %0 at every new iteration?
please correct me if i am wrong.

[shrikardongre]

@bane0101 also if
"So it is likely an issue in analysing memref footprint by multiple loads and stores."
this was the case, it shouldnt have worked after rerolling too as we are doing nothing different atleast thats what i think .

[bane0101]

@shrikardongre I hope you are doing well.

Yes, I agree the loads and stores in the first loop may not always necessarily be retained, and if preferred, can be store-forwarded using affine scalar replacement opt. But in this case lets assume it is needed, say, for some artificial reason like liveness of %alloc. The memref privatization post-fusion is still expected to emit legal output. The test case provided here is a minimal reproducer and is kind of a tiny part of a much larger IR.

The re-rolled version was created just for curious testing and observation. In this case, the privatization code has analyzed the liveness of %alloc correctly, and determined that not more than <1x1x1> is necessary for the computation. However, in the original case, the emitted code has illegal accesses in %alloc after reducing its size to <1x1x1>.

[bondhugula]

With recent changes, this bug is only reproducible under maximal fusion, i.e., with
bin/mlir-opt --affine-loop-fusion="maximal" ...
@shrikardongre - are you still working on this? If not, I can take this up. The issue here is that the private memref size/shape has been incorrectly determined.

[shrikardongre]

Hi, I tried to solve the issue but couldn't resolve it completely please feel free to take it up 😃.