commit 5b8757b4b881ccd8fbddb5ad7e93ee45e45cebb3
author Ben Vanik <ben.vanik@gmail.com> Fri Feb 17 11:54:19 2023 -0800
committer GitHub <noreply@github.com> Fri Feb 17 11:54:19 2023 -0800
tree 3d18d552d4b55287c0db35ebe9fb3c3f770433cc
parent 425efcd856a6a14351071b9ce63dcb5bfa6c00f1
parent e231b396520e94fc9609ede834733027e35c0f65

Adding `--iree-hal-substitute-executable=` flag. (#12240)

This allows for specifying one or more `executable_name=file.xxx` pairs
that each replace a `hal.executable` op with the `executable_name` with
the contents of `file.xxx`. .mlir/.mlirbc files are loaded and a
`hal.executable` with the matching name is used as a replacement while
any other file type will cause the original executable to be
externalized and linked with the specified file (.ptx/.spv/etc).

The additional `--iree-hal-substitute-executable-*-from=` flags allow
for scanning a directory for executables by name to build the
substitution mapping. Only files in the executable_name or
module_executable_name form will be substituted but we can extend this
in the future to support variant naming.

Because of phase ordering constraints around where codegen is able to
mutate host-related code such as workgroup count calculations there are
two flag sets:
  `--iree-hal-substitute-executable-source=name=file.xxx`
  `--iree-hal-substitute-executable-sources-from=path/`
 and
  `--iree-hal-substitute-executable-object=name=file.xxx`
  `--iree-hal-substitute-executable-objects-from=path/`
Sources are substituted immediately prior to benchmark generation and
when substituting target objects (.ptx, .spv, etc) require that it's ok
to skip codegen (workgroup count calculation is not dependent on root op
detection, etc). Objects are substituted immediately after codegen for
use in cases where codegen is to generate the host code. There are uses
for both depending on what the input IR is and what the developer wants
to modify (host code, device code, or both).

The primary developer workflows this covers are:
1. dump executable sources via `--iree-hal-dump-executable-sources-to=`
and modify them, potentially running any number of iree-opt passes,
before linking them back in to the original program they came from
2. author custom implementations ala the custom_dispatch sample in
target toolchains (.cu -> .ptx, .glsl -> .spv, etc) and use those in
full programs without needing to modify the compiler
3. do either of the above and use the substituted executable for
microbenchmarking via `--iree-hal-dump-executable-benchmarks-to=` (so
one can easily microbenchmark handwritten kernels)

Example usage:
```sh
# dump sources for a program
iree-compile ... \
  --iree-hal-dump-executable-sources-to=~/sources/
# <modify some of the sources>
# recompile with the new changes and substitute 2 executables
iree-compile ... \
    --iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
    --iree-hal-substitute-executable-source=_main_dispatch_1=~/sources/modified_dispatch_1.mlir
# same thing with search paths
iree-compile ... \
    --iree-hal-executable-object-search-path=~/sources/ \
    --iree-hal-substitute-executable-source=_main_dispatch_0=modified_dispatch_0.mlir
# same thing but matching all files by name as from dump-sources-to:
iree-compile ... \
    --iree-hal-substitute-executable-sources-from=~/sources/
```

This works with ptx/spv as well:
```sh
# dump ptx binaries
iree-compile ... \
    --iree-hal-dump-executable-binaries-to=~/binaries/
# <modify dispatch ptx>
# replace @_main_dispatch_0 with the external ptx file
iree-compile ... \
    --iree-hal-substitute-executable-object=_main_dispatch_0=~/binaries/modified_dispatch_0.ptx
```

It's also possible to iterate on microbenchmarks using the custom
sources/objects:
```sh
# dump all benchmarks with the substitution active
iree-compile ... \
    --iree-hal-substitute-executable-source=_main_dispatch_0=~/sources/modified_dispatch_0.mlir \
    --iree-hal-dump-executable-benchmarks-to=~/benchmarks/
# inspect benchmark for the dispatch and see the hello.world attr
# can use iree-compile to build the benchmark and then iree-benchmark-module
```

Progress on #12222 (the rest for linking alternative formats is
orthogonal).