Skip to content

Commit a672319

Browse files
matthias-springermatthias-springer
committed
[mlir][linalg][bufferize] Add BufferizableOpInterface
This commit adds a new op interface: BufferizableOpInterface. In the future, ops that implement this interface can be bufferized using Comprehensive Bufferize. Note: The interface methods of this interface correspond to the "op interface" in ComprehensiveBufferize.cpp. Differential Revision: https://reviews.llvm.org/D112974
1 parent 6503117 commit a672319

File tree

9 files changed

+335
-14
lines changed

9 files changed

+335
-14
lines changed

‎mlir/include/mlir/Dialect/Linalg/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
add_subdirectory(IR)
2+
add_subdirectory(Transforms)
23

34
set(LLVM_TARGET_DEFINITIONS Passes.td)
45
mlir_tablegen(Passes.h.inc -gen-pass-decls -name Linalg)

‎mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
//===- BufferizableOpInterface.h - Comprehensive Bufferize ------*- 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+
#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
10+
#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_
11+
12+
#include "mlir/IR/Builders.h"
13+
#include "mlir/IR/BuiltinTypes.h"
14+
#include "mlir/IR/Operation.h"
15+
#include "mlir/Support/LLVM.h"
16+
17+
namespace mlir {
18+
class BlockAndValueMapping;
19+
20+
namespace linalg {
21+
class AllocationCallbacks;
22+
class BufferizationAliasInfo;
23+
24+
/// Specify fine-grain relationship between buffers to enable more analysis.
25+
enum class BufferRelation {
26+
None,
27+
// TODO: ResultContainsOperand,
28+
// TODO: OperandContainsResult,
29+
Equivalent
30+
};
31+
} // namespace linalg
32+
} // namespace mlir
33+
34+
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h.inc"
35+
36+
#endif // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE_H_

