[Mlir-commits] [mlir] 3bfbc5d - [MLIR][Affine] Fix createPrivateMemRef in affine fusion

Diego Caballero llvmlistbot at llvm.org
Tue Aug 4 12:34:06 PDT 2020


Author: Diego Caballero
Date: 2020-08-04T12:17:48-07:00
New Revision: 3bfbc5df87cb0e736fe917e4a8d5166bc0b4ea79

URL: https://github.com/llvm/llvm-project/commit/3bfbc5df87cb0e736fe917e4a8d5166bc0b4ea79
DIFF: https://github.com/llvm/llvm-project/commit/3bfbc5df87cb0e736fe917e4a8d5166bc0b4ea79.diff

LOG: [MLIR][Affine] Fix createPrivateMemRef in affine fusion

Always define 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.

Reviewed By: bondhugula, andydavis1

Differential Revision: https://reviews.llvm.org/D85177

Added: 
    

Modified: 
    mlir/lib/Transforms/LoopFusion.cpp
    mlir/test/Transforms/loop-fusion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index dd7b7b83debd..ed79be02b816 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -921,21 +921,17 @@ static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
   // 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,

diff  --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index 7d91e780e9ed..b4eea34b41f1 100644
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -2634,3 +2634,32 @@ func @should_not_fuse_since_top_level_non_affine_users(%in0 : memref<32xf32>,
 // 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


        


More information about the Mlir-commits mailing list