luisagroup Manual AST Construction
Two approaches: DSL (Kernel2D<>, Callable<> lambdas) vs Manual AST (FunctionBuilder). Use manual AST for codegen, metaprogramming, programmatic kernel building.
Header: #include <luisa/ast/function_builder.h>
FunctionBuilder
using FuncBuilder = luisa::compute::detail::FunctionBuilder;
auto kernel = FuncBuilder::define_kernel([&]() { auto &cur = *FuncBuilder::current(); ... });
auto callable = FuncBuilder::define_callable([&]() { ... });
auto raster = FuncBuilder::define_raster_stage([&]() { ... });
All three return luisa::shared_ptr<const FuncBuilder>. define_kernel may duplicate the builder if outlined functions leak locals, so keep the returned pointer.
Built-in Variables
auto &cur = *FuncBuilder::current();
cur.dispatch_id(); // uint3
cur.dispatch_size(); // uint3
cur.thread_id(); // uint3
cur.block_id(); // uint3
cur.kernel_id(); // uint (indirect kernels only)
cur.warp_lane_id(); // uint
cur.warp_lane_count(); // uint
// Rasterization stage only
cur.raster_object_id(); // uint
cur.raster_barycentrics();// uint
Config
cur.set_block_size(uint3(16, 16, 1));
cur.set_name("my_kernel");
cur.set_variable_name(var_uid, "my_var");
auto name = cur.get_variable_name(var_uid);
cur.mark_variable_usage(var_uid, Usage::READ_WRITE);
Variables
// Arguments
cur.argument(Type::of<float3>()); // input value
auto ref = cur.reference(Type::of<uint2>()); // inout parameter
// Memory
cur.local(Type::of<float>()); // local variable
auto arr = cur.shared(Type::array(Type::of<float>(), count)); // shared array
// Constants
float4 v{1,2,3,4};
auto cdata = ConstantData::create(Type::of<float4>(), &v, sizeof(v));
auto c = cur.constant(cdata);
// Resources
cur.buffer(Type::of<Buffer<float>>()); // buffer
cur.texture(Type::of<Image<float>>()); // 2D texture
cur.texture(Type::of<Image3D<float>>()); // 3D texture
cur.accel(); // acceleration structure
cur.bindless_array(); // bindless array
Bindings
Binding methods create a bound argument and return its RefExpr*. Use that expression directly in the function body.
auto bound_buf = cur.buffer_binding(Type::of<Buffer<float>>(), handle, offset_bytes, size_bytes);
auto bound_tex = cur.texture_binding(Type::of<Image<float>>(), handle, level);
auto bound_ba = cur.bindless_array_binding(handle);
auto bound_acc = cur.accel_binding(handle);
// Example: read from a bound buffer
auto idx = cur.literal(Type::of<uint>(), 0u);
auto value = cur.call(Type::of<float>(), CallOp::BUFFER_READ, {bound_buf, idx});
Expressions
Literals
cur.literal(Type::of<float>(), 1.0f);
cur.literal(Type::of<int>(), 42);
cur.literal(Type::of<uint>(), 0u);
cur.literal(Type::of<bool>(), true);
cur.literal(Type::of<float2>(), float2(0.5f, 0.5f));
Binary Operations
Comparison operators are named, not symbolic.
cur.binary(Type::of<float>(), BinaryOp::ADD, a, b); // +
cur.binary(Type::of<float>(), BinaryOp::SUB, a, b); // -
cur.binary(Type::of<float>(), BinaryOp::MUL, a, b); // *
cur.binary(Type::of<float>(), BinaryOp::DIV, a, b); // /
cur.binary(Type::of<int>(), BinaryOp::MOD, a, b); // %
cur.binary(Type::of<uint>(), BinaryOp::BIT_AND, a, b); // &
cur.binary(Type::of<uint>(), BinaryOp::BIT_OR, a, b); // |
cur.binary(Type::of<uint>(), BinaryOp::BIT_XOR, a, b); // ^
cur.binary(Type::of<uint>(), BinaryOp::SHL, a, b); // <<
cur.binary(Type::of<uint>(), BinaryOp::SHR, a, b); // >>
cur.binary(Type::of<bool>(), BinaryOp::AND, a, b); // &&
cur.binary(Type::of<bool>(), BinaryOp::OR, a, b); // ||
cur.binary(Type::of<bool>(), BinaryOp::EQUAL, a, b); // ==
cur.binary(Type::of<bool>(), BinaryOp::NOT_EQUAL, a, b); // !=
cur.binary(Type::of<bool>(), BinaryOp::LESS, a, b); // <
cur.binary(Type::of<bool>(), BinaryOp::LESS_EQUAL, a, b); // <=
cur.binary(Type::of<bool>(), BinaryOp::GREATER, a, b); // >
cur.binary(Type::of<bool>(), BinaryOp::GREATER_EQUAL, a, b); // >=
Unary Operations
cur.unary(Type::of<float>(), UnaryOp::PLUS, value); // +
cur.unary(Type::of<float>(), UnaryOp::MINUS, value); // -
cur.unary(Type::of<bool>(), UnaryOp::NOT, value); // !
cur.unary(Type::of<uint>(), UnaryOp::BIT_NOT, value); // ~
Swizzle
Component indices are packed in 4-bit nibbles, lowest bits first:
// .xy from uint3: (0) | (1 << 4)
uint64_t swizzle_xy = (0ull) | (1ull << 4ull);
cur.swizzle(Type::of<uint2>(), coord_uint3, 2, swizzle_xy);
// .xyzw (all): 0x3210u
cur.swizzle(Type::of<float4>(), vec, 4, 0x3210u);
// .x: 0ull .y: 1ull<<4 .z: 2ull<<8 .w: 3ull<<12
Function Calls
// Built-in
cur.call(Type::of<float4>(), CallOp::MAKE_FLOAT4, {r, g, b, a});
cur.call(CallOp::TEXTURE_WRITE, {texture, coord, color});
cur.call(Type::of<float4>(), CallOp::TEXTURE_READ, {texture, coord});
// Buffer
cur.call(Type::of<float>(), CallOp::BUFFER_READ, {buffer, index});
cur.call(CallOp::BUFFER_WRITE, {buffer, index, value});
// Atomic (use AtomicRefNode)
auto ref = luisa::compute::detail::AtomicRefNode::create(buffer)
->access(index);
auto old = ref->operate(CallOp::ATOMIC_EXCHANGE, {new_value});
// Custom callable
cur.call(Function(callable.get()), {arg1, arg2});
Other Expressions
cur.cast(Type::of<float>(), CastOp::STATIC, int_value); // type cast
cur.cast(Type::of<float>(), CastOp::BITWISE, int_value); // bitwise reinterpretation
cur.access(Type::of<float>(), buffer_expr, index_expr); // array/buffer access
cur.member(Type::of<float>(), struct_expr, member_index); // struct member
cur.make_vector(Type::of<float4>(), luisa::vector{x, y, z, w}); // vector construction
cur.string_id("my_string"); // string ID -> uint64
cur.type_id(Type::of<float3>()); // type ID -> uint64
Statements
cur.assign(lhs_expr, rhs_expr);
// Control flow
cur.break_();
cur.continue_();
cur.return_(value_expr); // with value
cur.return_(); // void
// Use cur.with(scope, body) to populate nested scopes
auto if_stmt = cur.if_(cond);
cur.with(if_stmt->true_branch(), [&] { /* then */ });
cur.with(if_stmt->false_branch(), [&] { /* else */ });
auto loop_stmt = cur.loop_();
cur.with(loop_stmt->body(), [&] { /* loop body */ });
auto for_stmt = cur.for_(var, cond, step);
cur.with(for_stmt->body(), [&] { /* for body */ });
auto switch_stmt = cur.switch_(expr);
cur.with(switch_stmt->body(), [&] {
auto case0 = cur.case_(cur.literal(Type::of<int>(), 0));
cur.with(case0->body(), [&] { ...; cur.break_(); });
auto case1 = cur.case_(cur.literal(Type::of<int>(), 1));
cur.with(case1->body(), [&] { ...; cur.break_(); });
auto def = cur.default_();
cur.with(def->body(), [&] { ...; cur.break_(); });
});
auto ray_query_stmt = cur.ray_query_(query_expr);
cur.with(ray_query_stmt->on_triangle_candidate(), [&] { ... });
cur.with(ray_query_stmt->on_procedural_candidate(), [&] { ... });
auto ad_stmt = cur.autodiff_();
cur.with(ad_stmt->body(), [&] { ... });
// Print
cur.print_("value = {}", luisa::vector<const Expression *>{value_expr});
// Comment
cur.comment_("marker");
Type System
Getting Types
// Scalars
Type::of<float>(); Type::of<int>(); Type::of<uint>(); Type::of<bool>();
Type::of<half>(); Type::of<double>();
Type::of<short>(); Type::of<ushort>();
Type::of<int8_t>(); Type::of<uint8_t>();
Type::of<slong>(); Type::of<ulong>();
// Vectors
Type::of<float2>(); Type::of<float3>(); Type::of<float4>();
Type::of<int2>(); Type::of<int3>(); Type::of<int4>();
Type::of<uint2>(); Type::of<uint3>(); Type::of<uint4>();
Type::of<half2>(); Type::of<double4>(); // etc.
// Matrices
Type::of<float2x2>(); Type::of<float3x3>(); Type::of<float4x4>();
// Resources
Type::of<Buffer<float>>();
Type::of<Image<float>>(); Type::of<Image3D<float>>();
Type::of<Accel>(); Type::of<BindlessArray>();
Constructing Types
Type::vector(Type::of<float>(), 2); // float2
Type::matrix(4); // float4x4
Type::array(Type::of<float>(), 100); // float[100]
Type::structure({Type::of<float>(), Type::of<int>()}); // struct
Type::buffer(Type::of<float>()); // buffer<float>
Type::texture(Type::of<float>(), 2); // 2D texture
Type::texture(Type::of<float>(), 3); // 3D texture
Type::custom("MyOpaqueType");
Type::from("vector<float,4>"); // from string
Operators
BinaryOp
ADD, SUB, MUL, DIV, MOD // arithmetic
BIT_AND, BIT_OR, BIT_XOR // bitwise
SHL, SHR // shift
AND, OR // logical
LESS, GREATER, LESS_EQUAL, GREATER_EQUAL, EQUAL, NOT_EQUAL // comparison
UnaryOp
PLUS, MINUS, NOT, BIT_NOT
CallOp
The full set is defined in include/luisa/ast/op.h. Common groups:
// Vector construction
MAKE_FLOAT2/3/4, MAKE_INT2/3/4, MAKE_UINT2/3/4, MAKE_BOOL2/3/4
MAKE_SHORT2/3/4, MAKE_USHORT2/3/4, MAKE_LONG2/3/4, MAKE_ULONG2/3/4
MAKE_HALF2/3/4, MAKE_DOUBLE2/3/4, MAKE_BYTE2/3/4, MAKE_UBYTE2/3/4
MAKE_FLOAT2X2/3X3/4X4
// Buffer/Texture
BUFFER_READ, BUFFER_WRITE, BUFFER_SIZE, BUFFER_ADDRESS
BUFFER_VOLATILE_READ, BUFFER_VOLATILE_WRITE
BYTE_BUFFER_READ, BYTE_BUFFER_WRITE, BYTE_BUFFER_SIZE
TEXTURE_READ, TEXTURE_WRITE, TEXTURE_SIZE
TEXTURE2D_SAMPLE, TEXTURE2D_SAMPLE_LEVEL, TEXTURE2D_SAMPLE_GRAD, ...
// Atomic
ATOMIC_EXCHANGE, ATOMIC_COMPARE_EXCHANGE, ATOMIC_FETCH_ADD, ATOMIC_FETCH_SUB
ATOMIC_FETCH_AND, ATOMIC_FETCH_OR, ATOMIC_FETCH_XOR, ATOMIC_FETCH_MIN, ATOMIC_FETCH_MAX
// Bindless
BINDLESS_TEXTURE2D_SAMPLE, BINDLESS_TEXTURE2D_READ, BINDLESS_TEXTURE2D_SIZE
BINDLESS_BUFFER_READ, BINDLESS_BUFFER_WRITE, BINDLESS_BUFFER_SIZE, ...
UNIFORM_BINDLESS_*, TYPED_BINDLESS_*, TYPED_UNIFORM_BINDLESS_*
// Ray tracing
RAY_TRACING_TRACE_CLOSEST, RAY_TRACING_TRACE_ANY
RAY_TRACING_QUERY_ALL, RAY_TRACING_QUERY_ANY
RAY_TRACING_SET_INSTANCE_TRANSFORM, RAY_TRACING_SET_INSTANCE_VISIBILITY, ...
RAY_QUERY_WORLD_SPACE_RAY, RAY_QUERY_TRIANGLE_CANDIDATE_HIT,
RAY_QUERY_COMMIT_TRIANGLE, RAY_QUERY_TERMINATE, RAY_QUERY_PROCEED, ...
// Math
ALL, ANY, SELECT, CLAMP, SATURATE, LERP, SMOOTHSTEP, STEP
ABS, MIN, MAX, CLZ, CTZ, POPCOUNT, REVERSE
ISINF, ISNAN
SIN, COS, TAN, ASIN, ACOS, ATAN, ATAN2, SINH, COSH, TANH, ASINH, ACOSH, ATANH
EXP, EXP2, EXP10, LOG, LOG2, LOG10, POW, SQRT, RSQRT
CEIL, FLOOR, FRACT, TRUNC, ROUND, FMA, COPYSIGN
// Vector/Matrix
DOT, CROSS, LENGTH, LENGTH_SQUARED, NORMALIZE, FACEFORWARD, REFLECT, REFRACT
OUTER_PRODUCT, MATRIX_COMPONENT_WISE_MULTIPLICATION
DETERMINANT, TRANSPOSE, INVERSE
// Warp/Wave
WARP_IS_FIRST_ACTIVE_LANE, WARP_FIRST_ACTIVE_LANE, WARP_ACTIVE_ALL_EQUAL
WARP_ACTIVE_BIT_AND/OR/XOR, WARP_ACTIVE_COUNT_BITS, WARP_ACTIVE_MAX/MIN
WARP_ACTIVE_PRODUCT/SUM, WARP_ACTIVE_ALL/ANY, WARP_ACTIVE_BIT_MASK
WARP_PREFIX_SUM, WARP_PREFIX_PRODUCT, WARP_PREFIX_COUNT_BITS
WARP_READ_LANE, WARP_READ_FIRST_ACTIVE_LANE
// Sync
SYNCHRONIZE_BLOCK
// Rasterization
RASTER_DISCARD, RASTER_SET_Z_DEPTH,
RASTER_SET_Z_DEPTH_GREATER_EQUAL, RASTER_SET_Z_DEPTH_LESS_EQUAL
// Derivatives
DDX, DDY
// Indirect dispatch
INDIRECT_SET_DISPATCH_KERNEL, INDIRECT_SET_DISPATCH_COUNT
// Debugging/optimization
ASSERT, ASSUME, UNREACHABLE, FLATTEN, BRANCH, FORCE_CASE
// Clock
CLOCK
Usage Flags
enum struct Usage : uint32_t {
NONE = 0u, READ = 0x01u, WRITE = 0x02u, READ_WRITE = READ | WRITE
};
References must be marked explicitly:
cur.mark_variable_usage(ref->variable().uid(), Usage::READ_WRITE);
mark_variable_usage ORs flags, so it is safe to call multiple times.
Examples
Simple Kernel
auto kernel = FuncBuilder::define_kernel([&]() {
auto &cur = *FuncBuilder::current();
cur.set_block_size(uint3(16, 16, 1));
auto dispatch = cur.dispatch_id();
auto img = cur.texture(Type::of<Image<float>>());
auto color = cur.argument(Type::of<float4>());
auto coord = cur.swizzle(Type::of<uint2>(), dispatch, 2, (0ull) | (1ull << 4ull));
cur.call(CallOp::TEXTURE_WRITE, {img, coord, color});
});
Callable with Reference
auto callable = FuncBuilder::define_callable([&]() {
auto &cur = *FuncBuilder::current();
auto tex = cur.texture(Type::of<Image<float>>());
auto coord_ref = cur.reference(Type::of<uint2>());
cur.mark_variable_usage(coord_ref->variable().uid(), Usage::READ_WRITE);
auto color = cur.argument(Type::of<float3>());
auto alpha = cur.literal(Type::of<float>(), 1.0f);
auto value = cur.make_vector(Type::of<float4>(),
luisa::vector<const Expression *>{color, alpha});
cur.call(CallOp::TEXTURE_WRITE, {tex, coord_ref, value});
});
Kernel Calling Callable
auto kernel = FuncBuilder::define_kernel([&]() {
auto &cur = *FuncBuilder::current();
cur.set_block_size(uint3(16, 16, 1));
auto img = cur.texture(Type::of<Image<float>>());
auto color = cur.argument(Type::of<float3>());
auto coord_uint3 = cur.dispatch_id();
auto coord = cur.local(Type::of<uint2>());
cur.assign(coord, cur.swizzle(Type::of<uint2>(), coord_uint3, 2, (0ull) | (1ull << 4ull)));
cur.call(Function(callable.get()), {img, coord, color});
});
Swizzle Operations
auto kernel = FuncBuilder::define_kernel([&]() {
auto &cur = *FuncBuilder::current();
auto input = cur.argument(Type::of<float4>());
auto output = cur.reference(Type::of<float4>());
cur.mark_variable_usage(output->variable().uid(), Usage::READ_WRITE);
uint64_t swizzle_xyz = (0ull) | (1ull << 4ull) | (2ull << 8ull);
auto xyz = cur.swizzle(Type::of<float3>(), input, 3, swizzle_xyz);
auto w = cur.swizzle(Type::of<float>(), input, 1, 3ull); // .w
cur.assign(output, cur.make_vector(Type::of<float4>(), luisa::vector{x, w}));
});
Buffer Operations
auto kernel = FuncBuilder::define_kernel([&]() {
auto &cur = *FuncBuilder::current();
cur.set_block_size(uint3(256, 1, 1));
auto input_buf = cur.buffer(Type::of<Buffer<float>>());
auto output_buf = cur.buffer(Type::of<Buffer<float>>());
auto idx = cur.swizzle(Type::of<uint>(), cur.thread_id(), 1, 0ull);
auto value = cur.call(Type::of<float>(), CallOp::BUFFER_READ, {input_buf, idx});
auto scaled = cur.binary(Type::of<float>(), BinaryOp::MUL, value, cur.literal(Type::of<float>(), 2.0f));
auto result = cur.binary(Type::of<float>(), BinaryOp::ADD, scaled, cur.literal(Type::of<float>(), 1.0f));
cur.call(CallOp::BUFFER_WRITE, {output_buf, idx, result});
});
For Loop
auto kernel = FuncBuilder::define_kernel([&]() {
auto &cur = *FuncBuilder::current();
auto buf = cur.buffer(Type::of<Buffer<float>>());
auto i = cur.local(Type::of<uint>());
cur.assign(i, cur.literal(Type::of<uint>(), 0u));
auto ten = cur.literal(Type::of<uint>(), 10u);
auto cond = cur.binary(Type::of<bool>(), BinaryOp::LESS, i, ten);
auto step = cur.literal(Type::of<uint>(), 1u);
auto for_stmt = cur.for_(i, cond, step);
cur.with(for_stmt->body(), [&] {
auto idx = i;
auto v = cur.call(Type::of<float>(), CallOp::BUFFER_READ, {buf, idx});
cur.call(CallOp::BUFFER_WRITE, {buf, idx,
cur.binary(Type::of<float>(), BinaryOp::ADD, v,
cur.literal(Type::of<float>(), 1.0f))});
});
});
Key Rules
- Always use
Type::of<T>()for explicit types. - Mark reference usage with
mark_variable_usage(uid, Usage::READ_WRITE). - Swizzle: component indices in nibbles, lowest bits first.
- Use
FunctionBuilder::current()within define callbacks. - Statements and expressions are owned by
FunctionBuilder; do not delete them. - Set block size for compute kernels (typically
uint3(16, 16, 1)for 2D). - Use
cur.with(scope, body)to append statements intoif/loop/for/switch/ray_query/autodiffbodies. - Atomic operations require
AtomicRefNode, not raw buffer variables. print_takes a format string and aluisa::span/vectorof expressions, not an initializer list.BinaryOpcomparison names areEQUAL,NOT_EQUAL,LESS,GREATER,LESS_EQUAL,GREATER_EQUAL.