xref: /llvm-project/mlir/include/mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h (revision b0a4e958e85784cff46303c92b6a3a14b20fa1d8)
1 //===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
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 #ifndef MLIR_BUFFERIZATION_TRANSFORMS_FUNCBUFFERIZABLEOPINTERFACEIMPL_H
10 #define MLIR_BUFFERIZATION_TRANSFORMS_FUNCBUFFERIZABLEOPINTERFACEIMPL_H
11 
12 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
13 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
14 #include "mlir/Dialect/Func/IR/FuncOps.h"
15 #include "llvm/ADT/SmallVector.h"
16 
17 namespace mlir {
18 class DialectRegistry;
19 
20 namespace func {
21 class FuncOp;
22 } // namespace func
23 
24 namespace bufferization {
25 /// Helper function that returns all func.return ops in the given function.
26 SmallVector<func::ReturnOp> getReturnOps(func::FuncOp funcOp);
27 
28 namespace func_ext {
29 /// The state of analysis of a FuncOp.
30 enum class FuncOpAnalysisState { NotAnalyzed, InProgress, Analyzed };
31 
32 using func::FuncOp;
33 
34 /// Extra analysis state that is required for bufferization of function
35 /// boundaries.
36 struct FuncAnalysisState : public OneShotAnalysisState::Extension {
37   FuncAnalysisState(OneShotAnalysisState &state)
38       : OneShotAnalysisState::Extension(state) {}
39 
40   // Note: Function arguments and/or function return values may disappear during
41   // bufferization. Functions and their CallOps are analyzed and bufferized
42   // separately. To ensure that a CallOp analysis/bufferization can access an
43   // already bufferized function's analysis results, we store bbArg/return value
44   // indices instead of BlockArguments/OpOperand pointers.
45 
46   /// A set of block argument indices.
47   using BbArgIndexSet = DenseSet<int64_t>;
48 
49   /// A mapping of indices to indices.
50   using IndexMapping = DenseMap<int64_t, int64_t>;
51 
52   /// A mapping of indices to a list of indices.
53   using IndexToIndexListMapping = DenseMap<int64_t, SmallVector<int64_t>>;
54 
55   /// A mapping of ReturnOp OpOperand indices to equivalent FuncOp BBArg
56   /// indices.
57   DenseMap<FuncOp, IndexMapping> equivalentFuncArgs;
58 
59   /// A mapping of FuncOp BBArg indices to aliasing ReturnOp OpOperand indices.
60   DenseMap<FuncOp, IndexToIndexListMapping> aliasingReturnVals;
61 
62   /// A set of all read BlockArguments of FuncOps.
63   DenseMap<FuncOp, BbArgIndexSet> readBbArgs;
64 
65   /// A set of all written-to BlockArguments of FuncOps.
66   DenseMap<FuncOp, BbArgIndexSet> writtenBbArgs;
67 
68   /// Keep track of which FuncOps are fully analyzed or currently being
69   /// analyzed.
70   DenseMap<FuncOp, FuncOpAnalysisState> analyzedFuncOps;
71 
72   /// This function is called right before analyzing the given FuncOp. It
73   /// initializes the data structures for the FuncOp in this state object.
74   void startFunctionAnalysis(FuncOp funcOp);
75 };
76 
77 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
78 } // namespace func_ext
79 } // namespace bufferization
80 } // namespace mlir
81 
82 #endif // MLIR_BUFFERIZATION_TRANSFORMS_FUNCBUFFERIZABLEOPINTERFACEIMPL_H
83