ir_pipeline

LuisaCompute IR and XIR Pipeline

Overview

LuisaCompute has two intermediate representations:

Both paths start from the same AST (src/ast/) and feed into backend codegen.

---

Compiler Pipeline Flow

User C++ DSL / Python
         │
         ▼
┌─────────────────┐
│   DSL Tracing   │  src/dsl/ — captures C++ lambdas into AST
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│       AST       │  src/ast/ — expression/statement tree
└────────┬────────┘
         │
    ┌────┴────┐
    ▼         ▼
┌────────┐ ┌──────────┐
│  XIR   │ │   IR     │  src/xir/ (new) / src/ir/ (legacy)
│transl. │ │transl.   │
└───┬────┘ └────┬─────┘
    │           │
    ▼           ▼
┌─────────────────┐
│  Backend Codegen│  src/backends/<name>/
│  + Compilation  │
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│  GPU Execution  │  src/runtime/
└─────────────────┘

Rust IR path: src/rust/luisa_compute_ir/ performs additional transforms (autodiff, DCE, SSA, vectorize) before backend codegen.

XIR path: Pure C++ pipeline with ast2xir translator and XIR optimization passes.

---

AST → IR Translation (Legacy)

File: src/ir/ast2ir.cpp

1. AST → JSON: Uses to_json(function) to serialize AST to JSON
2. JSON → Rust IR: Calls Rust FFI functions:
   • ir::luisa_compute_ir_ast_json_to_ir_kernel() for kernels
   • ir::luisa_compute_ir_ast_json_to_ir_callable() for callables
   • ir::luisa_compute_ir_ast_json_to_ir_type() for types

AST Function → JSON String → Rust IR Module → CArc<KernelModule/CallableModule>

Key legacy IR classes (defined in Rust, used via C FFI):

• ir::KernelModule / ir::CallableModule
• ir::Node — IR nodes (intrusive linked list)
• ir::Instruction — Local, Call, Phi, Loop, If, Switch, RayQuery, AdScope, etc.
• ir::Type — Type representation
• ir::BasicBlock — Basic block structure

---

AST → XIR Translation

File: src/xir/translators/ast2xir.cpp

Direct C++ implementation using XIRBuilder API.

Translation Context (AST2XIRContext)

• Maps AST variables to XIR Values
• Tracks break/continue target blocks
• Handles autodiff adjoint variables
• Manages constant/literal caching

Expression Translation

Statement Translation

---

XIR Core Architecture

Value Hierarchy

Value (base class)
├── GlobalValue (module-scoped)
│   ├── Function
│   │   ├── KernelFunction (entry point with block size)
│   │   ├── CallableFunction (user-defined functions)
│   │   └── ExternalFunction (external linkage)
│   ├── Constant (literal values)
│   ├── Undefined (placeholder for undef)
│   └── SpecialRegister (builtin variables)
│       ├── SPR_ThreadID, SPR_BlockID, SPR_DispatchID, etc.
├── FunctionScopeValue (function-scoped)
│   └── BasicBlock (container for instructions)
└── BlockScopeValue (block-scoped)
    └── Instruction (computation)
        ├── TerminatorInstruction (control flow)
        │   ├── BranchInst
        │   ├── ConditionalBranchInst
        │   ├── IfInst, SwitchInst, LoopInst, etc.
        │   └── ReturnInst
        └── Non-terminator instructions

Key Classes

Module (include/luisa/xir/module.h)

• Container for all global values
• Creates/manages functions, constants, undefined values, special registers
• Uniquifies constants via hash table

Function (include/luisa/xir/function.h)

• ArgumentList — function parameters
• BasicBlockList — CFG nodes
• FunctionDefinition — for kernels/callables with body blocks
• Traversal: traverse_basic_blocks(), traverse_instructions()

BasicBlock (include/luisa/xir/basic_block.h)

• InstructionList — sequential instructions
• is_terminated(), terminator() — control flow queries
• traverse_predecessors(), traverse_successors() — CFG navigation

Instruction (include/luisa/xir/instruction.h)

• Derived via DerivedInstruction<> template
• is_terminator() — ends basic block
• control_flow_merge() — for structured control flow
• clone() — for pass transformations

