1//===- TestTensorCopyInsertion.cpp - Bufferization Analysis -----*- c++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
10#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
11#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
12#include "mlir/Pass/Pass.h"
13
14using namespace mlir;
15
16namespace {
17/// This pass runs One-Shot Analysis and inserts copies for all OpOperands that
18/// were decided to bufferize out-of-place. After running this pass, a
19/// bufferization can write to buffers directly (without making copies) and no
20/// longer has to care about potential read-after-write conflicts.
21///
22/// Note: By default, all newly inserted tensor copies/allocs (i.e., newly
23/// created `bufferization.alloc_tensor` ops) that do not escape block are
24/// annotated with `escape = false`. If `create-allocs` is unset, all newly
25/// inserted tensor copies/allocs are annotated with `escape = true`. In that
26/// case, they are not getting deallocated when bufferizing the IR.
27struct TestTensorCopyInsertionPass
28 : public PassWrapper<TestTensorCopyInsertionPass, OperationPass<ModuleOp>> {
29 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestTensorCopyInsertionPass)
30
31 TestTensorCopyInsertionPass() = default;
32 TestTensorCopyInsertionPass(const TestTensorCopyInsertionPass &pass)
33 : PassWrapper(pass) {}
34
35 void getDependentDialects(DialectRegistry &registry) const override {
36 registry.insert<bufferization::BufferizationDialect>();
37 }
38 StringRef getArgument() const final { return "test-tensor-copy-insertion"; }
39 StringRef getDescription() const final {
40 return "Module pass to test Tensor Copy Insertion";
41 }
42
43 void runOnOperation() override {
44 bufferization::OneShotBufferizationOptions options;
45 options.allowReturnAllocsFromLoops = allowReturnAllocsFromLoops;
46 options.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
47 if (mustInferMemorySpace) {
48 options.defaultMemorySpaceFn =
49 [](TensorType t) -> std::optional<Attribute> { return std::nullopt; };
50 }
51
52 bufferization::BufferizationState bufferizationState;
53
54 if (failed(bufferization::insertTensorCopies(getOperation(), options,
55 bufferizationState)))
56 signalPassFailure();
57 }
58
59 Option<bool> allowReturnAllocsFromLoops{
60 *this, "allow-return-allocs-from-loops",
61 llvm::cl::desc("Allows returning/yielding new allocations from a loop."),
62 llvm::cl::init(Val: false)};
63 Option<bool> bufferizeFunctionBoundaries{
64 *this, "bufferize-function-boundaries",
65 llvm::cl::desc("Bufferize function boundaries."), llvm::cl::init(Val: false)};
66 Option<bool> mustInferMemorySpace{
67 *this, "must-infer-memory-space",
68 llvm::cl::desc(
69 "The memory space of an memref types must always be inferred. If "
70 "unset, a default memory space of 0 is used otherwise."),
71 llvm::cl::init(Val: false)};
72};
73} // namespace
74
75namespace mlir::test {
76void registerTestTensorCopyInsertionPass() {
77 PassRegistration<TestTensorCopyInsertionPass>();
78}
79} // namespace mlir::test
80

source code of mlir/test/lib/Dialect/Bufferization/TestTensorCopyInsertion.cpp