[mlir] remove BufferAssignmentPlacer from BufferAssignmentOpConversionPattern

[dfki-ehna]

Summary

BufferPlacement has been removed, as allocations are no longer placed during the conversion.

[herhut]

Do you intend to remove the BufferAssignmentPlacer itself in a separate step?

mlir/include/mlir/Transforms/BufferPlacement.h

27

Can this also go?

175

If nullptr is not a usable default, then why have it anyway?

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

47

This should be a no-op as the rewriter should already be at the op positon.

mlir/lib/Transforms/BufferPlacement.cpp

880

This is a no-op, right?

mlir/test/lib/Transforms/TestBufferPlacement.cpp

59

Here, too.

[dfki-ehna]

BufferAssignmentPlacer is removed, and comments are addressed.

[herhut]

Thanks. Please submit a corresponding PR for TensorFlow before landing.

[herhut]

Also, as a nit, you could rephrase the summary to just read

BufferPlacement has been removed, as allocations are no longer placed during the conversion.

mlir/include/mlir/Transforms/BufferPlacement.h

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 
 /// Prepares a buffer placement phase. It can place (user-defined) alloc
 /// nodes. This simplifies the integration of the actual buffer-placement
 /// pass. Sample usage:
 ///   BufferAssignmentPlacer baHelper(regionOp);
 ///   -> determine alloc positions
 ///   auto allocPosition = baHelper.computeAllocPosition(value);
 ///   -> place alloc
 ///   allocBuilder.setInsertionPoint(positions.getAllocPosition());
 ///   <create alloc>

 ///     bufferAssignment->computeAllocPosition(resultOp);
 ///   }
 /// };
 template <typename SourceOp>
 class BufferAssignmentOpConversionPattern
     : public OpConversionPattern<SourceOp> {
 public:
   explicit BufferAssignmentOpConversionPattern(
-      MLIRContext *context, BufferAssignmentPlacer *bufferAssignment = nullptr,
-      BufferAssignmentTypeConverter *converter = nullptr,
+      MLIRContext *context, BufferAssignmentTypeConverter *converter = nullptr,
       PatternBenefit benefit = 1)
-      : OpConversionPattern<SourceOp>(context, benefit),
-        bufferAssignment(bufferAssignment), converter(converter) {
+      : OpConversionPattern<SourceOp>(context, benefit), converter(converter) {
     assert(converter && "The type converter has not been defined");
   }
 
 protected:
-  BufferAssignmentPlacer *bufferAssignment;
   BufferAssignmentTypeConverter *converter;
 };
 
 /// Converts the signature of the function using BufferAssignmentTypeConverter.
 /// Each result type of the function is kept as a function result or appended to
 /// the function arguments list based on ResultConversionKind for the converted
 /// result type.
 class BufferAssignmentFuncOpConverter
LogicalResult matchAndRewrite(CallOp, ArrayRef<Value>,
                                 ConversionPatternRewriter &) const;
 };
 
 /// Populates `patterns` with the conversion patterns of buffer
 /// assignment.
 template <typename ReturnOpSourceTy, typename ReturnOpTargetTy,
           typename CopyOpTy>
 static void populateWithBufferAssignmentOpConversionPatterns(
-    MLIRContext *context, BufferAssignmentPlacer *placer,
-    BufferAssignmentTypeConverter *converter,
+    MLIRContext *context, BufferAssignmentTypeConverter *converter,
     OwningRewritePatternList *patterns) {
   // clang-format off
   patterns->insert<
     BufferAssignmentCallOpConverter,
     BufferAssignmentFuncOpConverter,
     BufferAssignmentReturnOpConverter
       <ReturnOpSourceTy, ReturnOpTargetTy, CopyOpTy>
-  >(context, placer, converter);
+  >(context, converter);
   // clang-format on
 }
 } // end namespace mlir
 
 #endif // MLIR_TRANSFORMS_BUFFERPLACEMENT_H

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

matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
     Location loc = op.getLoc();
     ResultRange results = op.getOperation()->getResults();
     SmallVector<Value, 2> newArgs, newResults;
     newArgs.reserve(operands.size() + results.size());
     newArgs.append(operands.begin(), operands.end());
     newResults.reserve(results.size());
 
     // Update all types to memref types.
