[mlir] Fix invalid handling of AllocOp symbolOperands by SimplifyAllocConst.

[Hardcode84]

Summary

symbolOperands were completely ignored by SimplifyAllocConst. Also, I dindn't find any tests for this pattern. Also, slightly improved diagnostic message for verifyAllocLikeOp.

[Hardcode84]

newline

[Hardcode84]

Do not use getOperand

[Hardcode84]

Not sure who is the right person to review this

[Hardcode84]

someone?

[bondhugula]

Thanks for fixing this. Some minor comments.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

161

This is incorrectly named - should be dynamicSizes.

171

operand -> dynamicSize

mlir/test/Dialect/MemRef/canonicalize.mlir

226–230

There are tests for this pattern but they are in test/Transforms/canonicalize.mlir - @alloc_const_fold. Can you bring that test case here and keep either that or this?

234

Nit: ignore_symbols -> with_symbols ?

[Hardcode84]

rebase, move some tests from transforms/canonicalize, renames operand -> dynamicSize

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

if (static_cast<int64_t>(op.dynamicSizes().size()) !=
       memRefType.getNumDynamicDims())
     return op.emitOpError("dimension operand count does not equal memref "
                           "dynamic dimension count");
 
   unsigned numSymbols = 0;
   if (!memRefType.getAffineMaps().empty())
     numSymbols = memRefType.getAffineMaps().front().getNumSymbols();
   if (op.symbolOperands().size() != numSymbols)
-    return op.emitOpError(
-        "symbol operand count does not equal memref symbol count");
+    return op.emitOpError("symbol operand count does not equal memref symbol "
+                          "count: expected ")
+           << numSymbols << ", got " << op.symbolOperands().size();
 
   return success();
 }
 
 static LogicalResult verify(AllocOp op) { return verifyAllocLikeOp(op); }
 
 static LogicalResult verify(AllocaOp op) {
   // An alloca op needs to have an ancestor with an allocation scope trait.

 template <typename AllocLikeOp>
 struct SimplifyAllocConst : public OpRewritePattern<AllocLikeOp> {
   using OpRewritePattern<AllocLikeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AllocLikeOp alloc,
                                 PatternRewriter &rewriter) const override {
     // Check to see if any dimensions operands are constants.  If so, we can
     // substitute and drop them.
-    if (llvm::none_of(alloc.getOperands(), [](Value operand) {
+    if (llvm::none_of(alloc.dynamicSizes(), [](Value operand) {
           return matchPattern(operand, matchConstantIndex());
         }))
       return failure();
 
     auto memrefType = alloc.getType();
 
     // Ok, we have one or more constant operands.  Collect the non-constant ones
     // and keep track of the resultant memref type to build.
     SmallVector<int64_t, 4> newShapeConstants;
     newShapeConstants.reserve(memrefType.getRank());
     SmallVector<Value, 4> newOperands;
 
     unsigned dynamicDimPos = 0;
     for (unsigned dim = 0, e = memrefType.getRank(); dim < e; ++dim) {
       int64_t dimSize = memrefType.getDimSize(dim);
       // If this is already static dimension, keep it.
       if (dimSize != -1) {
         newShapeConstants.push_back(dimSize);
         continue;
       }
-      auto *defOp = alloc.getOperand(dynamicDimPos).getDefiningOp();
+      auto operand = alloc.dynamicSizes()[dynamicDimPos];
+      auto *defOp = operand.getDefiningOp();
       if (auto constantIndexOp = dyn_cast_or_null<ConstantIndexOp>(defOp)) {
         // Dynamic shape dimension will be folded.
         newShapeConstants.push_back(constantIndexOp.getValue());
       } else {
         // Dynamic shape dimension not folded; copy operand from old memref.
         newShapeConstants.push_back(-1);
-        newOperands.push_back(alloc.getOperand(dynamicDimPos));
+        newOperands.push_back(operand);
       }
       dynamicDimPos++;
     }
 
     // Create new memref type (which will have fewer dynamic dimensions).
     MemRefType newMemRefType =
         MemRefType::Builder(memrefType).setShape(newShapeConstants);
     assert(static_cast<int64_t>(newOperands.size()) ==
            newMemRefType.getNumDynamicDims());
 
     // Create and insert the alloc op for the new memref.
     auto newAlloc = rewriter.create<AllocLikeOp>(
-        alloc.getLoc(), newMemRefType, newOperands, alloc.alignmentAttr());
+        alloc.getLoc(), newMemRefType, newOperands, alloc.symbolOperands(),
+        alloc.alignmentAttr());
     // Insert a cast so we have the same type as the old alloc.
     auto resultCast =
         rewriter.create<CastOp>(alloc.getLoc(), newAlloc, alloc.getType());
 
     rewriter.replaceOp(alloc, {resultCast});
     return success();
   }
 };

mlir/test/Dialect/MemRef/canonicalize.mlir

 //  CHECK:   %[[cst:.+]] = memref.cast %arg
 //  CHECK:   memref.store %[[cst]]
 func @no_fold_of_store(%arg : memref<32xi8>, %holder: memref<memref<?xi8>>) {
   %0 = memref.cast %arg : memref<32xi8> to memref<?xi8>
   memref.store %0, %holder[] : memref<memref<?xi8>>
   return
 }
 
+// -----
+
+// CHECK-LABEL: func @simplify_alloc(
+//  CHECK: %[[mem1:.+]] = memref.alloc() : memref<1xi32>
+//  CHECK: %[[mem2:.+]] = memref.cast %[[mem1]] : memref<1xi32> to memref<?xi32>
+//  CHECK: return %[[mem2]] : memref<?xi32>
+func @simplify_alloc() -> memref<?xi32> {
+  %c1 = constant 1 : index
+  %0 = memref.alloc(%c1) : memref<?xi32>
+  return %0 : memref<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @simplify_alloc_ignore_symbols1(
+//  CHECK: %[[c1:.+]] = constant 1 : index
+//  CHECK: %[[mem1:.+]] = memref.alloc({{.*}})[%[[c1]], %[[c1]]] : memref<?xi32, #map>
+//  CHECK: return %[[mem1]] : memref<?xi32, #map>
+#map0 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+func @simplify_alloc_ignore_symbols1(%arg0 : index) -> memref<?xi32, #map0> {
+  %c1 = constant 1 : index
+  %0 = memref.alloc(%arg0)[%c1, %c1] : memref<?xi32, #map0>
+  return %0 : memref<?xi32, #map0>
+}
+
+// -----
+
+// CHECK-LABEL: func @simplify_alloc_ignore_symbols2(
+//  CHECK: %[[c1:.+]] = constant 1 : index
+//  CHECK: %[[mem1:.+]] = memref.alloc()[%[[c1]], %[[c1]]] : memref<1xi32, #map>
+//  CHECK: %[[mem2:.+]] = memref.cast %[[mem1]] : memref<1xi32, #map> to memref<?xi32, #map>
+//  CHECK: return %[[mem2]] : memref<?xi32, #map>
+#map0 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+func @simplify_alloc_ignore_symbols2() -> memref<?xi32, #map0> {
+  %c1 = constant 1 : index
+  %0 = memref.alloc(%c1)[%c1, %c1] : memref<?xi32, #map0>
+  return %0 : memref<?xi32, #map0>
+}