[MLIR][Affine] Fix createPrivateMemRef in affine fusion

[dcaballe]

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

[bondhugula]

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?

[dcaballe]

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]

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.

[andydavis1]

Thanks Diego!

[dcaballe]

mlir/test/Transforms/loop-fusion.mlir

2658

Thanks! Will do before landing.

mlir/lib/Transforms/LoopFusion.cpp

static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   // consumer loop nests to reduce their live range. Currently they are added
   // at the beginning of the function, because loop nests can be reordered
   // during the fusion pass.
   Value newMemRef = top.create<AllocOp>(forOp.getLoc(), newMemRefType);
 
   // Build an AffineMap to remap access functions based on lower bound offsets.
   SmallVector<AffineExpr, 4> remapExprs;
   remapExprs.reserve(rank);
-  unsigned zeroOffsetCount = 0;
   for (unsigned i = 0; i < rank; i++) {
-    if (auto constExpr = offsets[i].dyn_cast<AffineConstantExpr>())
-      if (constExpr.getValue() == 0)
-        ++zeroOffsetCount;
     auto dimExpr = b.getAffineDimExpr(outerIVs.size() + i);
 
     auto remapExpr =
         simplifyAffineExpr(dimExpr - offsets[i], outerIVs.size() + rank, 0);
     remapExprs.push_back(remapExpr);
   }
-  auto indexRemap = zeroOffsetCount == rank
-                        ? AffineMap()
-                        : AffineMap::get(outerIVs.size() + rank, 0, remapExprs,
-                                         forOp.getContext());
+
+  auto indexRemap =
+      AffineMap::get(outerIVs.size() + rank, 0, remapExprs, forOp.getContext());
+
   // Replace all users of 'oldMemRef' with 'newMemRef'.
   LogicalResult res =
       replaceAllMemRefUsesWith(oldMemRef, newMemRef, {}, indexRemap,
                                /*extraOperands=*/outerIVs,
                                /*symbolOperands=*/{},
                                /*domInstFilter=*/&*forOp.getBody()->begin());
   assert(succeeded(res) &&
          "replaceAllMemrefUsesWith should always succeed here");

mlir/test/Transforms/loop-fusion.mlir

func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,
   return
 }
 
 // CHECK:  affine.for
 // CHECK:    addf
 // CHECK:  affine.for
 // CHECK:    mulf
 // CHECK:    subf
+
+// -----
+
+// MAXIMAL-LABEL: func @fuse_minor_affine_map
+func @fuse_minor_affine_map(%in: memref<128xf32>, %out: memref<20x512xf32>) {
+  %tmp = alloc() : memref<128xf32>
+
+  affine.for %arg4 = 0 to 128 {
+    %ld = affine.load %in[%arg4] : memref<128xf32>
+    affine.store %ld, %tmp[%arg4] : memref<128xf32>
+  }
+
+  affine.for %arg3 = 0 to 20 {
+    affine.for %arg4 = 0 to 512 {
+      %ld = affine.load %tmp[%arg4 mod 128] : memref<128xf32>
+      affine.store %ld, %out[%arg3, %arg4] : memref<20x512xf32>
+    }
+  }
+
+  return
+}
+
+// TODO: The size of the private memref is not properly computed in the presence
+// of the 'mod' operation. It should be memref<1xf32> instead of
+// memref<128xf32>: https://bugs.llvm.org/show_bug.cgi?id=46973
+// MAXIMAL:       alloc() : memref<128xf32>
+// MAXIMAL:       affine.for
+// MAXIMAL-NEXT:    affine.for
+// MAXIMAL-NOT:   affine.for