Page MenuHomePhabricator

[mlir][ods] Allow sharding of op definitions
Needs ReviewPublic

Authored by Mogball on Jun 10 2022, 11:55 AM.

Details

Summary

This patch adds an option to gen-op-defs and gen-op-decls called op-shard-count that, when set to a value greater than 1, splits the op class definitions into N shards and generates N registration functions, i.e. calls to addOperation<...>, guarded by the definition GET_OP_DEFS_<N>. It will also generate a top-level registration hook registerFooDialectOperations(FooDialect *) that calls each of the sharded registration hooks.

Sharding the op definitions and registration hooks can greatly speed up compilation time, especially for dialects with large opsets. To assist in sharding, a new buildrule is added (gentbl_sharded_ops in bazel and add_mlir_sharded_ops) which will copy a template source file and compile it multiple times with each shard.

Fixes #51789

Diff Detail

Event Timeline

Mogball created this revision.Jun 10 2022, 11:55 AM
Mogball requested review of this revision.Jun 10 2022, 11:55 AM
Mogball edited the summary of this revision. (Show Details)Jun 10 2022, 11:58 AM
Mogball updated this revision to Diff 436270.Jun 12 2022, 9:56 PM

Add a build rule for automatically sharding ops.

Any feedback on the build setup is welcome, by the way.

rriddle added inline comments.Jun 13 2022, 12:09 AM
mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp
18

Can you drop this using?

81

Is the only thing this tool does? i.e. it just prints a #define GET_OP_DEFS_<N> line? I'm not really sure this tool pulls it's weight. Can you describe the file setup that you are trying to get from cmake?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp
3153–3164

Do we really need new generators? Could you not just detect and dispatch accordingly based on the shard count? i.e. if there is only one shard (or if the shard count option wasn't specified) do something unique?

3167

What is this used for?

mlir/tools/mlir-tblgen/OpGenHelpers.h
29–31

AFAICT you don't really need to copy into different std::vectors here, you're just chunking the input vector. It's much more optimal in that case to just use ArrayRef.

Mogball added inline comments.Jun 13 2022, 12:36 AM
mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp
81

Ideally, what I want is to compile the same source file with different definitions *in the same CMake/Bazel library*, i.e.

clang++ foo.cpp -DSHARD_0 -o foo0.o
clang++ foo.cpp -DSHARD_1 -o foo1.o
ld foodialect.o foo0.o. foo1.o

But I wasn't able to find a way to pull that off. CMake only allows setting per-file defines once per file and Bazel only allows local_defines per build target (I think... I'm not a build file wizard). Doing it with separate libraries would likely end up creating cyclic dependencies or otherwise make the dependency management of a dialect kind of tricky.

I originally used echo -e "#define ..\n${cat src.cpp}" but I'm not sure that works on windows.

Mogball updated this revision to Diff 436606.Jun 13 2022, 4:54 PM
Mogball marked 4 inline comments as done.

review comments

Mogball edited the summary of this revision. (Show Details)Jun 13 2022, 9:30 PM
rriddle added inline comments.Jun 29 2022, 12:38 AM
mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp
81

If it's just about generating files, file(GENERATE or configure_file( should probably serve what you need. Worst case, when generating the file we would just do something like:

#define GET_OP_DEFS_${OP_DEF_N}
#include "src.cpp"

With ${OP_DEF_N} referencing a variable we set in cmake.

Another option of course, would be to consider what it would take to let tablegen generate multiple files...

Mogball added inline comments.Jun 29 2022, 12:05 PM
mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp
81

I'll give that a try. Thanks!

The problem with generating multiple files is that they all still the correct includes...

Mogball updated this revision to Diff 442967.Jul 7 2022, 10:08 AM

review comments

Mogball updated this revision to Diff 442969.Jul 7 2022, 10:11 AM

wrong diff

Mogball updated this revision to Diff 442972.Jul 7 2022, 10:13 AM

wrong diff