kuhar commented 1 week ago

This is an uber-issue for making the tuner easier to maintain. The current implementation has a few issues that make the tuner library fragile and prone to getting out of sync with the IREE compiler. Specifically, we identified the following issues:

There are two ways to (re-)configure executable sources: a. By updating the lowering config and translation info in-situ. This is used when producing candidate dispatches using executable benchmarks as the source-of-truth. b. By using the transform dialect library script to match root ops and apply compilation info attributes to them. This is used during the model candidate compilation and benchmarking stage.

As a result, we have duplicate logic to apply configurations found by the constraint solver. The fix is to write a pass that strips existing configuration from executable sources, and then use transform dialect to re-configure them. This can be done as a separate invocation of iree-opt.
The MLIR processing is mostly string-based. While this allowed us to quickly prototype, it makes the code prone to getting out of sync with the IREE compiler. The lowering configs and translation info attributes are considered compiler-internals and there's no stability guarantee as to the exact structure and format of these attributes. As a result, every time the format changes, we have to update the parsing and printing logic in the tuner to match the new format in the compiler.

Here, the proposed solution is to expose these key attributes (translation info, compilation info, and MFMA intrinsic info) through python bindings. We already have it for the GPU pipeline options that can be used as a template for future bindings: https://github.com/iree-org/iree/pull/18840.
Make it easier to identify 'root ops'. We can make the IREE compiler annotate the root lingalg ops with a new attribute that the tuner can use to recognize them, without having to duplicate the compiler logic.
The Configuration representation is modeled after the requirements of the LLVMGPUVectorDistribute pipeline. This made it so that the surrounding code makes implicit assumptions about the problem representation. Instead, we should define an interface that allows us to support multiple compilation pipelines, such that the generated SMT constraints are specific to both the pipeline and the dispatch kind. Further, the constraint generation code should be decoupled from the parsing/printing code, such that projects like TKW can use just the constraint generation and benchmarking infra.
Move from two stages of compile-and-benchmark to just one. This made sense for SDXL where the best isolated dispatch does not necessarily perform best across the whole model, but it may not be necessary or even sufficiently general for other applications. This is related to the libtuner.TuningClient class; clients should be able to define their own tuning stages with libtuner providing the interface to specify the compilation and benchmarking commands.

Tasks

[x] Add an iree-opt pass to strip configuration from executable sources (incl. executable benchmarks) @bangtianliu
- https://github.com/iree-org/iree/pull/19069
[x] Expose key attributes via python bindings. @kuhar
[ ] Add a utility function to query supported MMA intrinsics and expose it to C API and python @bangtianliu
- https://github.com/iree-org/iree/pull/19124
[ ] Use MLIR types for types in the tuner @nithinsubbiah
[ ] Use IREE bindings for compilation info (incl., lowering_config and translation_info) (on a new branch) @kuhar
[ ] Make the tuner use python bindings to parse and print attributes
[ ] Update the tuner to generate candidate dispatches using the new iree-opt pass and transform dialect specs. @Max191
[ ] Modify IREE to annotate root ops with a new unit attribute.
[ ] Update the tuner to identify root ups using the new unit attribute produced by IREE.
[ ] Add support for the TileAndFuse pipeline. @bangtianliu and @Max191
[x] Move constraint generation logic out of the parsing/printing logic in candidate_gen.py. @kuhar
- https://github.com/nod-ai/SHARK-Platform/pull/531
- https://github.com/nod-ai/SHARK-Platform/pull/539
[ ] Use only one compile-benchmark stage in TuningCandidate. Update the existing example to adapt to this change.

kuhar commented 1 week ago

cc: @MaheshRavishankar @nithinsubbiah @bangtianliu @Max191

MaheshRavishankar commented 4 days ago

Thanks @kuhar . Could you also add a pointer to where the tuner lives today.

The tasks you described are related to making it easy to maintain the tuner. One more thing we need to think about is how to make it easy to maintain and update the tuner script. Specifically questions to answer are

1) We currently use the result of tuner (the TD script) during compilation. Do we maintain a default script that is somewhere in IREE codebase that gets loaded. 2) How does a developer update the tuner script if they need to (i.e they change the lowering config/translation info could there be a tool that "ports" the configuration over) 3) Could we have multiple TD scripts passed to the compiler. My answer to this would be no, but it might be useful to have a tool to "merge" TD scripts. This could also be the mechanism to update the tuner script. You create a new tuner script and "update".

For now just noting some top-of-mind questions.

kuhar commented 4 days ago

@MaheshRavishankar The tuner lives under the tuner directory in this repo. The questions you listed are very relevant but not to the tuner itself but the iree project -- I decided to decouple these two sets of concerns. I will answer your questions under an iree issue and add a link here.

nod-ai / shark-ai

[tuner] Reduce maintenance burden and prepare for more codegen pipelines #453

Tasks