[mlir][bufferization][NFC] Move EmptyTensorToAllocTensorPass

[springerm]

Summary

This change moves the pass from the Linalg dialect to the bufferization dialect.

Depends On D135129

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h

 /// Replace buffers that are returned from a function with an out parameter.
 /// Also update all call sites.
 LogicalResult promoteBufferResultsToOutParams(ModuleOp module);
 
 /// Creates a pass that drops memref function results that are equivalent to a
 /// function argument.
 std::unique_ptr<Pass> createDropEquivalentBufferResultsPass();
 
+/// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor.
+std::unique_ptr<Pass> createEmptyTensorToAllocTensorPass();
+
 /// Drop all memref function results that are equivalent to a function argument.
 LogicalResult dropEquivalentBufferResults(ModuleOp module);
 
 /// Creates a pass that finalizes a partial bufferization by removing remaining
 /// bufferization.to_tensor and bufferization.to_memref operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createFinalizingBufferizePass();
 
 /// Create a pass that bufferizes all ops that implement BufferizableOpInterface

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

let description = [{
 
     Note: If a bbArg buffer is not returned directly but casted to beforehand,
     the buffer is still considered equivalent.
   }];
   let constructor = "mlir::bufferization::createDropEquivalentBufferResultsPass()";
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
+def EmptyTensorToAllocTensor : Pass<"empty-tensor-to-alloc-tensor"> {
+  let summary = "Replace all empty ops by alloc_tensor ops.";
+  let description = [{
+    tensor.empty ops return a tensor of unspecified contents who's only purpose
+    is to carry the tensor shape. This pass converts such ops to
+    bufferization.alloc_tensor ops, which bufferize to buffer allocations.
+  }];
+  let constructor = "mlir::bufferization::createEmptyTensorToAllocTensorPass()";
+  let dependentDialects = ["tensor::TensorDialect"];
+}
+
 def OneShotBufferize : Pass<"one-shot-bufferize", "ModuleOp"> {
   let summary = "One-Shot Bufferize";
   let description = [{
     This pass bufferizes all ops that implement `BufferizableOpInterface`. It
     first performs an inplacability analysis on SSA use-def chains of tensor
     values to determine which OpOperands may bufferize in-place, i.e., without
     inserting a buffer copy. It then rewrites the IR, inserting a buffer
     allocation and copy for each OpOperand that was decided to bufferize

mlir/include/mlir/Dialect/Linalg/Passes.h

 createConvertLinalgToParallelLoopsPass();
 
 /// Create a pass to convert Linalg operations to affine.for loops and
 /// affine_load/affine_store accesses.
 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<func::FuncOp>>
 createConvertLinalgToAffineLoopsPass();
 
-/// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor.
-std::unique_ptr<Pass> createEmptyTensorToAllocTensorPass();
-
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgBufferizePass();
 
 /// Create a pass to convert named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createLinalgGeneralizationPass();
 

mlir/include/mlir/Dialect/Linalg/Passes.td

let description = [{
     This pass only converts ops that operate on ranked tensors. It can be
     run on op which contains linalg ops (most commonly a
     FunctionOpInterface op).
   }];
   let constructor = "mlir::createConvertElementwiseToLinalgPass()";
   let dependentDialects = ["linalg::LinalgDialect", "memref::MemRefDialect"];
 }
 
-def EmptyTensorToAllocTensor : Pass<"empty-tensor-to-alloc-tensor"> {
-  let summary = "Replace all empty ops by alloc_tensor ops.";
-  let description = [{
-    tensor.empty ops return a tensor of unspecified contents who's only purpose
-    is to carry the tensor shape. This pass converts such ops to
-    bufferization.alloc_tensor ops, which bufferize to buffer allocations.
-  }];
-  let constructor = "mlir::createEmptyTensorToAllocTensorPass()";
-}
-
 def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [
     Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
             /*default=*/"false",
            "Only folds the one-trip loops from Linalg ops on tensors "
            "(for testing purposes only)">

mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRBufferizationTransforms
   AllocTensorElimination.cpp
   Bufferize.cpp
   BufferDeallocation.cpp
   BufferOptimizations.cpp
   BufferResultsToOutParams.cpp
   BufferUtils.cpp
   BufferViewFlowAnalysis.cpp
   DropEquivalentBufferResults.cpp
+  EmptyTensorToAllocTensor.cpp
   FuncBufferizableOpInterfaceImpl.cpp
   OneShotAnalysis.cpp
   OneShotModuleBufferize.cpp
   TensorCopyInsertion.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Bufferization
 

mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorToAllocTensor.cpp

-//===- InitTensorToAllocTensor.cpp - Lower init_tensor to alloc_tensor ----===//
+//===- InitTensorToAllocTensor.cpp - Lower tensor.empty to alloc_tensor ---===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
+namespace bufferization {
 #define GEN_PASS_DEF_EMPTYTENSORTOALLOCTENSOR
-#include "mlir/Dialect/Linalg/Passes.h.inc"
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
+} // namespace bufferization
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::bufferization;
 using namespace mlir::tensor;
 
 namespace {
 struct EmptyTensorLoweringPattern : public OpRewritePattern<tensor::EmptyOp> {
   using OpRewritePattern<tensor::EmptyOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(tensor::EmptyOp op,
                                 PatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<bufferization::AllocTensorOp>(
         op, op.getType(), op.getDynamicSizes());
     return success();
   }
 };
 
 struct EmptyTensorToAllocTensor
-    : public impl::EmptyTensorToAllocTensorBase<EmptyTensorToAllocTensor> {
+    : public bufferization::impl::EmptyTensorToAllocTensorBase<
+          EmptyTensorToAllocTensor> {
   EmptyTensorToAllocTensor() = default;
 
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<tensor::TensorDialect, bufferization::BufferizationDialect>();
   }
 };
 } // namespace
 
 void EmptyTensorToAllocTensor::runOnOperation() {
   Operation *op = getOperation();
   RewritePatternSet patterns(op->getContext());
   patterns.insert<EmptyTensorLoweringPattern>(op->getContext());
   if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
     signalPassFailure();
 }
 
-std::unique_ptr<Pass> mlir::createEmptyTensorToAllocTensorPass() {
+std::unique_ptr<Pass>
+mlir::bufferization::createEmptyTensorToAllocTensorPass() {
   return std::make_unique<EmptyTensorToAllocTensor>();
 }

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   ElementwiseOpFusion.cpp
   ElementwiseToLinalg.cpp
   FusePadOpWithLinalgProducer.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Generalization.cpp
   Hoisting.cpp
   HoistPadding.cpp
-  InitTensorToAllocTensor.cpp
   InlineScalarOperands.cpp
   Interchange.cpp
   Loops.cpp
   LinalgStrategyPasses.cpp
   NamedOpConversions.cpp
   Promotion.cpp
   Split.cpp
   SplitReduction.cpp