Differential D138915 Diff 479618 mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp

Changeset View

Standalone View

mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp

This file was added.

				//===- TestTensorCopyInsertion.cpp - Bufferization Analysis ------ c++ --===//
				//
				// 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/Bufferization/IR/Bufferization.h"
				#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
				#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
				#include "mlir/Pass/Pass.h"

				using namespace mlir;

				namespace {
				/// This pass runs One-Shot Analysis and inserts copies for all OpOperands that
				/// were decided to bufferize out-of-place. After running this pass, a
				/// bufferization can write to buffers directly (without making copies) and no
				/// longer has to care about potential read-after-write conflicts.
				///
				/// Note: By default, all newly inserted tensor copies/allocs (i.e., newly
				/// created `bufferization.alloc_tensor` ops) that do not escape block are
				/// annotated with `escape = false`. If `create-allocs` is unset, all newly
				/// inserted tensor copies/allocs are annotated with `escape = true`. In that
				/// case, they are not getting deallocated when bufferizing the IR.
				struct TestTensorCopyInsertionPass
				: public PassWrapper<TestTensorCopyInsertionPass, OperationPass<ModuleOp>> {
				MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTensorCopyInsertionPass)

				TestTensorCopyInsertionPass() = default;
				TestTensorCopyInsertionPass(const TestTensorCopyInsertionPass &pass)
				: PassWrapper(pass) {}

				void getDependentDialects(DialectRegistry &registry) const override {
				registry.insert<bufferization::BufferizationDialect>();
				}
				StringRef getArgument() const final { return "test-tensor-copy-insertion"; }
				StringRef getDescription() const final {
				return "Module pass to test Tensor Copy Insertion";
				}

				void runOnOperation() override {
				bufferization::OneShotBufferizationOptions options;
				options.allowReturnAllocs = allowReturnAllocs;
				options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
				options.createDeallocs = createDeallocs;
				if (mustInferMemorySpace)
				options.defaultMemorySpace = None;
				if (failed(bufferization::insertTensorCopies(getOperation(), options)))
				signalPassFailure();
				}

				Option<bool> allowReturnAllocs{
				*this, "allow-return-allocs",
				llvm::cl::desc("Allows returning/yielding new allocations from a block."),
				llvm::cl::init(false)};
				Option<bool> bufferizeFunctionBoundaries{
				*this, "bufferize-function-boundaries",
				llvm::cl::desc("Bufferize function boundaries."), llvm::cl::init(false)};
				Option<bool> createDeallocs{
				*this, "create-deallocs",
				llvm::cl::desc("Specify if new allocations should be deallocated."),
				llvm::cl::init(true)};
				Option<bool> mustInferMemorySpace{
				*this, "must-infer-memory-space",
				llvm::cl::desc(
				"The memory space of an memref types must always be inferred. If "
				"unset, a default memory space of 0 is used otherwise."),
				llvm::cl::init(false)};
				};
				} // namespace

				namespace mlir::test {
				void registerTestTensorCopyInsertionPass() {
				PassRegistration<TestTensorCopyInsertionPass>();
				}
				} // namespace mlir::test