diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.h @@ -46,6 +46,9 @@ /// function argument. std::unique_ptr createDropEquivalentBufferResultsPass(); +/// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor. +std::unique_ptr createEmptyTensorToAllocTensorPass(); + /// Drop all memref function results that are equivalent to a function argument. LogicalResult dropEquivalentBufferResults(ModuleOp module); diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td --- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td +++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td @@ -168,6 +168,17 @@ 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 = [{ diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.h b/mlir/include/mlir/Dialect/Linalg/Passes.h --- a/mlir/include/mlir/Dialect/Linalg/Passes.h +++ b/mlir/include/mlir/Dialect/Linalg/Passes.h @@ -61,9 +61,6 @@ std::unique_ptr> createConvertLinalgToAffineLoopsPass(); -/// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor. -std::unique_ptr createEmptyTensorToAllocTensorPass(); - /// Create a pass to convert Linalg operations which work on tensors to use /// buffers instead. std::unique_ptr> createLinalgBufferizePass(); diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td --- a/mlir/include/mlir/Dialect/Linalg/Passes.td +++ b/mlir/include/mlir/Dialect/Linalg/Passes.td @@ -24,16 +24,6 @@ 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()"; diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt @@ -7,6 +7,7 @@ BufferUtils.cpp BufferViewFlowAnalysis.cpp DropEquivalentBufferResults.cpp + EmptyTensorToAllocTensor.cpp FuncBufferizableOpInterfaceImpl.cpp OneShotAnalysis.cpp OneShotModuleBufferize.cpp diff --git a/mlir/lib/Dialect/Linalg/Transforms/InitTensorToAllocTensor.cpp b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorToAllocTensor.cpp rename from mlir/lib/Dialect/Linalg/Transforms/InitTensorToAllocTensor.cpp rename to mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorToAllocTensor.cpp --- a/mlir/lib/Dialect/Linalg/Transforms/InitTensorToAllocTensor.cpp +++ b/mlir/lib/Dialect/Bufferization/Transforms/EmptyTensorToAllocTensor.cpp @@ -1,4 +1,4 @@ -//===- 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. @@ -6,7 +6,7 @@ // //===----------------------------------------------------------------------===// -#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" @@ -14,8 +14,10 @@ #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; @@ -35,7 +37,8 @@ }; struct EmptyTensorToAllocTensor - : public impl::EmptyTensorToAllocTensorBase { + : public bufferization::impl::EmptyTensorToAllocTensorBase< + EmptyTensorToAllocTensor> { EmptyTensorToAllocTensor() = default; void runOnOperation() override; @@ -55,6 +58,7 @@ signalPassFailure(); } -std::unique_ptr mlir::createEmptyTensorToAllocTensorPass() { +std::unique_ptr +mlir::bufferization::createEmptyTensorToAllocTensorPass() { return std::make_unique(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt --- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt @@ -15,7 +15,6 @@ Generalization.cpp Hoisting.cpp HoistPadding.cpp - InitTensorToAllocTensor.cpp InlineScalarOperands.cpp Interchange.cpp Loops.cpp