SwarmLang - A Triton-like Pythonic GPU Kernel Language (Learning Project)

A learning project to understand how languages like Triton work under the hood — from Python decorators all the way down to GPU machine code via LLVM.

Project Structure

.
├── src/
│   ├── frontend/          # Python DSL → AST
│   │   ├── __init__.py
│   │   ├── decorators.py  # @swarm.kernel decorator
│   │   ├── ast_nodes.py   # Custom AST node definitions
│   │   └── parser.py      # Python AST → SwarmLang AST
│   ├── codegen/           # AST → LLVM IR
│   │   ├── __init__.py
│   │   ├── ir_builder.py  # LLVM IR generation using llvmlite
│   │   └── types.py       # Type system (tensor types, scalar types)
│   ├── backend/           # LLVM IR → PTX/CUDA
│   │   ├── __init__.py
│   │   ├── ptx_compile.py # LLVM IR → PTX via NVPTX backend
│   │   └── cuda_launch.py # Launch PTX kernels on GPU
│   └── runtime/           # Runtime helpers
│       ├── __init__.py
│       ├── memory.py      # GPU memory management
│       └── grid.py        # Grid/block dimension helpers
├── examples/              # Example kernels
│   ├── vector_add.py
│   ├── softmax.py
│   └── matmul.py
├── tests/                 # Tests
│   ├── test_frontend.py
│   ├── test_codegen.py
│   └── test_e2e.py
├── requirements.txt
└── README.md

The Pipeline (How It All Connects)

  Python Code          Your Frontend        LLVM            GPU
  ──────────          ────────────         ────            ───
  @swarm.kernel  ───►  Parse Python  ───►  LLVM IR  ───►  PTX  ───►  Run on CUDA
  def add(...)         AST to our          (llvmlite)      (NVPTX     (pycuda /
                       custom AST                          backend)   cuda-python)

Stage 1: Frontend (Python → AST)

• User writes a kernel using @swarm.kernel decorator
• When the decorator fires, we inspect the function using inspect.getsource()
• We parse it with Python's ast module into a Python AST
• We walk the Python AST and convert it to our own SwarmLang AST nodes

Stage 2: Codegen (AST → LLVM IR)

• Walk our custom AST and emit LLVM IR using llvmlite
• Map tensor operations to LLVM vector ops or memory load/store patterns
• Insert GPU-specific intrinsics (thread_id, block_id, etc.)
• The output is an LLVM IR module (textual .ll or in-memory)

Stage 3: Backend (LLVM IR → PTX → GPU)

• Use LLVM's NVPTX backend (via llvmlite) to compile IR to PTX assembly
• PTX is NVIDIA's virtual ISA — it's like assembly for NVIDIA GPUs
• Load the PTX into the CUDA driver using pycuda or cuda-python
• Launch the kernel with grid/block dimensions

Prerequisites

Must Install

# Python 3.10+
pip install llvmlite          # LLVM bindings for Python
pip install pycuda            # CUDA driver API from Python (needs CUDA toolkit)
pip install numpy             # For host-side arrays

Must Have on System

• CUDA Toolkit (nvcc, libcuda) — https://developer.nvidia.com/cuda-downloads
• NVIDIA GPU (or use Google Colab for free GPU access)
• LLVM 14+ (llvmlite bundles its own, but understanding helps)

Optional but Helpful

pip install cuda-python       # Alternative to pycuda (NVIDIA's official binding)
pip install triton            # Install actual Triton to study/compare

How to Run

Step 1: Verify your setup

python -c "import llvmlite; print(llvmlite.__version__)"
python -c "import pycuda.autoinit; print('CUDA OK')"

Step 2: Run an example

python examples/vector_add.py

Step 3: Run tests

python -m pytest tests/ -v

Step 4: Inspect generated LLVM IR

python -c "from examples.vector_add import inspect_ir; inspect_ir()"
# This dumps the LLVM IR so you can read it

Step 5: Inspect generated PTX

python -c "from examples.vector_add import inspect_ptx; inspect_ptx()"
# This dumps the PTX assembly

Learning Workflow — What To Do

Phase 1: Understand the Frontend

1. Read src/frontend/decorators.py — understand how @swarm.kernel captures code
2. Read src/frontend/ast_nodes.py — understand the custom AST
3. Read src/frontend/parser.py — understand Python AST → custom AST conversion
4. Exercise: Add support for if/else in the parser
5. Exercise: Add support for while loops

Phase 2: Understand LLVM IR

1. Read the LLVM Language Reference (skim it)
2. Read src/codegen/ir_builder.py — see how AST becomes LLVM IR
3. Run inspect_ir() on examples and read the output
4. Exercise: Write LLVM IR by hand for a simple add function
5. Exercise: Add a new binary op (e.g., modulo) to the codegen

Phase 3: Understand the GPU Backend

1. Read NVIDIA's PTX ISA docs
2. Read src/backend/ptx_compile.py — see how IR becomes PTX
3. Read src/backend/cuda_launch.py — see how PTX runs on GPU
4. Exercise: Change grid dimensions and observe performance
5. Exercise: Add shared memory support

Phase 4: Build Something Real

1. Implement a working softmax kernel
2. Implement a basic matmul kernel
3. Compare performance with numpy and actual Triton
4. Exercise: Add basic autotuning (try different block sizes)

What NOT To Do

• Don't try to build a production compiler — this is for learning
• Don't skip reading the generated LLVM IR — that's where the learning is
• Don't ignore error messages from LLVM/CUDA — they teach you about constraints
• Don't start with complex kernels — get vector_add working first
• Don't worry about performance initially — correctness first
• Don't skip the exercises — modifying code teaches more than reading it
• Don't try to support all Python syntax — start with arithmetic + loads/stores

Key Concepts to Understand

Concept	What It Is	Where In Code
program_id	Like Triton's tl.program_id() — which block am I?	src/runtime/grid.py
BLOCK_SIZE	How many elements each GPU thread block processes	examples/vector_add.py
Thread hierarchy	Grid → Blocks → Threads	src/backend/cuda_launch.py
LLVM IR	Text representation of your kernel in LLVM's language	src/codegen/ir_builder.py
PTX	NVIDIA's assembly language for GPUs	src/backend/ptx_compile.py
Shared memory	Fast on-chip memory shared within a block	Future exercise

Useful Commands for Debugging

# Dump LLVM IR to file
python -c "from src.codegen.ir_builder import dump_ir; dump_ir()" > kernel.ll

# Read the IR (it's human-readable text)
cat kernel.ll

# If you have llc installed, compile IR to PTX manually
llc -march=nvptx64 -mcpu=sm_80 kernel.ll -o kernel.ptx

# Read the PTX
cat kernel.ptx

# Check CUDA device info
nvidia-smi
python -c "import pycuda.autoinit; import pycuda.driver as d; print(d.Device(0).name())"

Resources

• LLVM Language Reference
• llvmlite Documentation
• Triton Source Code — study the real thing
• PTX ISA Reference
• PyCUDA Documentation
• CUDA C Programming Guide