How gwokhiujin/avelang works
1. Overall shape
Ave is a mixed Python/C++ GPU-kernel compiler stack. The Python package avelang is the user-facing layer, while the C++ code implements parsing, IR construction, lowering passes, and backend code generation through an MLIR-based pipeline.
2. User-facing entry point
Kernel authors import avelang from Python and write “AveLang” kernels in a Pythonic style. The package exposes language constructs, runtime helpers, backend selection, specialization logic, and compilation/JIT entry points. This is the top boundary of the system.
3. Frontend and parsing
The C++ frontend in lib/Frontend parses AveLang source into an AST. lib/AST contains the language tree, visitors, context, and casting utilities. Diagnostics live in lib/Basic, so parse/type errors are reported before IR generation.
4. Language and type model
lib/IR/type_system.cc, type_promotion.cc, symbol_table.cc, symbol_scope.cc, and named_module.cc define the semantic layer that sits between parsing and MLIR generation. This subsystem resolves names, manages scopes, enforces tensor/layout rules, and handles type promotion and symbol lookup.
5. MLIR generation core
lib/IR/mlir_generator.cc and related files translate the AveLang AST and semantic model into MLIR. Supporting pieces such as expr_generator.cc, function_generator.cc, layout_operation.cc, constant_folder.cc, and mangler.cc shape the generated IR, preserve naming/symbol identity, and fold compile-time values.
6. Built-in modules and intrinsics
The IR layer also contains embedded built-in dialect content and intrinsic libraries. lib/IR/builtin_module.cc and lib/IR/Intrinsics/* provide AMDGPU and NVVM intrinsic definitions, signatures, and support code. These are the bridge from high-level kernel code to backend-specific GPU operations.
7. Custom AveLang dialect
Under lib/Dialect/AveLang, the project defines an MLIR dialect for AveLang-specific ops, types, and transforms. TableGen files (*.td) plus generated/handwritten dialect code establish the intermediate representation used before lowering to more standard MLIR or LLVM forms.
8. Lowering and normalization passes
lib/Dialect/AveLang/Transforms contains the first lowering stages: return normalization, stack allocation hoisting, lowering AveLang ops to memref, and lowering GPU ops to intrinsics. These passes are the main boundary where the high-level dialect becomes backend-oriented IR.
9. GPU-target pipeline
lib/Target/GPU provides target-agnostic GPU lowering support: outlining, linking intrinsic implementations, and lowering to LLVM. This layer is shared by backends and acts as the common GPU compilation backbone before vendor-specific codegen.
10. Vendor backends
There are two concrete GPU backends:
- lib/Target/NVVM for NVIDIA / CUDA-style codegen
- lib/Target/AMDGPU for AMD / ROCm / HIP-style codegen
Each backend has its own pipeline, backend driver, and tests, with AMDGPU additionally including ROCm installation detection and math/legalization passes.
11. Python compilation and runtime packages
The Python code is split into logical layers:
- python/avelang/language: language surface syntax/constructs
- python/avelang/compiler: code generation and compiler orchestration
- python/avelang/runtime: build, driver, JIT, and specialization
- python/avelang/backends: backend-facing compiler/driver wrappers for AMDGPU and NVIDIA
This is the orchestration layer that invokes the native compiler and runtime machinery.
12. Native Python bindings
python/src/avelang_bindings.cc provides the compiled extension that connects Python to the C++ implementation. This is the key interop boundary: Python constructs and runtime helpers call into native compilation and execution facilities rather than implementing the compiler in pure Python.
13. Tests and validation
The test/ tree covers language primitives, memory/layout behavior, JIT behavior, and end-to-end examples for GEMM/GEMV on both AMDGPU and NVIDIA. C++ tests also exist beside the IR and backend code. This indicates the repo validates both compiler correctness and target-specific code generation.
14. Docs, examples, and tooling
The docs/ site describes installation, tutorials, and language reference. benchmark/ and python/avelang_kernels/ provide performance and example kernels. Utility scripts in tools/ dump LLVM IR or assembly, which are useful for inspecting the compiler pipeline output.
15. Build and packaging
The project uses CMake for the native code, plus Python packaging via pyproject.toml and setup.py. The repository is therefore a hybrid source tree: native libraries and targets are built with CMake, while the Python package is installed as ave-lang and exposes the end-user API as avelang.
Main entry points: mlir_generator.cc MLIR Gen · GPU GPU Core · avelang_bindings.cc Native Bindings · backends Backend Wrappers
Components 20
Python stack
User-facing language, compiler orchestration, runtime, and backend wrappers.
- Language API · Pythonic AveLang syntax and built-in language constructs.
python/avelang/language - Compiler API · Python-side code generation and compiler entry points.
python/avelang/compiler - Runtime · Build, driver, JIT, and specialization helpers.
python/avelang/runtime - Backend Wrappers · Python wrappers for AMDGPU and NVIDIA compilation/driver flows.
python/avelang/backends - Native Bindings · Compiled extension connecting Python to the C++ implementation.
python/src/avelang_bindings.cc
C++ core
Frontend, IR, MLIR generation, dialects, and shared utilities.
- Frontend · Parses AveLang source into an AST.
lib/Frontend - AST · Language tree, visitors, casting, and AST context.
lib/AST - Diagnostics · Compile-time diagnostics and reporting.
lib/Basic - Type Model · Type promotion, symbols, scopes, and semantic module state.
lib/IR/type_system.cc - MLIR Gen · Translates AST and semantic state into MLIR.
lib/IR/mlir_generator.cc - IR Helpers · Expression/function generation, layout ops, mangling, folding, and parsing utilities.
lib/IR - Builtins · Embedded builtin modules plus AMDGPU and NVVM intrinsic definitions.
lib/IR/Intrinsics - AveLang Dialect · Custom AveLang ops, types, and dialect definitions.
lib/Dialect/AveLang/IR - Lowering Passes · Normalizes returns, hoists allocations, and lowers AveLang/GPU ops.
lib/Dialect/AveLang/Transforms
GPU targets
Shared GPU lowering plus NVIDIA and AMDGPU backend pipelines.
- GPU Core · Shared GPU outlining, intrinsic linking, and lowering to LLVM.
lib/Target/GPU - NVVM Backend · NVIDIA-specific GPU-to-NVVM pipeline and codegen.
lib/Target/NVVM - AMDGPU Backend · AMDGPU pipeline, ROCm detection, legalization, and codegen.
lib/Target/AMDGPU
Docs and tests
Documentation, examples, benchmarks, and validation coverage.
- Tests · End-to-end examples and primitive behavior coverage across targets.
test - Docs Site · Installation, tutorials, language reference, and developer notes.
docs - Benchmarks · Performance-oriented sample kernels and benchmark entry points.
benchmark
Connections 23
- Language API → Compiler API: kernel code · Language constructs feed compilation orchestration.
- Compiler API → Native Bindings: invoke · Python compilation calls into native bindings.
- Runtime → Native Bindings: execute · Runtime and JIT paths cross the native boundary.
- Backend Wrappers → Native Bindings: wrap · Backend wrappers delegate to the extension module.
- Native Bindings → Frontend: parse · Native layer enters the C++ compilation pipeline.
- Frontend → AST: build · Parser produces the AveLang AST.
- Frontend → Diagnostics: report · Parsing failures and diagnostics are emitted through the basic error layer.
- AST → Type Model: analyze · AST nodes are resolved into semantic types and scopes.
- Type Model → MLIR Gen: lower · Semantic resolution feeds MLIR generation.
- MLIR Gen → IR Helpers: use · Core generation relies on helper modules for expressions, functions, layouts, and folding.
- MLIR Gen → Builtins: inline · Generated IR incorporates builtin and GPU intrinsic modules.
- MLIR Gen → AveLang Dialect: emit · High-level AveLang semantics are materialized as a custom MLIR dialect.
- AveLang Dialect → Lowering Passes: lower · Dialect ops are normalized and lowered by AveLang transforms.
- Lowering Passes → GPU Core: handoff · Lowered GPU-oriented IR enters the shared GPU backend pipeline.
- Builtins → GPU Core: bridge · Intrinsic implementations connect high-level kernels to backend GPU ops.
- GPU Core → NVVM Backend: target · Shared GPU lowering feeds the NVIDIA backend.
- GPU Core → AMDGPU Backend: target · Shared GPU lowering feeds the AMD backend.
- Backend Wrappers → NVVM Backend: select · Python backend selection chooses the NVIDIA path.
- Backend Wrappers → AMDGPU Backend: select · Python backend selection chooses the AMDGPU path.
- NVVM Backend → Tests: verify · NVIDIA backend behavior is covered by target tests.
- AMDGPU Backend → Tests: verify · AMDGPU backend behavior is covered by target tests.
- Docs Site → Language API: describe · Documentation explains language constructs and usage.
- Benchmarks → Tests: exercise · Benchmarks and sample kernels validate performance-sensitive paths.