[WIP] Only insert memref_cast when needed during canonicalization.

[bondhugula]

Please include a brief commit summary. I only did a superficial review. Could you please include test cases?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

312

Doc comment with ///

313

Typo

322

More compact to use:

for (auto *userOp : oldOp->getUsers())

329

Likewise.

335

You could use replaceOpWithNewOp(oldOp, oldValue.getLoc(), ...)

[ftynse]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

311

MLIR/LLVM prefer static to anonymous namespaces for functions

320

Either auto * or Operation *.

330

I'm not certain setType is okay inside a rewrite pattern. Can you find another way?

335

Nit: you don't need oldValue.getLoc(), the location will be extracted from oldOp.

[rriddle]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

330

Should be fine if you wrap this loop in: rewriter.startRootUpdate()/rewriter.finalizeRootUpdate()

[mehdi_amini]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

322

Operation * instead of auto please. In general please avoid auto where is does not improve readability / when the type isn't obvious from the context.

[poechsel]

I prefer to abandon this revision as I believe the way I'm doing things is the wrong way to do them:

• This change is already breaking with some recent diffs of LLVM.
• I really dislike having to add a "blacklist" instead of a whitelist.
• I believe a new canonicalisation pattern would be more suited (ie less invasive).

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(AllocOp op) {
 
   // Verify that all operands are of type Index.
   for (auto operandType : op.getOperandTypes())
     if (!operandType.isIndex())
       return op.emitOpError("requires operands to be of type Index");
   return success();
 }
 
 namespace {
+// Create a memref_cast when needed to convert from newType to the type of
+// oldOp. Called when canonilizing a view or subview op that changed the type of
+// the memref (from a dynamic-sized memref to a statically-sized memref). Only
+// inserts the memref_cast when the corresponding memref is used as argument for
+// a function call or is returned directly.
+void replaceOpInsertMemrefCastWhenNeeded(PatternRewriter &rewriter,
+                                         Value oldValue, Value newValue) {
+  bool canSkipMemrefCast = true;
+  auto oldOp = oldValue.getDefiningOp();
+  auto uses = oldOp->getUses();
+  for (auto it = uses.begin(); it != uses.end(); it++) {
+    auto *user = it.getUser();
+    canSkipMemrefCast &= !(isa<CallOp>(user)
+                           || isa<CallIndirectOp>(user)
+                           || isa<ReturnOp>(user));
+  }
+  if (canSkipMemrefCast) {
+    for (auto it = uses.begin(); it != uses.end(); it++)
+      it->get().setType(newValue.getType());
+    rewriter.replaceOp(oldOp, {newValue});
+  } else {
+    auto castOp = rewriter.create<MemRefCastOp>(oldValue.getLoc(), newValue,
+                                                oldValue.getType());
+    rewriter.replaceOp(oldOp, {castOp});
+  }
+}
+
 /// Fold constant dimensions into an alloc operation.
 struct SimplifyAllocConst : public OpRewritePattern<AllocOp> {
   using OpRewritePattern<AllocOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(AllocOp alloc,
                                 PatternRewriter &rewriter) const override {
     // Check to see if any dimensions operands are constants.  If so, we can
     // substitute and drop them.
LogicalResult matchAndRewrite(AllocOp alloc,
     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<AllocOp>(alloc.getLoc(), newMemRefType,
                                              newOperands, IntegerAttr());
-    // Insert a cast so we have the same type as the old alloc.
-    auto resultCast = rewriter.create<MemRefCastOp>(alloc.getLoc(), newAlloc,
-                                                    alloc.getType());
-
-    rewriter.replaceOp(alloc, {resultCast});
+    replaceOpInsertMemrefCastWhenNeeded(rewriter, alloc, newAlloc);
     return success();
   }
 };
 
 /// Fold alloc operations with no uses. Alloc has side effects on the heap,
 /// but can still be deleted if it has zero uses.
 struct SimplifyDeadAlloc : public OpRewritePattern<AllocOp> {
   using OpRewritePattern<AllocOp>::OpRewritePattern;
for (auto size : llvm::enumerate(subViewOp.sizes())) {
       assert(defOp);
       staticShape[size.index()] = cast<ConstantIndexOp>(defOp).getValue();
     }
     MemRefType newMemRefType =
         MemRefType::Builder(subViewType).setShape(staticShape);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), subViewOp.offsets(),
         ArrayRef<Value>(), subViewOp.strides(), newMemRefType);
-    // Insert a memref_cast for compatibility of the uses of the op.
-    rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, newSubViewOp,
-                                              subViewOp.getType());
+    replaceOpInsertMemrefCastWhenNeeded(rewriter, subViewOp, newSubViewOp);
     return success();
   }
 };
 
 // Pattern to rewrite a subview op with constant stride arguments.
 class SubViewOpStrideFolder final : public OpRewritePattern<SubViewOp> {
 public:
   using OpRewritePattern<SubViewOp>::OpRewritePattern;
LogicalResult matchAndRewrite(SubViewOp subViewOp,
     }
     AffineMap layoutMap = makeStridedLinearLayoutMap(
         staticStrides, resultOffset, rewriter.getContext());
     MemRefType newMemRefType =
         MemRefType::Builder(subViewType).setAffineMaps(layoutMap);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), subViewOp.offsets(),
         subViewOp.sizes(), ArrayRef<Value>(), newMemRefType);
