Wave: Python Domain-Specific Language for High Performance Machine Learning

Motivation

Wave is a high-level programming language designed to accelerate the development and optimization of machine learning kernels. It aims to dramatically improve kernel author velocity in two critical dimensions:

Implementation Velocity

Wave enables rapid prototyping of new optimization ideas and algorithms through its high-level abstractions and symbolic programming model. Kernel authors can quickly express complex tensor operations and experiment with different optimization strategies without getting bogged down in low-level implementation details.

Performance Velocity

The language is designed to achieve high performance with minimal tuning effort. Through its declarative constraints and hardware-aware abstractions, Wave automatically generates optimized GPU code, allowing kernel authors to focus on algorithmic innovation rather than manual optimization.

Wave is particularly focused on the machine learning domain, with deep integration into the PyTorch ecosystem. This integration enables:

1. Seamless transition from PyTorch models to optimized GPU kernels
2. Easy experimentation with new ML algorithms and optimizations
3. Rapid deployment of custom kernels in production ML pipelines
4. Direct reuse of PyTorch’s tensor abstractions and operator semantics

The language bridges the gap between high-level ML programming and low-level GPU performance by providing a flexible and expressive programming model that maintains close control over hardware resources while enabling rapid innovation in the ML space.

Core Design Principles

Wave is built around several key design principles that guide its architecture and implementation:

• Symbolic Programming Model
• Heavy use of symbolic variables to represent tensor dimensions
• Symbolic expressions for memory access patterns and layouts
• Enables compile-time optimization and analysis
• Provides flexibility in expressing complex tensor operations
• Separation of Distribution and Computation
• Clear separation between distribution strategy and computation graph
• Distribution strategy defined through declarative constraints
• Computation graph expressed independently of parallelization
• Enables better optimization and code reuse
• Dual Programming Models
• Support for both tile-based and SIMT-based programming models
• Tile-based model for coarse-grained parallelism
• SIMT model for fine-grained vector operations
• Flexible mapping to different GPU architectures
• Hardware-Aware Abstractions
• Direct mapping to GPU hardware concepts (workgroups, waves, etc.)
• Explicit control over memory hierarchy
• Hardware-specific optimizations (MMA operations, memory layouts)
• Performance portability across different GPU architectures

For more information, visit the documentation here: https://wave-lang.readthedocs.io/en/latest/wave/wave.html

Quickstart

Wave supports both end users (who want to use Wave kernels in their ML workflows) and developers (who want to contribute to or extend Wave). Follow the relevant guide below to get started quickly.

Prequisites

Before installing Wave, ensure you have the following prerequisites:

• Python 3.10 or later
• PyTorch
• ROCm (for AMD GPU support)
• A compatible AMD GPU with ROCm support (MI250, MI300, etc.) or CPU

For Users

Install ROCm PyTorch Dependencies

Before installing Wave, ensure you have the appropriate ROCm-enabled PyTorch dependencies:

pip install -r pytorch-rocm-requirements.txt

Install Wave

You can then install Wave and its dependencies using pip:

pip install wave-lang

Try the Example Notebook

To get started, try the examples/jupyter/wave_gemm_example.ipynb notebook.

For Developers

Clone the Repository

git clone https://github.com/iree-org/wave.git
cd wave

Install Development Dependencies

If you don’t have Rust already installed on your system, install Rust using the following command:

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

Next, install the python dependencies. It is recommended to use a virtual environment.

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements-iree-pinned.txt
pip install -r pytorch-rocm-requirements.txt
pip install -e ".[dev]"

Note: If you do not have access to Instinct GPUs, you can still install Wave as above but with the CPU version of PyTorch:

pip install -r pytorch-cpu-requirements.txt

Currently, you can only run lit tests in this mode.

Run Tests

To verify your setup:

# Python tests
pytest -s tests/
# To run end to end tests, add --run-e2e flag

# Lit tests
lit lit_tests/ -v

Build Configuration Options

Wave’s build process can be customized using the following environment variables:

WAVE_BUILD_TYPE: Build type for CMake (default: Release)

export WAVE_BUILD_TYPE=Debug

WAVE_BUILD_WATER: Build the Water MLIR dialect and infrastructure (default: 0)

export WAVE_BUILD_WATER=1

WAVE_LLVM_DIR: Path to pre-built LLVM installation (optional)

export WAVE_LLVM_DIR=/path/to/llvm-install

If not set, LLVM will be built from source during installation.

WAVE_LLVM_REPO: Custom LLVM repository URL (default: https://github.com/llvm/llvm-project.git)

export WAVE_LLVM_REPO=https://github.com/custom/llvm-fork.git

WAVE_LLVM_BUILD_SHARED_LIBS: Build LLVM as shared libraries (default: OFF)

export WAVE_LLVM_BUILD_SHARED_LIBS=ON

WAVE_WATER__DIR: Path to an existing Water CMake build directory for symlink-based installation (optional)

export WAVE_WATER_DIR=/path/to/wave/water/build

When set, the installation will create symlinks to the build directory instead of building Water from source. This enables a fast C++ development workflow where changes to Water are immediately available after running ninja without re-running pip install. Requires the build directory to have been configured with -DWATER_ENABLE_PYTHON=ON -DMLIR_ENABLE_BINDINGS_PYTHON=ON. 1.

Important Environment Flags

• WAVE_CHECK_LEAKS: Check whether Wave correctly deallocates tensors it owns (default: 0).

Contributing

Please see CONTRIBUTING.md for guidelines on contributing to Wave.