| commit | 76aaedc3c9b10241b4f575235ff9660d870d0424 | [log] [tgz] |
|---|---|---|
| author | Ben Vanik <ben.vanik@gmail.com> | Thu Mar 23 08:57:15 2023 -0700 |
| committer | GitHub <noreply@github.com> | Thu Mar 23 08:57:15 2023 -0700 |
| tree | 5b106e5575254c88694a0fedeedd3ccab69e3a19 | |
| parent | 5eb892ac3ce323d3db5b1dc1724d57ba61f8d28d [diff] |
Reworking CUDA channel creation and plumbing group/ID. (#12695) It's now possible to inject a function that gets called prior to each channel created. This function can populate default values if needed or do ID initialization and exchange. To make it possible for the function to differentiate channel groups a new `group` field is exposed in the compiler and runtime. This is user-defined and defaults to nothing today but could be used to identify unique subsets of participants within a context. This fixes some issues with the existing implementation such as always bootstrapping a NCCL root and relying on magic environment variables in the core implementation. Hosting layers can now completely control how ID exchange happens. The PJRT plugin, for example, can have its own provider when it creates its CUDA device that interops with that stack. See `iree_hal_cuda_nccl_query_group_params` for an example approximating the existing behavior. In the future we may want to move channel creation into the provider such that we never create channels ourselves. This would allow NCCL (or other implementations) to be externalized, though there are several APIs we'd need to support the direct access to buffer resources and streams/graphs.
IREE (Intermediate Representation Execution Environment, pronounced as “eerie”) is an MLIR-based end-to-end compiler and runtime that lowers Machine Learning (ML) models to a unified IR that scales up to meet the needs of the datacenter and down to satisfy the constraints and special considerations of mobile and edge deployments.
See our website for project details, user guides, and instructions on building from source.
IREE is still in its early phase. We have settled down on the overarching infrastructure and are actively improving various software components as well as project logistics. It is still quite far from ready for everyday use and is made available without any support at the moment. With that said, we welcome any kind of feedback on any communication channels!
See our website for more information.
IREE is licensed under the terms of the Apache 2.0 License with LLVM Exceptions. See LICENSE for more information.