‎mlir/include/mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.td

Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
//===-- BufferizableOpInterface.td - Compreh. Bufferize ----*- tablegen -*-===//
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_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
10+
#define MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE
11+
12+
include "mlir/IR/OpBase.td"
13+
14+
def BufferizableOpInterface : OpInterface<"BufferizableOpInterface"> {
15+
let description = [{
16+
An op interface for Comprehensive Bufferization. Ops that implement this
17+
interface can be bufferized using Comprehensive Bufferization.
18+
}];
19+
let cppNamespace = "::mlir::linalg";
20+
let methods = [
21+
InterfaceMethod<
22+
/*desc=*/[{
23+
Return `true` if the given OpOperand bufferizes to a memory read. This
24+
method will never be called on OpOperands that do not have a tensor
25+
type.
26+
27+
Note: It is always safe to consider an OpOperand as a memory read,
28+
even if it does actually not read; however, this can introduce
29+
unnecessary out-of-place bufferization decisions. The analysis of
30+
Comprehensive Bufferize considers OpOperands of unknown ops (that do
31+
not implement this interface) as reading OpOperands.
32+
}],
33+
/*retType=*/"bool",
34+
/*methodName=*/"bufferizesToMemoryRead",
35+
/*args=*/(ins "OpOperand &":$opOperand),
36+
/*methodBody=*/"",
37+
/*defaultImplementation=*/[{
38+
// Does not have to be implemented for ops without tensor OpOperands.
39+
llvm_unreachable("bufferizesToMemoryRead not implemented");
40+
}]
41+
>,
42+
InterfaceMethod<
43+
/*desc=*/[{
44+
Return `true` if the given OpOperand bufferizes to a memory write.
45+
This method will never be called on OpOperands that do not have a
46+
tensor type.
47+
48+
Note: It is always safe to consider an OpOperand as a memory write,
49+
even if it does actually not write; however, this can introduce
50+
unnecessary out-of-place bufferization decisions. The analysis of
51+
Comprehensive Bufferize considers OpOperands of unknown ops (that do
52+
not implement this interface) as writing OpOperands.
53+
}],
54+
/*retType=*/"bool",
55+
/*methodName=*/"bufferizesToMemoryWrite",
56+
/*args=*/(ins "OpOperand &":$opOperand),
57+
/*methodBody=*/"",
58+
/*defaultImplementation=*/[{
59+
// Does not have to be implemented for ops without tensor OpOperands.
60+
llvm_unreachable("bufferizesToMemoryWrite not implemented");
61+
}]
62+
>,
63+
// TODO: Simplify this interface by removing `bufferizesToAliasOnly` and
64+
// `getInplaceableOpResult`. Instead, always use `getAliasingOpResult`. If
65+
// `getAliasingOpResult` returns a non-null value, we know that an alias
66+
// is created. If `bufferizesToMemoryRead` and `bufferizesToMemoryWrite`
67+
// return `false`, we know that the operands "bufferizes to alias only".
68+
InterfaceMethod<
69+
/*desc=*/[{
70+
Return `true` if the given OpOperand creates an alias but does neither
71+
read nor write. This implies that `bufferizesToMemoryRead` and
72+
`bufferizesToMemoryWrite` must return `false`. This method will never
73+
be called on OpOperands that do not have a tensor type.
74+
75+
Examples of such ops are `tensor.extract_slice` and `tensor.cast`.
76+
}],
77+
/*retType=*/"bool",
78+
/*methodName=*/"bufferizesToAliasOnly",
79+
/*args=*/(ins "OpOperand &":$opOperand),
80+
/*methodBody=*/"",
81+
/*defaultImplementation=*/[{
82+
// For better debugging, run `bufferizesToMemoryWrite`, which fires an
83+
// assertion when called on an op that should not have tensor
84+
// OpOperands.
85+
(void) cast<BufferizableOpInterface>($_op.getOperation())
86+
.bufferizesToMemoryWrite(opOperand);
87+
// Return `false` by default, as most ops are not "alias only".
88+
return false;
89+
}]
90+
>,
91+
InterfaceMethod<
92+
/*desc=*/[{
93+
Return the OpResult that can bufferize in-place with a given
94+
OpOperand. Return a null value if the OpOperand cannot bufferize
95+
in-place. This method will never be called on OpOperands that do not
96+
have a tensor type.
97+
}],
98+
/*retType=*/"OpResult",
99+
/*methodName=*/"getInplaceableOpResult",
100+
/*args=*/(ins "OpOperand &":$opOperand),
101+
/*methodBody=*/"",
102+
/*defaultImplementation=*/[{
103+
// Does not have to be implemented for ops without tensor OpOperands.
104+
llvm_unreachable("getInplaceableOpResult not implemented");
105+
}]
106+
>,
107+
InterfaceMethod<
108+
/*desc=*/[{
109+
Return the OpResult that aliases with a given OpOperand when
110+
bufferized in-place. This is a superset of `getInplaceableOpResult`.
111+
This method will never be called on OpOperands that do not have a
112+
tensor type.
113+
}],
114+
/*retType=*/"OpResult",
115+
/*methodName=*/"getAliasingOpResult",
116+
/*args=*/(ins "OpOperand &":$opOperand),
117+
/*methodBody=*/"",
118+
/*defaultImplementation=*/[{
119+
return cast<BufferizableOpInterface>($_op.getOperation())
120+
.getInplaceableOpResult(opOperand);
121+
}]
122+
>,
123+
InterfaceMethod<
124+
/*desc=*/[{
125+
Return the OpOperands that alias with a given OpResult when
126+
bufferized in-place. This method will never be called on OpResults
127+
that do not have a tensor type.
128+
129+
Note: This method can return multiple OpOperands, indicating that the
130+
given OpResult may at runtime alias with any of the OpOperands. This
131+
is useful for branches and for ops such as `std.select`.
132+
}],
133+
/*retType=*/"SmallVector<OpOperand *>",
134+
/*methodName=*/"getAliasingOpOperand",
135+
/*args=*/(ins "OpResult":$opResult),
136+
/*methodBody=*/"",
137+
/*defaultImplementation=*/[{
138+
// Does not have to be implemented for ops without tensor OpResults.
139+
llvm_unreachable("getInplaceableOpResult not implemented");
140+
}]
141+
>,
142+
InterfaceMethod<
143+
/*desc=*/[{
144+
Return the buffer relation between the given OpOperand and its
145+
aliasing OpResult when bufferized in-place. Most OpOperands have an
146+
"equivalence" relation.
147+
148+
TODO: Support other relations such as "OpOperand is included in
149+
OpResult".
150+
}],
151+
/*retType=*/"BufferRelation",
152+
/*methodName=*/"bufferRelation",
153+
/*args=*/(ins "OpOperand &":$opOperand),
154+
/*methodBody=*/"",
155+
/*defaultImplementation=*/[{
156+
// Does not have to be implemented for ops without tensor OpOperands.
157+
llvm_unreachable("bufferRelation not implemented");
158+
}]
159+
>,
160+
// TODO: Simplify method signature: Pass an OpBuilder and a
161+
// BufferizationState object.
162+
InterfaceMethod<
163+
/*desc=*/[{
164+
Bufferize this op, i.e., rewrite it into a memref-based equivalent.
165+
`bvm` maps tensor values to memref values and this method should map
166+
tensor results to memref results after creating/modifying ops.
167+
}],
168+
/*retType=*/"LogicalResult",
169+
/*methodName=*/"bufferize",
170+
/*args=*/(ins "OpBuilder &":$b,
171+
"BlockAndValueMapping &":$bvm,
172+
"BufferizationAliasInfo &":$aliasInfo,
173+
"AllocationCallbacks &":$allocationFn),
174+
/*methodBody=*/"",
175+
/*defaultImplementation=*/[{
176+
llvm_unreachable("bufferize not implemented");
177+
return failure();
178+
}]
179+
>,
180+
];
181+
}
182+
183+
#endif // MLIR_DIALECT_LINALG_TRANSFORMS_BUFFERIZABLEOPINTERFACE