+    rewriter.setInsertionPoint(op);
     for (auto result : results) {
       auto type = result.getType().cast<ShapedType>();
       assert(type && "tensor to buffer conversion expects ranked results");
       if (!type.hasStaticShape())
         return rewriter.notifyMatchFailure(
             op, "dynamic shapes not currently supported");
       auto memrefType = MemRefType::get(type.getShape(), type.getElementType());
-
-      // Compute alloc position and insert a custom allocation node.
-      OpBuilder::InsertionGuard guard(rewriter);
-      rewriter.restoreInsertionPoint(
-          bufferAssignment->computeAllocPosition(result));
       auto alloc = rewriter.create<AllocOp>(loc, memrefType);
       newArgs.push_back(alloc);
       newResults.push_back(alloc);
     }
 
     // Generate a new linalg operation that works on buffers.
     auto linalgOp = rewriter.create<linalg::GenericOp>(
         loc, llvm::None, newArgs, rewriter.getI64IntegerAttr(operands.size()),
matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
     rewriter.replaceOp(op, newResults);
     return success();
   }
 };
 
 /// Populate the given list with patterns to convert Linalg operations on
 /// tensors to buffers.
 static void populateConvertLinalgOnTensorsToBuffersPattern(
-    MLIRContext *context, BufferAssignmentPlacer *placer,
-    BufferAssignmentTypeConverter *converter,
+    MLIRContext *context, BufferAssignmentTypeConverter *converter,
     OwningRewritePatternList *patterns) {
   populateWithBufferAssignmentOpConversionPatterns<
-      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, placer,
-                                                      converter, patterns);
-  patterns->insert<GenericOpConverter>(context, placer, converter);
+      mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, converter,
+                                                      patterns);
+  patterns->insert<GenericOpConverter>(context, converter);
 }
 
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct ConvertLinalgOnTensorsToBuffers
     : public LinalgOnTensorsToBuffersBase<ConvertLinalgOnTensorsToBuffers> {
   void runOnOperation() override {
     MLIRContext &context = getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
+    target.addLegalOp<ModuleOp>();
+    target.addLegalOp<ModuleTerminatorOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return converter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(
         Optional<ConversionTarget::DynamicLegalityCallbackFn>(
             isLegalOperation));
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp funcOp) {
              llvm::none_of(funcOp.getType().getResults(),
                            [&](Type type) { return type.isa<MemRefType>(); }) &&
              converter.isLegal(&funcOp.getBody());
     });
 
     converter.setResultConversionKind<RankedTensorType, MemRefType>(
         BufferAssignmentTypeConverter::AppendToArgumentsList);
 
-    // Walk over all the functions to apply buffer assignment.
-    getOperation().walk([&](FuncOp function) -> WalkResult {
-      OwningRewritePatternList patterns;
-      BufferAssignmentPlacer placer(function);
-      populateConvertLinalgOnTensorsToBuffersPattern(&context, &placer,
-                                                     &converter, &patterns);
-
-      // Applying full conversion
-      return applyFullConversion(function, target, patterns);
-    });
+    populateConvertLinalgOnTensorsToBuffersPattern(&context, &converter,
+                                                   &patterns);
+    if (failed(applyFullConversion(this->getOperation(), target, patterns)))
+      this->signalPassFailure();
   }
 };
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertLinalgOnTensorsToBuffersPass() {
   return std::make_unique<ConvertLinalgOnTensorsToBuffers>();
 }

mlir/lib/Transforms/BufferPlacement.cpp

for (auto operand : operands) {
     newOperands.append(values.begin(), values.end());
   }
 
   // Create the new result types for the new `CallOp` and a mapping from the old
   // result to new value(s).
   SmallVector<Type, 2> newResultTypes;
   SmallVector<ResultMapping, 4> mappings;
   mappings.resize(callOp.getNumResults());
