Package: shark-turbine
SHARK TurbineTurbine is the set of development tools that theSHARK Teamis building for deploying all of our models for deployment to the cloud and devices. We
are building it as we transition from our TorchScript-era 1-off export and compilation
to a unified approach based on PyTorch 2 and Dynamo. While we use it heavily ourselves, it
is intended to be a general purpose model compilation and execution tool.Turbine provides three primary tools:AOT Export: For compiling one or morenn.Modules to compiled, deployment
ready artifacts. This operates via both asimple one-shot export APIfor simple models and an underlyingadvanced APIfor complicated models
and accessing the full features of the runtime.Eager Execution: Atorch.compilebackend is provided and a Turbine Tensor/Device
is available for more native, interactive use within a PyTorch session.Turbine Kernels: (coming soon) A union of theTritonapproach andPallasbut based on
native PyTorch constructs and tracing. It is intended to complement for simple
cases where direct emission to the underlying, cross platform, vector programming model
is desirable.Under the covers, Turbine is based heavily onIREEandtorch-mlirand we use it to drive evolution
of both, upstreaming infrastructure as it becomes timely to do so.Seethe roadmapfor upcoming work and places to contribute.Contact UsTurbine is under active development. If you would like to participate as it comes online,
please reach out to us on the#turbinechannel of thenod-ai Discord server.Quick Start for UsersInstall from source:pip install shark-turbine
# Or editable: pip install -e .The above does install some unecessary cuda/cudnn packages for cpu use. To avoid this you
can specify pytorch-cpu and install via:pip install --index-url https://download.pytorch.org/whl/cpu \
    -r pytorch-cpu-requirements.txt \
    -r torchvision-requirements.txt
pip install shark-turbine(or follow the "Developers" instructions below for installing from head/nightly)Try one of the samples:Generally, we use Turbine to produce valid, dynamic shaped Torch IR (from thetorch-mlir torchdialectwith various approaches to handling globals). Depending on the use-case and status of the
compiler, these should be compilable via IREE with--iree-input-type=torchfor
end to end execution. Dynamic shape support in torch-mlir is a work in progress,
and not everything works at head with release binaries at present.AOT MLP With Static ShapesAOT MLP with a dynamic batch sizeAOT llama2:
Dynamic sequence length custom compiled module with state management internal to the model.Eager MNIST withtorch.compileDevelopersGetting Up and RunningIf only looking to develop against this project, then you need to install Python
deps for the following:PyTorchiree-compiler (with Torch input support)iree-runtimeThe pinned deps at HEAD require pre-release versions of all of the above, and
therefore require additional pip flags to install. Therefore, to satisfy
development, we provide arequirements.txtfile which installs precise
versions and has all flags. This can be installed prior to the package:Installing into a venv is highly recommended.pip install --upgrade -r requirements.txt
pip install --upgrade -e .[torch-cpu-nightly,testing]Run tests:pytestUsing a development compilerIf doing native development of the compiler, it can be useful to switch to
source builds for iree-compiler and iree-runtime.In order to do this, check outIREEand
follow the instructions tobuild from source, making
sure to specifyadditional options:-DIREE_BUILD_PYTHON_BINDINGS=ON -DPython3_EXECUTABLE="$(which python)"Configuring PythonUninstall existing packages:pip uninstall iree-compiler
pip uninstall iree-runtimeCopy the.envfile fromiree/to this source directory to get IDE
support and add to your path for use from your shell:source .env && export PYTHONPATH