[mlir][tensor][bufferize] Fix deallocation of GenerateOp/FromElementsOp

[springerm]

Summary

Both ops allocate a buffer. There were cases in which the buffer was not deallocated.

Depends On D130463

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 template <typename OpTy, typename... Args>
 OpTy replaceOpWithNewBufferizedOp(RewriterBase &rewriter, Operation *op,
                                   Args &&...args) {
   auto newOp = rewriter.create<OpTy>(op->getLoc(), std::forward<Args>(args)...);
   replaceOpWithBufferizedValues(rewriter, op, newOp->getResults());
   return newOp;
 }
 
+/// Return `true` if the buffer of given OpResult should be deallocated. This
+/// function should be called during `BufferizableOpInterface::bufferize`
+/// implementations that allocate a new buffer for the given OpResult.
+bool shouldDeallocateOpResult(OpResult opResult,
+                              const BufferizationOptions &options);
+
 /// Return a MemRefType to which the type of the given value can be bufferized.
 ///
 /// If possible, op bufferization implementations should not use this function
 /// and instead infer precise memref types for tensor results by themselves.
 ///
 /// Unless a layout map was specified, `options.unknownTypeConverterFn`
 /// determines what kind of layout map will be used. For best composability
 /// (without copies), the fully dynamic layout map is used by default.

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

for (OpOperand *use : uses) {
       if (use->getOwner() != copy->getDefiningOp())
         rewriter.updateRootInPlace(use->getOwner(), [&]() { use->set(*copy); });
     }
   }
 
   return success();
 }
 
