Page MenuHomePhabricator

[MLIR][Affine] Fix createPrivateMemRef in affine fusion
ClosedPublic

Authored by dcaballe on Aug 3 2020, 7:27 PM.

Details

Summary

I'm hitting an assert[1] in affine loop fusion when invoking replaceAllMemRefUsesWith from createPrivateMemRef.
It seems like we expect any affine.load/affine.store's map in the loop nest to always take all the loop nest indexes
as inputs. Unfortunately, that's not aways true. For example, the following affine.load map takes in only one index
out of two:

affine.for %arg3 = 0 to 20 {
  affine.for %arg4 = 0 to 512 {
    %ld = affine.load %tmp[%arg4 mod 128] : memref<128xf32>
    ...
  }
}

Concretely, we pass outerIVs (%arg3 for the prev example) as extra operands in replaceAllMemRefUsesWith [2].
These extra operands are blindly used as map indexes in the memref replacement[3] regardless of the number of indexes
actually needed in the map.

[1] https://github.com/llvm/llvm-project/blob/master/mlir/lib/Transforms/Utils/Utils.cpp#L171
[2] https://github.com/llvm/llvm-project/blob/master/mlir/lib/Transforms/LoopFusion.cpp#L942
[3] https://github.com/llvm/llvm-project/blob/master/mlir/include/mlir/Transforms/Utils.h#L39

The current fix always defines a remapping for the memref replacement (indexRemap) with the proper number of inputs,
including all the outerIVs, so that the number of inputs and the operands provided for the map don't mismatch.
Another option would be to filter the indexes that are not used out of the extra operands provided to
replaceAllMemRefUsesWith, but the current fix looks cleaner to me.

Please, let me know if you have any other idea for the fix.

Thanks!
Diego

Diff Detail

Event Timeline

dcaballe created this revision.Aug 3 2020, 7:27 PM
Herald added a project: Restricted Project. · View Herald TranscriptAug 3 2020, 7:27 PM
dcaballe requested review of this revision.Aug 3 2020, 7:27 PM
bondhugula requested changes to this revision.Aug 4 2020, 12:52 AM

Thanks for spotting this @dcaballe. That was clearly an incorrect assumption in the simplification being done for the all zero offset case which was thought to not require a remapping at all.

mlir/test/Transforms/loop-fusion.mlir
2658

Do you also want to FileCheck the private memref's alloc / size?

This revision now requires changes to proceed.Aug 4 2020, 12:52 AM
dcaballe added inline comments.Aug 4 2020, 7:50 AM
mlir/test/Transforms/loop-fusion.mlir
2658

The private memref alloc is not properly computed atm. I opened a bug for this since it's a separate problem: https://bugs.llvm.org/show_bug.cgi?id=46973. It seems like the memory region is not properly computed in the presence of 'mod' operations but I haven't had time to investigate this further. I can add a check for alloc() : memref<128xf32> for now, if you want.

bondhugula accepted this revision.Aug 4 2020, 9:48 AM
bondhugula added inline comments.
mlir/test/Transforms/loop-fusion.mlir
2658

Yes, sure - that would be good. You could add a comment pointing to the filed bug that this has to be fixed/improved.

This revision is now accepted and ready to land.Aug 4 2020, 9:48 AM
andydavis1 accepted this revision.Aug 4 2020, 10:19 AM

Thanks Diego!

dcaballe added inline comments.Aug 4 2020, 11:39 AM
mlir/test/Transforms/loop-fusion.mlir
2658

Thanks! Will do before landing.

This revision was automatically updated to reflect the committed changes.