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

[kuhar]

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:

1. 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.

2. 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: Reland #18804 iree-org/iree#18840.

3. 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.

4. 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.

5. 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

• Add an iree-opt pass to strip configuration from executable sources (incl. executable benchmarks) @bangtianliu
  • [tuner] add an iree-opt pass to strip configuration from executable sources iree-org/iree#19069
• Expose key attributes via python bindings. @kuhar
  • iree_gpu Python bindings (GPUPipelineOptionsAttr) iree-org/iree#18804
  • [python][tuner] Add bindings for MMAIntrinsic iree-org/iree#19095
  • [python][tuner] Add bindings for lowering config iree-org/iree#19096
  • [python] Simplify iree_gpu dialect bindings tests. NFC. iree-org/iree#19104
  • [Codegen] Update translation_info attribute assembly format. NFC. iree-org/iree#19107
  • [python][tuner] Set up bindings for iree_codegen iree-org/iree#19108
  • [python][tuner] Add bindings for iree_codegen.translation_info iree-org/iree#19128
  • [python][tuner] Add bindings for iree_codegen.compilation_info iree-org/iree#19129
• Add a utility function to query supported MMA intrinsics and expose it to C API and python @bangtianliu
  • [tuner]: Add a utility function to query supported MMA intrinsics iree-org/iree#19124
• Use MLIR types for types in the tuner @kuhar
  • [tuner] Use ir.(Integer|Float)Type for element types #554
• Use IREE attributes for MFMA intrinsics 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
• Move constraint generation logic out of the parsing/printing logic in candidate_gen.py. @kuhar
  • [tuner] Clean up candidate generation code #508
  • [tuner] Fix typing issues in libtuner. NFC. #526
  • [tuner][CI] Enable mypy type checking #530
  • [tuner] Make candidate_gen more modular. NFC. #531
  • [tuner] Move constraint generation out of canddiate_gen. NFC. #539
• Use only one compile-benchmark stage in TuningCandidate. Update the existing example to adapt to this change.

[kuhar]

cc: @MaheshRavishankar @nithinsubbiah @bangtianliu @Max191

[MaheshRavishankar]

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]

@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.