[MLIR][MemRef] Normalize memref.alloc ops with non trivial layout map

Author: arnab-polymage

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -118,9 +118,11 @@ struct SemiAffineExprFlattener : public AffineExprFlattener {
       // with a positive value." (see AffineExprKind in AffineExpr.h). If this
       // assumption does not hold constraints (added above) are a contradiction.
 
+      return success();
+    } else if (localExpr.getKind() == AffineExprKind::Mul) {
+      (void)localVarCst.appendVar(VarKind::Local);
       return success();
     }
-
     // TODO: Support other semi-affine expressions.
     return failure();
   }
@@ -163,7 +165,6 @@ getFlattenedAffineExprs(ArrayRef<AffineExpr> exprs, unsigned numDims,
 
     return success();
   };
-
   if (addConservativeSemiAffineBounds) {
     SemiAffineExprFlattener flattener(numDims, numSymbols);
     return flattenExprs(flattener);
@@ -229,7 +230,8 @@ LogicalResult FlatLinearConstraints::composeMatchingMap(AffineMap other) {
   assert(other.getNumSymbols() == getNumSymbolVars() && "symbol mismatch");
 
   std::vector<SmallVector<int64_t, 8>> flatExprs;
-  if (failed(flattenAlignedMapAndMergeLocals(other, &flatExprs)))
+  if (failed(flattenAlignedMapAndMergeLocals(
+          other, &flatExprs, /*addConservativeSemiAffineBounds=*/true)))
     return failure();
   assert(flatExprs.size() == other.getNumResults());
 
@@ -796,8 +798,6 @@ LogicalResult FlatLinearConstraints::flattenAlignedMapAndMergeLocals(
                << "composition unimplemented for semi-affine maps\n");
     return failure();
   }
-
-  // Add localCst information.
   if (localCst.getNumLocalVars() > 0) {
     unsigned numLocalVars = getNumLocalVars();
     // Insert local dims of localCst at the beginning.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

@@ -1786,7 +1786,6 @@ static void createNewDynamicSizes(MemRefType oldMemRefType,
   }
 }
 
-// TODO: Currently works for static memrefs with a single layout map.
 template <typename AllocLikeOp>
 LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
   MemRefType memrefType = allocOp->getType();
@@ -1799,7 +1798,6 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
     // Either memrefType already had an identity map or the map couldn't be
     // transformed to an identity map.
     return failure();
-
   Value oldMemRef = allocOp->getResult();
 
   SmallVector<Value, 4> symbolOperands(allocOp->getSymbolOperands());
@@ -1819,8 +1817,40 @@ LogicalResult mlir::affine::normalizeMemRef(AllocLikeOp *allocOp) {
         b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType, newDynamicSizes,
                               allocOp->getAlignmentAttr());
   } else {
-    newAlloc = b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
-                                     allocOp->getAlignmentAttr());
+    mlir::ValueRange dynamicSizes = allocOp->getDynamicSizes();
+    mlir::ValueRange symbolOperands = allocOp->getSymbolOperands();
+    ArrayRef<int64_t> newShape = newMemRefType.getShape();
+    ArrayRef<int64_t> oldShape = memrefType.getShape();
+    SmallVector<Value> mapOperands(oldShape.size() + symbolOperands.size());
+    SmallVector<Value> dimensionOperands;
+    unsigned dimId = 0, symId = 0;
+    // Collect all the map operands of `allocOp` (both dynamic sizes and symbol
+    // operands), which will help us to compute the dynamic sizes of the new
+    // alloc op we are going to create.
+    for (unsigned i = 0, e = oldShape.size(); i < e; i++) {
+      if (oldShape[i] == ShapedType::kDynamic)
+        mapOperands[i] = dynamicSizes[dimId++];
+      else
+        mapOperands[i] =
+            b.create<arith::ConstantIndexOp>(allocOp->getLoc(), oldShape[i]);
+    }
+    for (unsigned i = oldShape.size(), e = mapOperands.size(); i < e; i++)
+      mapOperands[i] = symbolOperands[symId++];
+    // Compute the dynamic sizes operands for the new alloc op. If `newShape` is
+    // dynamic along a dimension, compute its shape using the layout map and
+    // dynamic sizes and symbol operands of the old `allocOp`.
+    for (unsigned i = 0, e = newShape.size(); i < e; i++) {
+      if (newShape[i] != ShapedType::kDynamic)
+        continue;
+      AffineExpr resExpr = layoutMap.getResult(i);
+      auto resMap = AffineMap::get(layoutMap.getNumDims(),
+                                   layoutMap.getNumSymbols(), resExpr);
+      dimensionOperands.push_back(
+          b.create<AffineApplyOp>(allocOp->getLoc(), resMap, mapOperands));
+    }
+    newAlloc =
+        b.create<AllocLikeOp>(allocOp->getLoc(), newMemRefType,
+                              dimensionOperands, allocOp->getAlignmentAttr());
   }
   // Replace all uses of the old memref.
   if (failed(replaceAllMemRefUsesWith(oldMemRef, /*newMemRef=*/newAlloc,
@@ -1868,11 +1898,8 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
 
   // Normalize only static memrefs and dynamic memrefs with a tiled-layout map
   // for now.
-  // TODO: Normalize the other types of dynamic memrefs.
   SmallVector<std::tuple<AffineExpr, unsigned, unsigned>> tileSizePos;
   (void)getTileSizePos(layoutMap, tileSizePos);
-  if (memrefType.getNumDynamicDims() > 0 && tileSizePos.empty())
-    return memrefType;
 
   // We have a single map that is not an identity map. Create a new memref
   // with the right shape and an identity layout map.
@@ -1894,7 +1921,6 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
   unsigned newRank = layoutMap.getNumResults();
   if (failed(fac.composeMatchingMap(layoutMap)))
     return memrefType;
-  // TODO: Handle semi-affine maps.
   // Project out the old data dimensions.
   fac.projectOut(newRank, fac.getNumVars() - newRank - fac.getNumLocalVars());
   SmallVector<int64_t, 4> newShape(newRank);
@@ -1910,14 +1936,14 @@ MemRefType mlir::affine::normalizeMemRefType(MemRefType memrefType) {
     // For a static memref and an affine map with no symbols, this is
     // always bounded. However, when we have symbols, we may not be able to
     // obtain a constant upper bound. Also, mapping to a negative space is
-    // invalid for normalization.
-    if (!ubConst.has_value() || *ubConst < 0) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "can't normalize map due to unknown/invalid upper bound");
+    // invalid for normalization. If dimension of new memrefType is dynamic,
+    // the value is `ShapedType::kDynamic`.
+    if (!ubConst.has_value())
+      newShape[d] = ShapedType::kDynamic;
+    else if (*ubConst >= 0)
+      newShape[d] = *ubConst + 1;
+    else
       return memrefType;
-    }
-    // If dimension of new memrefType is dynamic, the value is -1.
-    newShape[d] = *ubConst + 1;
   }
 
   // Create the new memref type after trivializing the old layout map.

mlir/lib/Dialect/MemRef/Transforms/NormalizeMemRefs.cpp

@@ -445,8 +445,10 @@ void NormalizeMemRefs::normalizeFuncOpMemRefs(func::FuncOp funcOp,
           if (oldMemRefType == newMemRefType)
             continue;
           // TODO: Assume single layout map. Multiple maps not supported.
+          // TODO: Semi-affine layout not supported.
           AffineMap layoutMap = oldMemRefType.getLayout().getAffineMap();
-          if (failed(replaceAllMemRefUsesWith(oldMemRef,
+          if (!layoutMap.getResult(0).isPureAffine() ||
+              failed(replaceAllMemRefUsesWith(oldMemRef,
                                               /*newMemRef=*/newMemRef,
                                               /*extraIndices=*/{},
                                               /*indexRemap=*/layoutMap,

mlir/test/Dialect/Affine/memref-bound-check.mlir

@@ -124,13 +124,14 @@ func.func @mod_floordiv_nested() {
   return
 }
 
-// CHECK-LABEL: func @test_semi_affine_bailout
-func.func @test_semi_affine_bailout(%N : index) {
+// CHECK-LABEL: func @test_semi_affine_access
+func.func @test_semi_affine_access(%N : index) {
   %B = memref.alloc() : memref<10 x i32>
   affine.for %i = 0 to 10 {
     %idx = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%i)[%N]
     %y = affine.load %B[%idx] : memref<10 x i32>
-    // expected-error@-1 {{getMemRefRegion: compose affine map failed}}
+    // expected-error@-1 {{'affine.load' op memref out of upper bound access along dimension #1}}
+    // expected-error@-2 {{'affine.load' op memref out of lower bound access along dimension #1}}
   }
   return
 }