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 | |
14 | using namespace mlir; |
15 | |
16 | namespace { |
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. |
27 | struct 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 ®istry) 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 | if (failed(result: bufferization::insertTensorCopies(getOperation(), options))) |
52 | signalPassFailure(); |
53 | } |
54 | |
55 | Option<bool> allowReturnAllocsFromLoops{ |
56 | *this, "allow-return-allocs-from-loops" , |
57 | llvm::cl::desc("Allows returning/yielding new allocations from a loop." ), |
58 | llvm::cl::init(Val: false)}; |
59 | Option<bool> bufferizeFunctionBoundaries{ |
60 | *this, "bufferize-function-boundaries" , |
61 | llvm::cl::desc("Bufferize function boundaries." ), llvm::cl::init(Val: false)}; |
62 | Option<bool> mustInferMemorySpace{ |
63 | *this, "must-infer-memory-space" , |
64 | llvm::cl::desc( |
65 | "The memory space of an memref types must always be inferred. If " |
66 | "unset, a default memory space of 0 is used otherwise." ), |
67 | llvm::cl::init(Val: false)}; |
68 | }; |
69 | } // namespace |
70 | |
71 | namespace mlir::test { |
72 | void registerTestTensorCopyInsertionPass() { |
73 | PassRegistration<TestTensorCopyInsertionPass>(); |
74 | } |
75 | } // namespace mlir::test |
76 | |