[mlir][linalg][bufferize][NFC] Call `bufferize` on all bufferizable ops

[springerm]

Summary

Call this function even on ops that do not have tensor OpOperands/OpResults.

Depends On D114055

[mehdi_amini]

It isn't obvious to me that this is NFC?

[springerm]

In D114056#3138984, @mehdi_amini wrote:

It isn't obvious to me that this is NFC?

For the current interface implementations, this change is NFC. Before this change, it was guaranteed that the bufferize method was never called on ops that do not need bufferization (because they have no tensor results/operands). After this change, bufferize is called for every bufferizable op. If there is no work to do, the interface method simply returns success. This is in preparation of changing the way that operations are traversed in one of the subsequent commits.

I'll update the commit message when landing this change.

[springerm]

Note, I may change this back to the original behavior again later. This change is mostly to keep the commits small during an overall larger refactoring.

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

 mlir::linalg::comprehensive_bufferize::bufferizeOp(Operation *op,
                                                    BufferizationState &state) {
   OpBuilder b(op->getContext());
 
   // Skip BufferCast and TensorLoad ops.
   if (isa<memref::BufferCastOp, memref::TensorLoadOp>(op))
     return success();
 
-  // Check if op has tensor results or operands.
-  auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
-  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
-  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
-  if (!hasTensorResult && !hasTensorOperand)
-    return success();
-
   // Bufferize using `BufferizableOpInterface`.
   b.setInsertionPoint(op);
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
     return bufferizableOp.bufferize(b, state);
 
-  // Other op with tensors. No bufferization method specified.
+  // Check if op has tensor results or operands.
+  auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
+  bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
+  bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+  if (hasTensorResult && hasTensorOperand)
     return op->emitError() << "unsupported op with tensors";
+
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Bufferization-specific BlockAndValueMapping support with debugging.
 //===----------------------------------------------------------------------===//
 
 /// Wrapper for better debugging.
 void mlir::linalg::comprehensive_bufferize::BufferizationState::mapBuffer(

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

struct TransferReadOpInterface
 
   OpResult getAliasingOpResult(Operation *op, OpOperand &opOperand) const {
     return OpResult();
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto transferReadOp = cast<vector::TransferReadOp>(op);
+    if (!transferReadOp.getShapedType().isa<TensorType>())
+      return success();
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
     // TransferReadOp always reads from the bufferized op.source().
-    assert(transferReadOp.getShapedType().isa<TensorType>() &&
-           "only tensor types expected");
     Value v = state.lookupBuffer(transferReadOp.source());
     transferReadOp.sourceMutable().assign(v);
     return success();
   }
 };
 
 struct TransferWriteOpInterface
     : public BufferizableOpInterface::ExternalModel<TransferWriteOpInterface,
struct TransferWriteOpInterface
 
   BufferRelation bufferRelation(Operation *op, OpOperand &opOperand) const {
     return BufferRelation::Equivalent;
   }
 
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
+    if (!writeOp.getShapedType().isa<TensorType>())
+      return success();
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
     // Create a new transfer_write on buffer that doesn't have a return value.
     // Leave the previous transfer_write to dead code as it still has uses at
     // this point.
-    assert(writeOp.getShapedType().isa<TensorType>() &&
-           "only tensor types expected");
     Value resultBuffer = getResultBuffer(b, op->getResult(0), state);
     if (!resultBuffer)
       return failure();
     b.create<vector::TransferWriteOp>(
         writeOp.getLoc(), writeOp.vector(), resultBuffer, writeOp.indices(),
         writeOp.permutation_map(),
         writeOp.in_bounds() ? *writeOp.in_bounds() : ArrayAttr());
     state.mapBuffer(op->getResult(0), resultBuffer);