[mlir][linalg] Fix bufferize_to_allocation error checking

[springerm]

Summary

bufferize_to_allocation does not supports ops with regions, unless bufferize_destination_only is set. In that case, only the operand is replaced with an allocation and wrapped in a to_tensor op. The error checking was too strict.

[springerm]

Closed by b76a180d3b1db1e5ad6fdddfdef0d01788a6a0da.

mlir/lib/Dialect/Linalg/Transforms/ConvertToDestinationStyle.cpp

Value linalg::bufferizeToAllocation(
   // Only bufferizable ops are supported.
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
   BufferizationOptions bufferizationOptions;
   AnalysisState state(bufferizationOptions);
 
 #ifndef NDEBUG
+  if (!options.bufferizeDestinationOnly) {
     // Ops with nested tensor ops are not supported yet. At the moment, this
     // function just bufferizes the given op itself, but not its body.
     op->walk([&](Operation *nestedOp) {
       if (op == nestedOp)
         return;
       if (llvm::any_of(nestedOp->getOperands(),
                        [](Value v) { return v.getType().isa<TensorType>(); }))
         llvm_unreachable("ops with nested tensor ops are not supported yet");
       if (llvm::any_of(nestedOp->getResults(),
                        [](Value v) { return v.getType().isa<TensorType>(); }))
         llvm_unreachable("ops with nested tensor ops are not supported yet");
     });
+  }
 #endif // NDEBUG
 
   // Gather tensor results.
   SmallVector<OpResult> tensorResults;
   for (OpResult result : op->getResults()) {
     if (!result.getType().isa<TensorType>())
       continue;
     // Unranked tensors are not supported
for (const AliasingOpOperand &operand : aliasingOperands) {
       for (OpOperand &resultUse : result.getUses())
         resultUses.push_back(&resultUse);
     }
   }
   for (OpOperand &operand : op->getOpOperands()) {
     if (!state.bufferizesToMemoryWrite(operand))
       continue;
     if (!isa<RankedTensorType>(operand.get().getType()))
-      return nullptr;
+      continue;
     addOutOfPlaceOperand(&operand);
   }
   // TODO: Support multiple buffers.
   if (outOfPlaceOperands.size() != 1)
     return nullptr;
 
   // Allocate buffers.
   OpBuilder::InsertionGuard g(rewriter);

mlir/test/Dialect/Linalg/transform-op-bufferize-to-allocation.mlir

func.func @tensor_insert_destination(%t: tensor<?x10xindex>, %idx: index, %v: index) -> tensor<?x10xindex> {
   return %r : tensor<?x10xindex>
 }
 
 transform.sequence failures(propagate) {
 ^bb1(%arg1: !transform.any_op):
   %0 = transform.structured.match ops{["tensor.insert"]} in %arg1 : (!transform.any_op) -> !transform.any_op
   %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only} : !transform.any_op
 }
+
+// -----
+
+// CHECK-LABEL: func @scf_for_destination(
+//  CHECK-SAME:     %[[t:.*]]: tensor<?x10xindex>
+//       CHECK:   %[[alloc:.*]] = memref.alloc(%{{.*}}) : memref<?x10xindex, 4>
+//       CHECK:   memref.tensor_store %[[t]], %[[alloc]]
+//       CHECK:   %[[t2:.*]] = bufferization.to_tensor %[[alloc]] restrict writable
+//       CHECK:   %[[for:.*]] = scf.for {{.*}} iter_args(%{{.*}} = %[[t2]])
+//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   return %[[for]]
+func.func @scf_for_destination(%t: tensor<?x10xindex>, %lb: index, %ub: index, %step: index) -> tensor<?x10xindex> {
+  %r = scf.for %iv = %lb to %ub step %step iter_args(%a = %t) -> tensor<?x10xindex> {
+    %b = "test.foo"(%a) : (tensor<?x10xindex>) -> (tensor<?x10xindex>)
+    scf.yield %b : tensor<?x10xindex>
+  }
+  return %r : tensor<?x10xindex>
+}
+
+transform.sequence failures(propagate) {
+^bb1(%arg1: !transform.any_op):
+  %0 = transform.structured.match ops{["scf.for"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+  %2, %new = transform.structured.bufferize_to_allocation %0 {memory_space = 4, bufferize_destination_only} : !transform.any_op
+}