+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 (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;
+}
+
 //===----------------------------------------------------------------------===//
 // OpFilter
 //===----------------------------------------------------------------------===//
 
 bool OpFilter::isOpAllowed(Operation *op) const {
   // All other ops: Allow/disallow according to filter.
   bool isAllowed = !hasAllowRule();
   for (const Entry &entry : entries) {

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

LogicalResult AllocTensorOp::bufferize(RewriterBase &rewriter,
 
   // Create memory copy (if any).
   if (getCopy()) {
     if (failed(options.createMemCpy(rewriter, loc, copyBuffer, *alloc)))
       return failure();
   }
 
   // Should the buffer be deallocated?
-  AnalysisState analysisState(options);
-  bool dealloc;
-  if (op->hasAttr(BufferizationDialect::kEscapeAttrName)) {
-    // AllocTensorOp has one result.
-    ArrayAttr escapeAttr =
-        op->getAttr(BufferizationDialect::kEscapeAttrName).cast<ArrayAttr>();
-    dealloc = !escapeAttr[0].cast<BoolAttr>().getValue();
-  } else {
-    // No "escape" annotation found.
-    if (options.createDeallocs) {
-      // Perform an ad-hoc analysis.
-      dealloc = !analysisState.isTensorYielded(getResult());
-    } else {
-      dealloc = false;
-    }
-  }
+  bool dealloc =
+      shouldDeallocateOpResult(getResult().cast<OpResult>(), options);
 
   // Replace op.
   replaceOpWithBufferizedValues(rewriter, getOperation(), *alloc);
 
   // Create buffer deallocation (if requested).
   if (!dealloc)
     return success();
 

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

createStores(rewriter, loc, dim + 1, buffer, shape, constants, elementIt,
                  indices);
   }
 }
 
 /// Bufferization of tensor.from_elements.
 struct FromElementsOpInterface
     : public BufferizableOpInterface::ExternalModel<FromElementsOpInterface,
                                                     tensor::FromElementsOp> {
+
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto fromElementsOp = cast<tensor::FromElementsOp>(op);
+    // Should the buffer be deallocated?
+    bool dealloc = shouldDeallocateOpResult(
+        fromElementsOp.getResult().cast<OpResult>(), options);
 
     // TODO: Implement memory space for this op.
     if (options.defaultMemorySpace != static_cast<unsigned>(0))
       return op->emitError("memory space not implemented yet");
 
     // Allocate a buffer for the result.
     Location loc = op->getLoc();
     auto tensorType = fromElementsOp.getType().cast<RankedTensorType>();
     auto shape = tensorType.getShape();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
-    AnalysisState analysisState(options);
-    FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
-        rewriter, loc, fromElementsOp.getResult(),
-        analysisState.isTensorYielded(fromElementsOp.getResult()), options,
+    FailureOr<Value> tensorAlloc =
+        allocateTensorForShapedValue(rewriter, loc, fromElementsOp.getResult(),
+                                     /*escape=*/!dealloc, options,
                                      /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
     auto memrefType =
         MemRefType::get(tensorType.getShape(), tensorType.getElementType());
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, *tensorAlloc);
 
     // Case: tensor<0xelem_type>.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
 
     // Traverse all `elements` and create `memref.store` ops.
     auto elementIt = fromElementsOp.getElements().begin();
     SmallVector<Value, 2> indices(tensorType.getRank(), constants[0]);
     createStores(rewriter, loc, /*dim=*/0, buffer, shape, constants, elementIt,
                  indices);
 
     replaceOpWithBufferizedValues(rewriter, op, buffer);
+
     return success();
   }
 };
 
 /// Bufferization of tensor.generate.
 struct GenerateOpInterface
     : public BufferizableOpInterface::ExternalModel<GenerateOpInterface,
                                                     tensor::GenerateOp> {
+
+  bool bufferizesToAllocation(Operation *op, OpResult opResult) const {
+    return true;
+  }
+
   LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                           const BufferizationOptions &options) const {
     auto generateOp = cast<tensor::GenerateOp>(op);
+    // Should the buffer be deallocated?
+    bool dealloc = shouldDeallocateOpResult(
+        generateOp.getResult().cast<OpResult>(), options);
 
     // TODO: Implement memory space for this op.
     if (options.defaultMemorySpace != static_cast<unsigned>(0))
       return op->emitError("memory space not implemented yet");
 
     auto tensorType = generateOp.getType().cast<RankedTensorType>();
     // Allocate memory.
     Location loc = op->getLoc();
     // TODO: Create alloc_tensor ops during TensorCopyInsertion.
-    AnalysisState analysisState(options);
-    FailureOr<Value> tensorAlloc = allocateTensorForShapedValue(
-        rewriter, loc, generateOp.getResult(),
-        analysisState.isTensorYielded(generateOp.getResult()), options,
+    FailureOr<Value> tensorAlloc =
+        allocateTensorForShapedValue(rewriter, loc, generateOp.getResult(),
+                                     /*escape=*/!dealloc, options,
                                      /*copy=*/false);
     if (failed(tensorAlloc))
       return failure();
     auto memrefType =
         MemRefType::get(tensorType.getShape(), tensorType.getElementType());
     Value buffer = rewriter.create<bufferization::ToMemrefOp>(
         op->getLoc(), memrefType, *tensorAlloc);
 
     // Collect loop bounds.
LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     // about creating that.
     Operation *elementYield = parallelBody->getTerminator()->getPrevNode();
     rewriter.setInsertionPointAfter(elementYield);
     rewriter.replaceOpWithNewOp<memref::StoreOp>(
         elementYield, elementYield->getOperands()[0], buffer,
         parallelBody->getArguments());
 
     replaceOpWithBufferizedValues(rewriter, op, buffer);
+
     return success();
   }
 };
 
 /// Bufferization of tensor.insert. Replace with memref.store.
 struct InsertOpInterface
     : public BufferizableOpInterface::ExternalModel<InsertOpInterface,
                                                     tensor::InsertOp> {

mlir/test/Dialect/Tensor/one-shot-bufferize.mlir

%result = scf.foreach_thread (%thread_idx) in (%num_threads) -> tensor<200x100xf32> {
         // CHECK: memref.subview %{{.*}}[1, %{{.*}}] [1, 1] [1, 1] : memref<200x100xf32, #[[$MAP1]]> to memref<1xf32, #[[$MAP0]]>
         tensor.parallel_insert_slice %1 into %out[1, %thread_idx][1, 1][1, 1] :
           tensor<1xf32> into tensor<200x100xf32>
       }
   }
   // CHECK: }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @dealloc_generate_buffer
+func.func @dealloc_generate_buffer(%arg: tensor<*xf32>, %sz: index, %idx: index)
+  -> index
+{
+  // CHECK: memref.alloc
+  // CHECK: scf.parallel
+  // CHECK: memref.load
+  // CHECK: memref.dealloc
+  %0 = tensor.generate %sz {
+  ^bb0(%i : index):
+    %elem = tensor.dim %arg, %i : tensor<*xf32>
+    tensor.yield %elem : index
+  } : tensor<?xindex>
+  %r = tensor.extract %0[%idx] : tensor<?xindex>
+  return %r : index
+}