| # Custom Dispatches |
| |
| "Dispatches" in IREE are parallelized device function calls where "device" may |
| be a CPU task system, a GPU, or a dedicated accelerator. Parallelism is a |
| first-class concept in this model and plugging custom dispatches into IREE |
| requires reasoning about these function calls as if they were parallelized and |
| dispatched across a 3D grid (as on GPUs or CPU task systems). Note that a |
| degenerate case of grid dispatch is a grid size of 1x1x1 which turns the |
| dispatches into simple (if inefficient) device-side function calls. |
| |
| In normal workflows the IREE compiler forms the dispatch functions by way of |
| fusion and then uses code generation ("codegen") to translate the dispatch |
| functions into backend-specific forms like PTX, SPIR-V, or LLVM IR. It's |
| possible to augment this translation by either bypassing these steps entirely |
| and providing the already translated representation of the dispatch function or |
| extending the code generation portion by calling out to external functions from |
| within the generated dispatch (sometimes "microkernels" or "device libraries"). |
| |
| There's ongoing work across the core IREE compiler and specific backends to |
| enable more extension points and ways to connect to them from frontends or |
| compiler transforms. This current set of samples demonstrates a very early |
| version of this extensibility that can be used to tunnel through dispatch |
| workloads by bypassing code generation (in the case of PTX/SPIR-V) or coarsely |
| interoperating (CPU function calls). In its current form it is intended for |
| things that would traditionally be custom ops in ML frameworks and produces much |
| smaller, hermetic, retargetable, and optimizable programs than traditional |
| custom ops can as there's nearly zero performance delta between what the |
| compiler can dispatch and what the user decides to dispatch. |
| |
| ## Approaches |
| |
| In the fullness of time all backends will support all approaches but currently |
| there are limitations and these samples only cover the supported cases: |
| |
| | | CPU | CUDA | Metal | SPIR-V | WGSL | |
| |--------------------|:------------------:|:------------------:|:-----:|:------------------:|:---------------:| |
| | Static Functions | :white_check_mark: | TBD | TBD | TBD | TBD | |
| | Dynamic Functions | :white_check_mark: | TBD | TBD | :grey_question: | :grey_question: | |
| | Static Dispatches | TBD | :white_check_mark: | TBD | :white_check_mark: | TBD | |
| | Dynamic Dispatches | TBD | TBD | TBD | TBD | TBD | |
| | Commands | TBD | TBD | TBD | TBD | TBD | |
| |
| ### Statically-linked Function Calls |
| |
| **Overview**: user defines functions in .c files, compiles them with specific |
| settings to .o files, emits calls to the functions in IR interleaved with other |
| IR, and lets the IREE compiler link the objects into the final binary. |
| |
| **Workflow**: |
| |
| ``` |
| +-------------+ +---------------------+ +--------------+ |
| | functions.c | -> clang -+-> | functions_aarch64.o | -+ | example.mlir | |
| +-------------+ | +---------------------+ | +--------------+ |
| | +---------------------+ | v |
| +-> | functions_x86_64.o | -+----> iree-compile |
| +---------------------+ v |
| +--------------+ |
| hermetic | example.vmfb | |
| +--------------+ |
| ``` |
| |
| **Samples**: |
| |
| * CPU: [custom_dispatch/cpu/embedded/](./cpu/embedded/) (.c -> .o) |
| |
| This approach is usable both from frontends as well as something that can be |
| synthesized by compiler transforms ("replace op X with call to extern fX and |
| link in object f.o") and the object files referenced can be generated on the |
| fly and embedded into the IR to allow for compile-time specialization. |
| |
| This is the preferred method of extension as it preserves IREE's ability to |
| specialize executables, optimize aggressively across call boundaries, |
| hermetically deploy the custom code without runtime changes, and portably target |
| multiple platforms and architectures. |
| |
| ### Dynamically-linked Function Calls |
| |
| **Overview**: user defines functions in any source language with a compatible C |
| ABI, wires them up and links them in their runtime binary, declares the |
| externally available functions in IR, and emits calls to the functions in IR |
| interleaved with other IR. |
| |
| **Workflows**: |
| |
| Statically-linked into the hosting runtime: |
| ``` |
| +--------------+ +--------------+ |
| | example.mlir | | runtime srcs | -+ |
| +--------------+ +--------------+ | |
| +--------------+ v +--------------+ | |
| | declarations | ----> iree-compile | imports.c | -+-> custom runtime |
| +--------------+ v +--------------+ ^ |
| +--------------+ | |
| | example.vmfb | - - - - - - - - - - - - - - - + |
| +--------------+ |
| ``` |
| |
| Dynamically-linked into the hosting runtime via system libraries: |
| ``` |
| +----------+ +---------------+ +--------------+ |
| | plugin.c | -> | plugin.so/dll |-+ | example.mlir | |
| +----------+ +---------------+ | +--------------+ |
| | v |
| | iree-compile |
| | v |
| | +--------------+ |
| | | example.vmfb | (non-hermetic) |
| | +--------------+ |
| | | |
| +-----+-----+ |
| v |
| +-----------------+ |
| | iree-run-module | |
| +-----------------+ |
| ``` |
| |
| Dynamically-linked into the hosting runtime via portable embedded ELFs: |
| |
| ``` |
| +----------+ +-------------------+ +--------------+ |
| | plugin.c | -+-> | plugin_aarch64.so | -+ | example.mlir | |
| +----------+ | +-------------------+ | +--------------+ |
| | +-------------------+ | v |
| +-> | plugin_x86_64.so | -+ iree-compile |
| +-------------------+ | v |
| +------------+ | +--------------+ |
| | plugin.sos | <--+ | example.vmfb | (non-hermetic) |
| +------------+ +--------------+ |
| | | |
| +----------+------------+ |
| v |
| +-----------------+ |
| | iree-run-module | |
| +-----------------+ |
| ``` |
| |
| **Samples**: |
| |
| * CPU (plugins): [custom_dispatch/cpu/plugin/](./cpu/plugin/) (.c -> .so/.sos) |
| |
| Unlike the other approaches this requires runtime device support for dynamic |
| linking and introduces complexity to the user as they must be careful to version |
| their input programs and their runtime libraries themselves. IREE's CPU backend |
| does provide basic support for optional imports such that users can emit their |
| calls and add fallbacks but otherwise such behavior falls on the user to |
| implement. |
| |
| ### Statically-linked Dispatch Functions |
| |
| **Overview**: user produces the final dispatch executable binary themselves, |
| declares it in IR, and dispatches it. |
| |
| **Workflow**: |
| |
| ``` |
| +------------+ +-------------------+ +--------------+ |
| | kernels.cu | -> nvcc -+-> | kernels_sm_52.ptx | -+ | example.mlir | |
| +------------+ | +-------------------+ | +--------------+ |
| | +-------------------+ | v |
| +-> | kernels_sm_80.ptx | -+----> iree-compile |
| +-------------------+ v |
| +--------------+ |
| hermetic | example.vmfb | |
| +--------------+ |
| ``` |
| |
| **Samples**: |
| |
| * CUDA/PTX: [custom_dispatch/cuda/kernels/](./cuda/kernels/) (.cu -> .ptx) |
| * Vulkan/SPIR-V: [custom_dispatch/vulkan/shaders/](./vulkan/shaders/) (.glsl -> .spv) |
| |
| Here IREE is used for scheduling the work and ensuring that buffers and |
| parameters are passed to the dispatch function but otherwise treats them as |
| opaque. This disables some IREE optimizations but the functions are still able |
| to be scheduled concurrently and with parallelization. Many custom kernels can |
| often be implemented like this instead of needing much heavier-weight runtime |
| custom calls that prevent the asynchronous scheduling that IREE uses. |
| |
| The dispatch functions are embedded into the IREE compiler outputs such that no |
| runtime changes are required and compiled programs are hermetic. Multi-targeting |
| is enabled by allowing the user to provide binaries for the target devices and |
| architectures they are compiling for. |
| |
| ### Dynamically-linked Dispatch Functions |
| |
| **Overview**: user defines functions in a target-specific source language, |
| compiles them into target-specific libraries, wires them up and links them in |
| their runtime binary, declares the externally available dispatch executable in |
| IR, and emits calls to the functions in IR interleaved with other IR. |
| |
| **Workflow**: |
| |
| ``` |
| +--------------+ +--------------+ |
| | example.mlir | | runtime srcs | -+ |
| +--------------+ +--------------+ | |
| +--------------+ v +--------------+ | |
| | declarations | ----> iree-compile | dispatches.c | -+-> custom runtime |
| +--------------+ v +--------------+ ^ |
| +--------------+ | |
| | example.vmfb | - - - - - - - - - - - - - - - + |
| +--------------+ |
| ``` |
| |
| **Samples**: plumbing required. |
| |
| Similar to statically-linked dispatch functions IREE is doing the scheduling and |
| managing resources but deferring the entire dispatch logic to externally-sourced |
| executables. In contrast to the static linking approach this has the compiler |
| emit references to runtime-provided target-specific executables that must be |
| built into the runtime. This means that deployment gets more complicated as |
| compiler-produced outputs are no longer hermetic and users must handle |
| versioning and platform constraints themselves. |
| |
| ### Custom Commands |
| |
| **Overview**: user writes custom command buffer operations in VM modules, |
| declares them in IR, dispatches them, and then links their custom modules into |
| the runtime. |
| |
| **Workflow**: |
| |
| ``` |
| +--------------+ +--------------+ |
| | example.mlir | | runtime srcs | -+ |
| +--------------+ +--------------+ | |
| +--------------+ v +--------------+ | |
| | declarations | ----> iree-compile | module.c | -+-> custom runtime |
| +--------------+ v +--------------+ ^ |
| +--------------+ | |
| | example.vmfb | - - - - - - - - - - - - - - - + |
| +--------------+ |
| ``` |
| |
| **Samples**: plumbing required. |
| |
| Though more work than the other approaches this allows for a host-side call that |
| can produce new transfer or execution commands. The compiler effectively treats |
| these calls as if they were custom versions of `hal.command_buffer.dispatch`/ |
| `iree_hal_command_buffer_dispatch` and the runtime custom module receives the |
| device, command buffer, and push constants/bindings. Portable modules can |
| use the HAL APIs to schedule more commands (multiple dispatches, transfers, |
| collectives, etc) while backend-specific ones can crack open the HAL objects and |
| get the internals with all the normal API stability caveats. |
| |
| An example use case would be calling a CUDA library that takes a `CUstream` as |
| an argument. The user would define their custom module with a call that uses the |
| `iree/hal/drivers/cuda/api.h` to cast the command buffer to a stream (using |
| graph capture when the command buffer is constructing a graph), make the call, |
| and then return. This only works if the library is designed for asynchronous and |
| deferred execution - if the library makes large allocations, schedules blocking |
| operations, or has side-effecting behavior then full custom modules must be |
| used (see the [custom_module/async/](/samples/custom_module/async/) sample). |
| |
| When at all possible the dispatch function call or dispatch function |
| substitution approaches should be used instead and in many cases that is |
| sufficient for most workloads not involving other libraries. |
| |
| ### Compile Time Inlining Custom Function Calls |
| |
| **Overview**: user defines functions with MLIR dialects IREE is able to ingest |
| paired with a matcher and replacement pattern. The matcher runs as preprocessing |
| and calls into the replacement pattern for all successful matches. The |
| replacement pattern imports a function from the externally |
| ABI, wires them up and links them in their runtime binary, declares the |
| externally available functions in IR, and emits calls to the functions in IR |
| interleaved with other IR. |
| |
| **Workflows**: |
| |
| Statically matched and imported external functions |
| ``` |
| +--------------+ |
| | example.mlir | |
| +--------------+ +--------------+ +--------------+ |
| | (one of the | | functions + | v |
| | above static | ----> | matchers + | ----> iree-compile |
| | workflows) | | replace.mlir | v |
| +--------------+ +--------------+ +--------------+ |
| | example.vmfb | |
| +--------------+ |
| ``` |
| |
| **Samples**: |
| |
| * CPU: [custom_dispatch/cpu/embedded/](./cpu/embedded/) (.c -> .o) |
| * [custom_dispatch/cpu/embedded/](./cpu/embedded/example_transform_spec.mlir) (.mlir) |
| * Vulkan/SPIR-V: [custom_dispatch/vulkan/shaders/](./vulkan/shaders/) (.glsl -> .spv) |
| * [custom_dispatch/vulkan/shaders/](./vulkan/shaders/example_transform_spec.mlir) (.mlir) |
| |
| The above two samples build on top of a couple of the static workflows shown |
| above, but should work with any of the other approaches. The idea is to separate |
| the custom kernel from the target module to be compiled, allowing integration of |
| custom dispatches with default IREE codegen without the need to build a custom |
| set of compiler tools around IREE to generate the necessary IR. |
| |
| There are a number of possible points at which the match and replace can happen; |
| the above shows it after import + input conversion. Other plugin points are |
| possible (e.g. before input conversion or after global optimization), but |
| currently are missing some ergonomics on the available matchers. |
| |
| ### Others |
| |
| Most other situations are covered by [custom modules](/samples/custom_module/). |
| These can still be asynchronous and scheduled in device order but incur |
| additional overheads and deployment complexity as custom runtimes are required. |
