[mlir] [bufferization] Fix dealloc errors.

[cxy-1993]

Summary

This patch solves problem that bufferization pass generate dealloc on inappropriate location.

Let me know if you need RFC on this patch, thanks for your time.

In fact, this patch can not pass test. If input IR is:

 mlir
#map = affine_map<(d0) -> (d0 * 5)>                                                
func.func @ternimator_use_not_deallocated(%arg0: tensor<10x10xf32>) -> tensor<10x10xf32> {
  %0 = scf.forall (%arg1, %arg2) in (2, 2) shared_outs(%arg3 = %arg0) -> (tensor<10x10xf32>) {
    %1 = bufferization.alloc_tensor() : tensor<5x5xf32>                            
    %2 = affine.apply #map(%arg1)                                                  
    %3 = affine.apply #map(%arg2)                                                  
    scf.forall.in_parallel {                                                       
      tensor.parallel_insert_slice %1 into %arg3[%2, %3] [5, 5] [1, 1] : tensor<5x5xf32> into tensor<10x10xf32>
    }                                                                              
  }                                                                                
  return %0 : tensor<10x10xf32>                                                    
}

when bufferize with param:

-one-shot-bufferize="allow-unknown-ops copy-before-write"

will add copy before tensor.parallel_insert_slice and cause failure. Please kindly show me how to solve this problem, thanks.

[springerm]

This is tricky... I think the ParallelInsertSliceOp should not bufferize to a memory write. Then, a copy would never be inserted there. I.e., ParallelInsertSliceOpInterface::bufferizesToMemoryWrite should always return false. Can you give that a try and see if it fails any tests? I haven't fully thought this through yet and may be missing something...

[springerm]

Can you try this patch: https://reviews.llvm.org/D148408. I ran into a similar issue today and I think this should fix it.

[cxy-1993]

In D147790#4270628, @springerm wrote:

Can you try this patch: https://reviews.llvm.org/D148408. I ran into a similar issue today and I think this should fix it.

Thanks for your reply, Sorry for not replying to the message in time.
This patch have solved my problem, thanks.

As https://reviews.llvm.org/D148408 , is any op contains in a terminator region (e.g. : tensor.parallel_insert_slice ) should move user before terminator after bufferize?

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

 
 bool bufferization::shouldDeallocateOpResult(
     OpResult opResult, const BufferizationOptions &options) {
   Operation *op = opResult.getOwner();
   assert(options.dynCastBufferizableOp(op).bufferizesToAllocation(opResult) &&
          "expected that op allocates");
 
   AnalysisState analysisState(options);
+  if (analysisState.isTensorYielded(opResult))
+    return false;
+
   if (op->hasAttr(BufferizationDialect::kEscapeAttrName)) {
     // AllocTensorOp has one result.
     ArrayAttr escapeAttr =
         op->getAttr(BufferizationDialect::kEscapeAttrName).cast<ArrayAttr>();
     return !escapeAttr[0].cast<BoolAttr>().getValue();
   }
 
   // No "escape" annotation found.
-  if (options.createDeallocs) {
-    // Perform an ad-hoc analysis.
-    return !analysisState.isTensorYielded(opResult);
-  }
-
-  return false;
+  return options.createDeallocs;
 }
 
 //===----------------------------------------------------------------------===//
 // OpFilter
 //===----------------------------------------------------------------------===//
 
 bool OpFilter::isOpAllowed(Operation *op) const {
   // All other ops: Allow/disallow according to filter.
while (!worklist.empty()) {
     if (!options.dynCastBufferizableOp(op))
       continue;
 
     // We cannot analyze through ToMemrefOps, so we have to conservatively
     // assume that the value is yielded.
     if (isa<ToMemrefOp>(op))
       return true;
 
-    // Check if the op is returning/yielding.
-    if (isRegionReturnLike(op))
+    // Check if the op is recursively returning/yielding.
+    Operation *parentOp = op;
+    do {
+      if (isRegionReturnLike(parentOp) ||
+          parentOp->hasTrait<OpTrait::IsTerminator>())
         return true;
+    } while ((parentOp = parentOp->getParentOp()));
 
     // Add all aliasing OpResults to the worklist.
     // Note: In the absence of detailed analysis information (e.g., there may be
     // no function call analysis information), this `getAliasingOpResult` is
     // conservative and may report additional OpResults as potentially aliasing.
     for (AliasingOpResult alias : getAliasingOpResults(*operand))
       for (OpOperand &use : alias.opResult.getUses())
         worklist.push_back(&use);

mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

func.func @read_of_alias(%t: tensor<100xf32>, %pos1: index, %pos2: index,
   // CHECK-TOP-DOWN-ANALYSIS: memref.copy
   // CHECK-TOP-DOWN-ANALYSIS: memref.store %{{.*}}, %[[alloc]]
   %0 = tensor.insert %f into %t[%pos1] : tensor<100xf32>
   %1 = tensor.extract_slice %t[%pos2][%sz][1] : tensor<100xf32> to tensor<?xf32>
   %2 = tensor.extract %1[%pos3] : tensor<?xf32>
   %3 = tensor.extract %0[%pos3] : tensor<100xf32>
   return %2, %3 : f32, f32
 }
+
+// -----
+
+// CHECK-LABEL: func @ternimator_use_not_deallocated
+#map = affine_map<(d0) -> (d0 * 5)>
+func.func @ternimator_use_not_deallocated(%arg0: tensor<10x10xf32>) -> tensor<10x10xf32> {
+  // CHECK: %[[alloc:.*]] = memref.alloc
+  // CHECK: memref.copy {{.*}} %[[alloc]]
+  // CHECK: scf.forall ({{.*}}) in (2, 2) {
+  %0 = scf.forall (%arg1, %arg2) in (2, 2) shared_outs(%arg3 = %arg0) -> (tensor<10x10xf32>) {
+    // CHECK: %[[local_alloc:.*]] = memref.alloc
+    // CHECK-NOT: memref.dealloc
+    // CHECK: %[[subview:.*]] = memref.subview %[[alloc]]
+    // CHECK: memref.copy %[[local_alloc]], %[[subview]]
+    %1 = bufferization.alloc_tensor() : tensor<5x5xf32>
+    %2 = affine.apply #map(%arg1)
+    %3 = affine.apply #map(%arg2)
+    scf.forall.in_parallel {
+      tensor.parallel_insert_slice %1 into %arg3[%2, %3] [5, 5] [1, 1] : tensor<5x5xf32> into tensor<10x10xf32>
+    }
+  }
+  return %0 : tensor<10x10xf32>
+}