+  rewriter.setInsertionPoint(callOp);
   for (auto result : llvm::enumerate(callOp.getResults())) {
     SmallVector<Type, 2> originTypes;
     converter->tryDecomposeType(result.value().getType(), originTypes);
     auto &resultMapping = mappings[result.index()];
     for (Type origin : originTypes) {
       Type converted = converter->convertType(origin);
       auto kind = converter->getResultConversionKind(origin, converted);
       if (kind == BufferAssignmentTypeConverter::KeepAsFunctionResult) {
         newResultTypes.push_back(converted);
         // The result value is not yet available. Its index is kept and it is
         // replaced with the actual value of the new `CallOp` later.
         resultMapping.addMapping(newResultTypes.size() - 1);
       } else {
         // kind = BufferAssignmentTypeConverter::AppendToArgumentsList
-        OpBuilder::InsertionGuard guard(rewriter);
-        rewriter.restoreInsertionPoint(
-            bufferAssignment->computeAllocPosition(result.value()));
         MemRefType memref = converted.dyn_cast<MemRefType>();
         if (!memref)
           return callOp.emitError("Cannot allocate for a non-Memref type");
         Value alloc = rewriter.create<AllocOp>(loc, memref);
         newOperands.push_back(alloc);
         resultMapping.addMapping(alloc);
       }
     }

mlir/test/lib/Transforms/TestBufferPlacement.cpp

matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
       Location loc = op.getLoc();
       ResultRange results = op.getOperation()->getResults();
       SmallVector<Value, 2> newArgs, newResults;
       newArgs.reserve(operands.size() + results.size());
       newArgs.append(operands.begin(), operands.end());
       newResults.reserve(results.size());
 
       // Update all types to memref types.
+      rewriter.setInsertionPoint(op);
       for (auto result : results) {
         ShapedType type = result.getType().cast<ShapedType>();
         assert(type && "Generic operations with non-shaped typed results are "
                        "not currently supported.");
         if (!type.hasStaticShape())
           return rewriter.notifyMatchFailure(
               op, "dynamic shapes not currently supported");
         auto memrefType =
             MemRefType::get(type.getShape(), type.getElementType());
-
-        // Compute alloc position and insert a custom allocation node.
-        OpBuilder::InsertionGuard guard(rewriter);
-        rewriter.restoreInsertionPoint(
-            bufferAssignment->computeAllocPosition(result));
         auto alloc = rewriter.create<AllocOp>(loc, memrefType);
         newArgs.push_back(alloc);
         newResults.push_back(alloc);
       }
 
       // Generate a new linalg operation that works on buffers.
       auto linalgOp = rewriter.create<linalg::GenericOp>(
           loc, llvm::None, newArgs, rewriter.getI64IntegerAttr(operands.size()),
matchAndRewrite(linalg::GenericOp op, ArrayRef<Value> operands,
       // Replace the results of the old Generic Op with the results of the new
       // one.
       rewriter.replaceOp(op, newResults);
       return success();
     }
   };
 
   void populateTensorLinalgToBufferLinalgConversionPattern(
-      MLIRContext *context, BufferAssignmentPlacer *placer,
-      BufferAssignmentTypeConverter *converter,
+      MLIRContext *context, BufferAssignmentTypeConverter *converter,
       OwningRewritePatternList *patterns) {
     populateWithBufferAssignmentOpConversionPatterns<
-        mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, placer,
-                                                        converter, patterns);
-    patterns->insert<GenericOpConverter>(context, placer, converter);
+        mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp>(context, converter,
+                                                        patterns);
+    patterns->insert<GenericOpConverter>(context, converter);
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<TestDialect>();
     registry.insert<linalg::LinalgDialect>();
   }
 
   void runOnOperation() override {
     MLIRContext &context = this->getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
     target.addLegalOp<MakeTupleOp>();
     target.addLegalOp<GetTupleElementOp>();
+    target.addLegalOp<ModuleOp>();
+    target.addLegalOp<ModuleTerminatorOp>();
 
     // Mark all Linalg operations illegal as long as they work on tensors.
     auto isLegalOperation = [&](Operation *op) {
       return converter.isLegal(op);
     };
     target.addDynamicallyLegalDialect<linalg::LinalgDialect>(isLegalOperation);
 
     // Mark Standard Return operations illegal as long as one operand is tensor.
converter.addDecomposeValueConversion([](OpBuilder &builder, Location loc,
       for (unsigned i = 0, e = resultType.size(); i < e; ++i) {
         Value res = builder.create<GetTupleElementOp>(
             loc, resultType.getType(i), value, builder.getI32IntegerAttr(i));
         values.push_back(res);
       }
       return success();
     });
 
-    // Walk over all the functions to apply buffer assignment.
-    this->getOperation().walk([&](FuncOp function) -> WalkResult {
-      OwningRewritePatternList patterns;
-      BufferAssignmentPlacer placer(function);
-      populateTensorLinalgToBufferLinalgConversionPattern(
-          &context, &placer, &converter, &patterns);
-
-      // Applying full conversion
-      return applyFullConversion(function, target, patterns);
-    });
+    populateTensorLinalgToBufferLinalgConversionPattern(&context, &converter,
+                                                        &patterns);
+    if (failed(applyFullConversion(this->getOperation(), target, patterns)))
+      this->signalPassFailure();
   };
 };
 } // end anonymous namespace
 
 namespace mlir {
 void registerTestBufferPlacementPreparationPass() {
   PassRegistration<
       TestBufferPlacementPreparationPass</*allowMemrefFunctionResults=*/false>>(