[mlir][linalg][bufferize] Fix bufferize bug where non-tensor ops are not skipped

[springerm]

Summary

Depends On D113961

[springerm]

update

[nicolasvasilache]

This is almost NFC but please improve the title and add a proper description.

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

if (!bufferizedFunctionTypes)
           "null bufferizedFunctionTypes when bufferizing CallOpInterface");
     return bufferize(b, callOp, state, *bufferizedFunctionTypes);
   }
 
   // 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`.
   if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
     return bufferizableOp.bufferize(b, state);
 
   // Other op with tensors. No bufferization method specified.
-  auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
-  if (any_of(op->getOperandTypes(), isaTensor) ||
-      any_of(op->getResultTypes(), isaTensor))
-    return op->emitError() << "unsupported op with tensors";
-  return success();
 }
 
 static LogicalResult bufferizeFuncOpInternals(
     FuncOp funcOp, BufferizationState &state,
     DenseMap<FuncOp, FunctionType> &bufferizedFunctionTypes) {
   LLVM_DEBUG(llvm::dbgs() << "\n\n");
   LDBG("Begin BufferizeFuncOpInternals:\n" << funcOp << '\n');
   OpBuilder b(funcOp->getContext());
struct TransferReadOpInterface
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto transferReadOp = cast<vector::TransferReadOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
-    if (transferReadOp.getShapedType().isa<MemRefType>())
-      return failure();
-
     // 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
   LogicalResult bufferize(Operation *op, OpBuilder &b,
                           BufferizationState &state) const {
     auto writeOp = cast<vector::TransferWriteOp>(op);
 
     // Take a guard before anything else.
     OpBuilder::InsertionGuard g(b);
     b.setInsertionPoint(op);
 
-    if (writeOp.getShapedType().isa<MemRefType>())
-      return failure();
-
     // 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);

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

 func @main() -> tensor<4xi32> {
   // expected-error @+1 {{unsupported op with tensors}}
   %r = scf.execute_region -> tensor<4xi32> {
     %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
     scf.yield %A: tensor<4xi32>
   }
   return %r: tensor<4xi32>
 }
-
-// -----
-
-func @main() -> i32 {
-  %c0 = arith.constant 0: index
-  // expected-error @+1 {{expected result-less scf.execute_region containing op}}
-  %r = scf.execute_region -> i32 {
-    %A = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
-    %e = tensor.extract %A[%c0]: tensor<4xi32>
-    scf.yield %e: i32
-  }
-  return %r: i32
-}