Skip to content

Commit 83d0616

Browse files
author
Jeff Niu
authored
[mlir][ods] Add documentation on how to use sharded op definitions (NFC) (#89664)
This adds explanations and instructions on how to set up a dialect for sharded op definitions to the MLIR documentation.
1 parent ae22ac9 commit 83d0616

File tree

1 file changed

+94
-0
lines changed

1 file changed

+94
-0
lines changed

mlir/docs/DefiningDialects/Operations.md

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1114,6 +1114,100 @@ void process(AddOp op, ArrayRef<Value> newOperands) {
11141114
}
11151115
```
11161116
1117+
#### Sharded Operation Definitions
1118+
1119+
Large dialects with many operations may struggle with C++ compile time of
1120+
generated op definitions, due to large compilation units. `mlir-tblgen`
1121+
provides the ability to shard op definitions by splitting them up evenly
1122+
by passing `-op-shard-count` to `-gen-op-defs` and `-gen-op-decls`. The tool
1123+
will generate a single include file for the definitions broken up by
1124+
`GET_OP_DEFS_${N}` where `${N}` is the shard number. A shard can be compiled in
1125+
a single compilation unit by adding a file like this to your dialect library:
1126+
1127+
```c++
1128+
#include "mlir/IR/Operation.h"
1129+
// Add any other required includes.
1130+
1131+
// Utilities shared by generated op definitions: custom directive parsers,
1132+
// printers, etc.
1133+
#include "OpUtils.h"
1134+
1135+
#define GET_OP_DEFS_0
1136+
#include "MyDialectOps.cpp.inc"
1137+
```
1138+
1139+
Note: this requires restructing shared utility functions within the dialect
1140+
library so they can be shared by multiple compilation units. I.e. instead of
1141+
defining `static` methods in the same source file, you should declare them in a
1142+
shared header and define them in their own source file.
1143+
1144+
The op registration hooks are also sharded, because the template instantiation
1145+
can take a very long time to compile. Operations should be registered in your
1146+
dialect like:
1147+
1148+
```c++
1149+
void MyDialect::initialize() {
1150+
registerMyDialectOperations(this);
1151+
}
1152+
```
1153+
1154+
CMake and Bazel functions are included to make sharding dialects easier.
1155+
Assuming you have organized your operation utility functions into their own
1156+
header, define a file that looks like the one above, but without the `#define`:
1157+
1158+
```c++
1159+
// MyDialectOps.cpp
1160+
#include "mlir/IR/Operation.h"
1161+
1162+
#include "OpUtils.h"
1163+
1164+
#include "MyDialectOps.cpp.inc"
1165+
```
1166+
1167+
In CMake, remove the manual `mlir_tablegen` invocations and replace them with:
1168+
1169+
```cmake
1170+
set(LLVM_TARGET_DEFINITIONS MyDialectOps.td)
1171+
add_sharded_ops(MyDialectOps 8) # shard the op definitions by 8
1172+
1173+
add_mlir_library(MyDialect
1174+
MyDialect.cpp
1175+
MyDialectOpDefs.cpp
1176+
${SHARDED_SRCS}
1177+
1178+
DEPENDS
1179+
MLIRTestOpsShardGen
1180+
)
1181+
```
1182+
1183+
This will automatically duplicate the `MyDialectOps.cpp` source file and add the
1184+
`#define` up the number of shards indicated.
1185+
1186+
It is recommended that any out-of-line op member functions (like verifiers) be
1187+
defined in a separate source file. In this example, it is called
1188+
`MyDialectOpDefs.cpp`.
1189+
1190+
In Bazel, remove the `-gen-op-defs` and `-gen-op-decls` invocations, and add
1191+
1192+
```bazel
1193+
gentbl_sharded_ops(
1194+
name = "MyDialectOpSrcs",
1195+
hdr_out = "MyDialectOps.h.inc",
1196+
shard_count = 8,
1197+
sharder = "//mlir:mlir-src-sharder",
1198+
src_file = "MyDialectOps.cpp",
1199+
src_out = "MyDialectOps.cpp.inc",
1200+
tblgen = "//mlir:mlir-tblgen",
1201+
td_file = "MyDialectOps.td",
1202+
deps = [":MyDialectOpsTdFiles"],
1203+
)
1204+
1205+
cc_library(
1206+
name = "MyDialect",
1207+
srcs = glob(["MyDialect/*.cpp"]) + [":MyDialectOpSrcs"]
1208+
)
1209+
```
1210+
11171211
## Constraints
11181212

11191213
Constraint is a core concept in table-driven operation definition: operation

0 commit comments

Comments
 (0)