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;
TestTensorCopyInsertionPass__anonfd3215680111::TestTensorCopyInsertionPass32 TestTensorCopyInsertionPass(const TestTensorCopyInsertionPass &pass)
33 : PassWrapper(pass) {}
34
getDependentDialects__anonfd3215680111::TestTensorCopyInsertionPass35 void getDependentDialects(DialectRegistry ®istry) const override {
36 registry.insert<bufferization::BufferizationDialect>();
37 }
getArgument__anonfd3215680111::TestTensorCopyInsertionPass38 StringRef getArgument() const final { return "test-tensor-copy-insertion"; }
getDescription__anonfd3215680111::TestTensorCopyInsertionPass39 StringRef getDescription() const final {
40 return "Module pass to test Tensor Copy Insertion";
41 }
42
runOnOperation__anonfd3215680111::TestTensorCopyInsertionPass43 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(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(false)};
59 Option<bool> bufferizeFunctionBoundaries{
60 *this, "bufferize-function-boundaries",
61 llvm::cl::desc("Bufferize function boundaries."), llvm::cl::init(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(false)};
68 };
69 } // namespace
70
71 namespace mlir::test {
registerTestTensorCopyInsertionPass()72 void registerTestTensorCopyInsertionPass() {
73 PassRegistration<TestTensorCopyInsertionPass>();
74 }
75 } // namespace mlir::test
76