Value & Use (include/luisa/xir/value.h, use.h)

• SSA Form: Each value has a UseList tracking all users
• replace_all_uses_with() — for substitution
• is_lvalue() — for memory operations (alloca, gep results)

---

XIR Instruction Set

Control Flow Instructions

Memory Instructions

SSA Instructions

Arithmetic Instructions (ArithmeticOp)

Unary: UNARY_MINUS, UNARY_BIT_NOT

Binary: BINARY_ADD, SUB, MUL, DIV, MOD, BIT_AND, BIT_OR, BIT_XOR, SHIFT_LEFT, SHIFT_RIGHT, ROTATE_LEFT, ROTATE_RIGHT, comparison ops

Math: ABS, MIN, MAX, CLAMP, SATURATE, LERP, SMOOTHSTEP, trigonometric, exponential, logarithmic

Vector/Matrix: DOT, CROSS, NORMALIZE, MATRIX_COMP_MUL, MATRIX_LINALG_MUL, MATRIX_DETERMINANT, MATRIX_TRANSPOSE, MATRIX_INVERSE

Aggregate: AGGREGATE, SHUFFLE, EXTRACT, INSERT

Resource Instructions

Atomic Instructions (AtomicOp)

EXCHANGE, COMPARE_EXCHANGE, FETCH_ADD, FETCH_SUB, FETCH_AND, FETCH_OR, FETCH_XOR, FETCH_MIN, FETCH_MAX

Thread Group Instructions (ThreadGroupOp)

Warp operations, block synchronization, shader execution reordering, raster quad derivatives

Ray Query Instructions

RayQueryLoopInst, RayQueryDispatchInst, RayQueryObjectReadInst, RayQueryObjectWriteInst, RayQueryPipelineInst

Autodiff Instructions

AutodiffScopeInst, AutodiffIntrinsicInst (requires_gradient, gradient, accumulate_gradient, backward, detach)

---

XIR Optimization Passes

Location: src/xir/passes/

Core Passes

Analysis Passes

Transformation Passes

---

Control Flow Representation

XIR uses structured control flow with explicit merge blocks:

IfInst:
  - condition: Value*
  - true_block: BasicBlock*
  - false_block: BasicBlock*
  - merge_block: BasicBlock*

LoopInst:
  - prepare_block: BasicBlock*
  - body_block: BasicBlock*
  - update_block: BasicBlock*
  - merge_block: BasicBlock*

This design:

• Maintains SSA property
• Enables structured transformations
• Maps well to GPU shader control flow requirements
• Supports PHI nodes at merge points

---

Metadata System

File: include/luisa/xir/metadata.h

Metadata types (DerivedMetadataTag):

• NAME — Symbol names
• LOCATION — Source file/line information
• COMMENT — Debug comments
• CURVE_BASIS — Ray tracing curve basis

Applied via MetadataListMixin to Values and Instructions.

---

JSON Serialization

IR (Legacy)

• Uses ast2json for AST serialization
• Rust IR can be serialized/deserialized
• C API: luisa_compute_ir_ast_json_to_ir_*

XIR

Files: src/xir/translators/xir2json.cpp, json2xir.cpp

Uses yyjson library for module serialization, cross-process communication, and debugging.

---

Key Design Patterns

1. Intrusive Lists: ManagedIntrusiveList for efficient node management
2. CRTP: DerivedValue<>, DerivedInstruction<>, DerivedFunction<>
3. Visitor Pattern: traverse_basic_blocks(), traverse_instructions()
4. Builder Pattern: XIRBuilder for instruction creation
5. Mixin Pattern: MetadataListMixin, ControlFlowMergeMixin, InstructionOpMixin
6. Use-Def Chains: Use objects track value users for SSA

---

Adding a New XIR Pass

1. Create pass implementation in src/xir/passes/<name>.cpp
2. Create header in include/luisa/xir/passes/<name>.h
3. Register in src/xir/CMakeLists.txt
4. Follow existing pass conventions:
   • Accept Module & or Function &
   • Return bool indicating whether changes were made
   • Use XIRBuilder for instruction insertion
   • Call replace_all_uses_with() for value substitution