[mlir][Bufferize] NFC - Introduce areCastCompatible assertions to catch misformed CastOp early

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizableOpInterface.cpp

mlir::linalg::comprehensive_bufferize::BufferizationState::createAlloc(
   // Note: getAllocationTypeAndShape also sets the insertion point.
   MemRefType allocMemRefType =
       getAllocationTypeAndShape(b, loc, shapedValue, dynShape);
   FailureOr<Value> allocated = createAlloc(b, loc, allocMemRefType, dynShape);
   if (failed(allocated))
     return failure();
   Value casted = allocated.getValue();
   if (memRefType && memRefType != allocMemRefType) {
+    assert(memref::CastOp::areCastCompatible(allocated.getValue().getType(),
+                                             memRefType) &&
+           "createAlloc: cast incompatible");
     casted = b.create<memref::CastOp>(loc, memRefType, allocated.getValue());
   }
 
   if (deallocMemref) {
     // 2. Create memory deallocation.
     b.setInsertionPoint(allocated.getValue().getParentBlock()->getTerminator());
     createDealloc(b, loc, allocated.getValue());
   }

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/BufferizationInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Fold to_memref(to_tensor(x)) to x. Insert a cast if necessary.
     if (auto toTensorOp =
             toMemrefOp.tensor().getDefiningOp<bufferization::ToTensorOp>()) {
       Value buffer = toTensorOp.memref();
 
       // Insert cast in case to_memref(to_tensor(x))'s type is different from
       // x's type.
-      if (toTensorOp.memref().getType() != toMemrefOp.getType())
+      if (toTensorOp.memref().getType() != toMemrefOp.getType()) {
+        assert(memref::CastOp::areCastCompatible(buffer.getType(),
+                                                 toMemrefOp.getType()) &&
+               "ToMemrefOp::bufferize : cast incompatible");
         buffer = rewriter.create<memref::CastOp>(toMemrefOp.getLoc(), buffer,
                                                  toMemrefOp.getType());
+      }
       replaceOpWithBufferizedValues(rewriter, toMemrefOp, buffer);
       return success();
     }
 
     return failure();
   }
 };
 

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

for (unsigned idx = 0; idx < numArgs; ++idx) {
     // Get the buffer type from the bufferized function type.
     Type memrefType = bufferizedFuncType.getInput(idx);
     Value memref = frontBlock.addArgument(memrefType);
     OpBuilder b(funcOp->getContext());
     b.setInsertionPointToStart(&frontBlock);
     // Replace all uses of bbArg through a ToMemRefOp by a memref::CastOp.
     for (auto &use : llvm::make_early_inc_range(bbArg.getUses())) {
       if (auto toMemrefOp =
               dyn_cast<bufferization::ToMemrefOp>(use.getOwner())) {
+        assert(memref::CastOp::areCastCompatible(
+                   memref.getType(), toMemrefOp.memref().getType()) &&
+               "bufferizeFuncOpBoundary: cast incompatible");
         auto castOp = b.create<memref::CastOp>(
             funcOp.getLoc(), toMemrefOp.memref().getType(), memref);
         toMemrefOp.memref().replaceAllUsesWith(castOp);
       }
     }
     // Replace all remaining uses by a to_tensor.
     if (!bbArg.use_empty()) {
       auto toTensorOp =
for (const auto &it : llvm::enumerate(funcOp.getType().getInputs())) {
       argumentTypes.push_back(desiredMemrefType);
 
       // If funcOp's body is not empty, change the bbArg type and propagate.
       if (!funcOp.body().empty()) {
         BlockArgument bbArg = funcOp.getArgument(argNumber);
         bbArg.setType(desiredMemrefType);
         OpBuilder b(bbArg.getContext());
         b.setInsertionPointToStart(bbArg.getOwner());
+        assert(memref::CastOp::areCastCompatible(bbArg.getType(), memrefType) &&
+               "layoutPostProcessing: cast incompatible");
         // Cast back to the original memrefType and let it canonicalize.
         Value cast =
             b.create<memref::CastOp>(funcOp.getLoc(), memrefType, bbArg);
         bbArg.replaceAllUsesExcept(cast, cast.getDefiningOp());
       }
 
       // Cast to desired buffer type on all callers to `funcOp`.
       // TODO: on the callee side, this may even have to trigger a copy to
       // change the layout. For now let the memref::CastOp fail to verify in
       // such cases.
       auto castArg = [&](Operation *caller) {
         OpBuilder b(caller);
+        assert(
+            memref::CastOp::areCastCompatible(
+                caller->getOperand(argNumber).getType(), desiredMemrefType) &&
+            "layoutPostProcessing.2: cast incompatible");
         Value newOperand = b.create<memref::CastOp>(
             funcOp.getLoc(), desiredMemrefType, caller->getOperand(argNumber));
         operandsPerCaller.find(caller)->getSecond().push_back(newOperand);
       };
       foreachCaller(callerMap, funcOp, castArg);
     }
 
     // Set operands with cast buffer on all callers to `funcOp`.
for (OpOperand &opOperand : callOp->getOpOperands()) {
       // Caller / callee type mistmatch is handled with a CastOp.
       auto memRefType = bufferizedFuncType.getInput(idx);
       // Since we don't yet have a clear layout story, to_memref may
       // conservatively turn tensors into more dynamic memref than necessary.
       // If the memref type of the callee fails, introduce an extra memref.cast
       // that will either canonicalize away or fail compilation until we can do
       // something better.
       if (buffer.getType() != memRefType) {
+        assert(
+            memref::CastOp::areCastCompatible(buffer.getType(), memRefType) &&
+            "CallOp::bufferize: cast incompatible");
         Value castBuffer = rewriter.create<memref::CastOp>(callOp.getLoc(),
                                                            memRefType, buffer);
         buffer = castBuffer;
       }
       newOperands[idx] = buffer;
     }
 
     // 4. Create the new CallOp.

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/TensorInterfaceImpl.cpp

if (auto rankedTensorType = resultTensorType.isa<RankedTensorType>()) {
       resultMemRefType =
           getContiguousMemRefType(resultTensorType, layout, memorySpace);
     } else {
       resultMemRefType =
           getUnrankedMemRefType(resultTensorType.getElementType(), memorySpace);
     }
 
     // Replace the op with a memref.cast.
+    assert(memref::CastOp::areCastCompatible(resultBuffer->getType(),
+                                             resultMemRefType) &&
+           "CallOp::bufferize: cast incompatible");
     replaceOpWithNewBufferizedOp<memref::CastOp>(rewriter, op, resultMemRefType,
                                                  *resultBuffer);
 
     return success();
   }
 };
 
 /// Bufferization of tensor.dim. Replace with memref.dim.