-    // Insert a memref_cast for compatibility of the uses of the op.
-    rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, newSubViewOp,
-                                              subViewOp.getType());
+    replaceOpInsertMemrefCastWhenNeeded(rewriter, subViewOp, newSubViewOp);
     return success();
   }
 };
 
 // Pattern to rewrite a subview op with constant offset arguments.
 class SubViewOpOffsetFolder final : public OpRewritePattern<SubViewOp> {
 public:
   using OpRewritePattern<SubViewOp>::OpRewritePattern;
LogicalResult matchAndRewrite(SubViewOp subViewOp,
 
     AffineMap layoutMap = makeStridedLinearLayoutMap(
         resultStrides, staticOffset, rewriter.getContext());
     MemRefType newMemRefType =
         MemRefType::Builder(subViewType).setAffineMaps(layoutMap);
     auto newSubViewOp = rewriter.create<SubViewOp>(
         subViewOp.getLoc(), subViewOp.source(), ArrayRef<Value>(),
         subViewOp.sizes(), subViewOp.strides(), newMemRefType);
-    // Insert a memref_cast for compatibility of the uses of the op.
-    rewriter.replaceOpWithNewOp<MemRefCastOp>(subViewOp, newSubViewOp,
-                                              subViewOp.getType());
+    replaceOpInsertMemrefCastWhenNeeded(rewriter, subViewOp, newSubViewOp);
     return success();
   }
 };
 
 } // end anonymous namespace
 
 void SubViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                             MLIRContext *context) {
LogicalResult matchAndRewrite(ViewOp viewOp,
     (void)dynamicOffsetOperandCount; // unused in opt mode
     assert(static_cast<int64_t>(newOperands.size()) ==
            dynamicOffsetOperandCount + newMemRefType.getNumDynamicDims());
 
     // Create new ViewOp.
     auto newViewOp = rewriter.create<ViewOp>(viewOp.getLoc(), newMemRefType,
                                              viewOp.getOperand(0), newOperands);
     // Insert a cast so we have the same type as the old memref type.
-    rewriter.replaceOpWithNewOp<MemRefCastOp>(viewOp, newViewOp,
-                                              viewOp.getType());
+    replaceOpInsertMemrefCastWhenNeeded(rewriter, viewOp, newViewOp);
     return success();
   }
 };
 
 struct ViewOpMemrefCastFolder : public OpRewritePattern<ViewOp> {
   using OpRewritePattern<ViewOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ViewOp viewOp,