‎mlir/include/mlir/Dialect/Linalg/Transforms/CMakeLists.txt

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
set(LLVM_TARGET_DEFINITIONS BufferizableOpInterface.td)
2+
mlir_tablegen(BufferizableOpInterface.h.inc -gen-op-interface-decls)
3+
mlir_tablegen(BufferizableOpInterface.cpp.inc -gen-op-interface-defs)
4+
add_public_tablegen_target(MLIRBufferizableOpInterfaceIncGen)
5+
add_dependencies(mlir-headers MLIRBufferizableOpInterfaceIncGen)

‎mlir/include/mlir/Dialect/Linalg/Transforms/ComprehensiveBufferize.h

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,6 @@ namespace linalg {
3434
/// uses BufferizationAliasInfo.
3535
class BufferizationAliasInfo {
3636
public:
37-
/// Specify fine-grain relationship between buffers to enable more analysis.
38-
enum class BufferRelation {
39-
None,
40-
// TODO: ResultContainsOperand,
41-
// TODO: OperandContainsResult,
42-
Equivalent
43-
};
44-
4537
explicit BufferizationAliasInfo(Operation *rootOp);
4638

4739
/// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the

‎mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
//===- BufferizableOpInterface.cpp - Comprehensive Bufferize --------------===//
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/Linalg/Transforms/BufferizableOpInterface.h"
10+
11+
namespace mlir {
12+
namespace linalg {
13+
14+
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.cpp.inc"
15+
16+
} // namespace linalg
17+
} // namespace mlir

‎mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,38 @@
1+
set(LLVM_OPTIONAL_SOURCES
2+
Bufferize.cpp
3+
BufferizableOpInterface.cpp
4+
CodegenStrategy.cpp
5+
ComprehensiveBufferize.cpp
6+
Detensorize.cpp
7+
Distribution.cpp
8+
DropUnitDims.cpp
9+
ElementwiseOpFusion.cpp
10+
ElementwiseToLinalg.cpp
11+
Fusion.cpp
12+
FusionOnTensors.cpp
13+
Generalization.cpp
14+
Hoisting.cpp
15+
HoistPadding.cpp
16+
InlineScalarOperands.cpp
17+
Interchange.cpp
18+
Loops.cpp
19+
LinalgStrategyPasses.cpp
20+
Promotion.cpp
21+
Tiling.cpp
22+
Transforms.cpp
23+
Vectorization.cpp
24+
)
25+
26+
add_mlir_dialect_library(MLIRBufferizableOpInterface
27+
BufferizableOpInterface.cpp
28+
29+
DEPENDS
30+
MLIRBufferizableOpInterfaceIncGen
31+
32+
LINK_LIBS PUBLIC
33+
MLIRIR
34+
)
35+
136
add_mlir_dialect_library(MLIRLinalgTransforms
237
Bufferize.cpp
338
CodegenStrategy.cpp
@@ -32,6 +67,7 @@ add_mlir_dialect_library(MLIRLinalgTransforms
3267
MLIRAffineUtils
3368
MLIRAnalysis
3469
MLIRArithmetic
70+
MLIRBufferizableOpInterface
3571
MLIRComplex
3672
MLIRInferTypeOpInterface
3773
MLIRIR

‎mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,7 @@
112112
#include "PassDetail.h"
113113
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
114114
#include "mlir/Dialect/Linalg/Passes.h"
115+
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterface.h"
115116
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
116117
#include "mlir/Dialect/Linalg/Utils/Utils.h"
117118
#include "mlir/Dialect/MemRef/IR/MemRef.h"
@@ -138,8 +139,6 @@ using namespace mlir;
138139
using namespace linalg;
139140
using namespace tensor;
140141

141-
using BufferRelation = BufferizationAliasInfo::BufferRelation;
142-
143142
#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
144143
#define LDBG(X) LLVM_DEBUG(DBGS() << X)
145144

0 commit comments

Comments
 (0)