FlyDSL Python DSL

FlyDSL Python DSL#

The flydsl package provides the Python front-end for authoring GPU kernels with explicit layout algebra.

Core module#

Expression API (flydsl.expr)#

The flydsl.expr module (imported as fx) provides the high-level Python API for constructing Fly IR, including layout construction, tiled copies, tensor operations, and kernel definitions.

import flydsl.expr as fx

Layout construction#

  • fx.make_layout(shape, stride) – create a layout from shape and stride tuples

  • fx.make_shape(*dims) – create a shape tuple

  • fx.make_stride(*strides) – create a stride tuple

  • fx.make_coord(*coords) – create a coordinate tuple

  • fx.make_ordered_layout(shape, order) – layout with explicit mode ordering

  • fx.make_identity_layout(shape) – identity layout (strides = prefix products)

Layout inspection#

  • fx.size(layout) – total number of elements

  • fx.cosize(layout) – codomain size

  • fx.rank(layout) – number of modes

  • fx.depth(layout) – nesting depth

  • fx.get_shape(layout) – extract shape tuple

  • fx.get_stride(layout) – extract stride tuple

  • fx.get_scalar(int_tuple) – extract the scalar from a single-leaf int tuple; use fx.get_(int_tuple, mode).unpack() for per-mode access

Layout algebra#

  • fx.composition(a, b) – compose two layouts

  • fx.complement(layout, codomain_size) – complementary layout

  • fx.right_inverse(layout) – right inverse

  • fx.coalesce(layout) – coalesce contiguous modes

  • fx.recast_layout(layout, old_type, new_type) – recast layout for type change

Layout products and divides#

  • fx.logical_divide(tensor, tiler) – partition tensor by tiler layout

  • fx.zipped_divide, fx.tiled_divide, fx.flat_divide – divide variants

  • fx.logical_product(a, b) – layout product

  • fx.zipped_product, fx.tiled_product, fx.flat_product – product variants

  • fx.raked_product(thr_layout, val_layout) – interleaved (raked) product

  • fx.blocked_product(a, b) – blocked product

Coordinate mapping#

  • fx.crd2idx(coord, layout) – coordinate to linear index

  • fx.idx2crd(idx, layout) – linear index to coordinate

  • fx.slice(tensor, slices) – slice a tensor by coordinates or None

  • fx.get_(int_tuple, mode) – select a mode while preserving its structured result; use .unpack() when a scalar result is required. The older fx.get spelling is deprecated

Memory operations#

  • fx.make_rmem_tensor(shape_or_layout, dtype) – allocate register-file memory

  • fx.memref_load(memref, indices) – scalar load from memref

  • fx.memref_store(value, memref, indices) – scalar store to memref

  • fx.memref_load_vec(memref) – load entire register as a vector

  • fx.memref_store_vec(vec, memref) – store vector to register memref

  • fx.make_fragment_layout_like(tensor) – compute the corresponding fragment layout

  • fx.make_fragment_like(tensor) – allocate register fragment with same layout

Copy and GEMM#

  • fx.make_copy_atom(instr, dtype) – create a CopyAtom from instruction descriptor

  • fx.make_mma_atom(instr) – create an MmaAtom from an MMA op type (the op type carries the dtype, for example fx.rocdl.MFMA(16, 16, 4, fx.Float32))

  • fx.make_tile(*layouts) – build a tile from layouts (variadic)

  • fx.make_tiled_copy(copy_atom, layout_tv, tile_mn) – build a TiledCopy

  • fx.make_tiled_mma(mma_atom, …) – build a TiledMma

  • fx.copy(copy_atom, src, dst, pred=None) – execute a copy (with optional predicate mask)

  • fx.gemm(mma_atom, d, a, b, c) – execute matrix multiply-accumulate (accumulator passed as both d and c)

  • fx.copy_atom_call(atom, src, dst) – invoke a single copy atom

  • fx.mma_atom_call(atom, d, a, b, c) – invoke a single MMA atom

Derived tiled operations (flydsl.expr.derived)#

High-level classes for tiled copy and MMA partitioning:

  • CopyAtom (flydsl.expr.typing) – single hardware copy instruction descriptor

  • MmaAtom (flydsl.expr.typing) – single MMA instruction descriptor (MFMA)

  • CopyAtomType, MmaAtomType – atom type wrappers exported by flydsl.expr.primitive

  • TiledCopy – multi-thread tiled copy; use get_slice(tid) → ThrCopy

  • TiledMma – multi-thread tiled MMA; use get_slice(tid) → ThrMma

  • ThrCopy – per-thread copy view: partition_S(src), partition_D(dst), retile(t)

  • ThrMma – per-thread MMA view: partition_A(a), partition_B(b), partition_C(c)

  • make_layout_tv(thr, val) – build thread-value layout

  • make_tiled_copy_A/B/C(copy_atom, tiled_mma) – create TiledCopy matched to MMA operands

  • fx.gather(copy_atom, base_iter, offset_tensor, dst_tensor, *, pred=None) – indexed load dst = base[offset] via a copy atom (offset tensor is (TV, Rest...))

  • fx.scatter(copy_atom, src_tensor, base_iter, offset_tensor, *, pred=None) – indexed store base[offset] = src (see examples/05-gather_scatter.py)

Type annotations#

  • fx.Tensor – GPU tensor argument

  • fx.Constexpr[int] – compile-time constant

  • fx.Int32 – dynamic int32 argument

  • fx.Float32, fx.Float16, fx.BFloat16 – scalar types

  • fx.Float8E4M3FN, fx.Float8E4M3FNUZ, fx.Float8E5M2 – FP8 scalar types

  • fx.Stream – GPU stream argument

  • fx.T – type namespace (T.f32, T.f16, T.bf16, T.i8, T.index, etc.)

  • fx.Basis(value, modes) / fx.E(*modes) – basis-stride leaves for by-mode layout construction (E(0) → 1E0)

  • fx.SyncScope – target-neutral LLVM sync scopes (SyncScope.System, SyncScope.SingleThread); AMDGPU scopes live in flydsl.expr.rocdl.enum.SyncScope

GPU intrinsics (flydsl.expr.gpu)#

  • fx.thread_idx – thread index (Tuple3D with .x, .y, .z)

  • fx.block_idx – block index

  • fx.block_dim – block dimensions

  • fx.grid_dim – grid dimensions

  • fx.gpu.barrier() – workgroup barrier synchronization

  • fx.gpu.smem_space() – shared memory (LDS) address space attribute

Arithmetic and numeric types#

Prefer typed DSL values and operator-overloaded arithmetic:

import flydsl.expr as fx
from flydsl.expr.typing import Vector as Vec

c = fx.Int64(42)
v = fx.Int32(idx)
f = fx.Float32(1.0)
r = cond.select(a, b)
y = (x + 1) * scale

Preferred APIs:

  • fx.Int32(value), fx.Int64(value), fx.Float32(value) – typed constants and casts (use fx.Int64 for index/offset values and loop bounds; fx.Index is deprecated)

  • ArithValue / Numeric operators – +, -, *, /, %, <<, >>

  • cond.select(true_val, false_val) – ternary select when cond is an ArithValue

  • arith.cmpi(predicate, lhs, rhs) – integer comparison

  • arith.cmpf(predicate, lhs, rhs) – float comparison

  • fx.maxnumf(a, b) – float maximum returning the non-NaN operand (libm fmax); preserves the DSL type of a

  • fx.minnumf(a, b) – float minimum returning the non-NaN operand (libm fmin); preserves the DSL type of a

  • Chained comparisons (lo <= x < hi) are supported inside traced kernels and lower to combined cmp + and.

You can apply fastmath flags ambiently to a block or per-op:

with fx.fastmath(fx.FastMathFlags.fast):
    y = a * b + c          # float operators/math funcs inherit the flags
    z = fx.exp(a, fastmath="contract")   # explicit arg overrides the ambient scope

  • fx.fastmath(flags) – context manager applying fastmath to float ops built in the block; nests and restores on exit

  • fx.FastMathFlags – flag enum (fast, contract, reassoc, …; combine with |)

  • Direct ``arith.addf(…, fastmath=…)`` / ``arith.AddFOp(…, fastmath=…)`` – per-op flags where an ambient scope is not desired

Vector values (flydsl.expr.typing.Vector)#

  • Vec.from_elements(elements, dtype) – construct vector from scalars

  • Vec.filled(shape, value, dtype) – splat vector

  • Vec(value)[i] – extract element

  • Vec(value).bitcast(dtype) – bitcast vector element type

  • Vec(value).to(dtype) – convert vector element type

  • Vec(value).store(memref, indices) – store vector to memref

ROCDL operations (flydsl.expr.rocdl)#

AMD-specific operations for ROCm:

The exported universal and architecture namespaces are governed by API Stability. Importable target helpers or generated upstream ROCDL builders that are absent from the export chain remain unstable even when they are useful for source-tree kernel development.

  • fx.rocdl.make_buffer_tensor(tensor) – create buffer resource from tensor (CDNA buffer copy)

  • fx.rocdl.BufferCopy32b / BufferCopy128b – buffer copy instruction atoms

  • fx.rocdl.MFMA(m, n, k, elem_ty_ab, elem_ty_acc=None) – MFMA instruction atom constructor (CDNA3/CDNA4; 4th arg is the A/B element type; accumulator defaults to f32)

  • fx.rocdl.WMMA(m, n, k, elem_ty_ab, elem_ty_acc=None, **kwargs) – WMMA MMA atom constructor (arch-dispatched: gfx11 / gfx120x / gfx1250). elem_ty_b optionally selects a different B operand type. gfx1250 supports f32(K4), f16/bf16(K32), fp8/bf8(K64/128), i8(K64), i4(K32); integer paths take sign_a / sign_b / clamp. gfx120x (RDNA4) supports 16x16x16 f16/bf16 and every fp8(E4M3FN)/bf8(E5M2) A/B combination to f32, on the v8 operand ABI

  • fx.rocdl.WMMAScale(m, n, k, elem_ty_a, elem_ty_b=None, elem_ty_acc=None, *, opsel_a=0, opsel_b=0, mod_c=0, reuse_a=False, reuse_b=False, block_size=32) – unstable, target-specific gfx1250 MX-scaled WMMA (E8M0 block scale, f8/f6/f4; 16x16x128 or 32x16x128 fp4-only). Per-operand scales are atom state (scale_a / scale_b)

  • fx.rocdl.make_tdm_atom(tensor, tensor_extents, strides=None, *, num_warps, …) – unstable, target-specific gfx1250 TDM (Tensor Data Mover) async Global↔LDS whole-tile copy atom (rank 1-5); the global base comes from the copy_atom_call operand pointer, while the per-dim extent (OOB), stride, imm_offset, and MCAST workgroup_mask are atom state. fx.rocdl.TDM(rank, num_warps, ...) builds the atom type only. Advance the K-loop tile with fx.copy(atom, gt, dst, imm_offset=...)

  • fx.rocdl.sched_mfma(cnt) – insert MFMA scheduling barrier

  • fx.rocdl.sched_vmem(cnt) – insert VMEM scheduling barrier

  • fx.rocdl.sched_dsrd(cnt) – insert DS read scheduling barrier

  • fx.rocdl.sched_dswr(cnt) – insert DS write scheduling barrier

  • mfma_f32_16x16x16f16, mfma_f32_16x16x16bf16_1k, etc. – generated direct MFMA intrinsics; these low-level builders are unstable

Compiler API (flydsl.compiler)#

import flydsl.compiler as flyc

  • @flyc.kernel – decorator for GPU kernel functions

  • @flyc.jit – decorator for host-side JIT launch functions

  • flyc.from_dlpack(tensor) – convert DLPack-compatible tensors (PyTorch, etc.) to FlyDSL

  • JitArgumentRegistry – registry for custom argument type adapters

  • flydsl.compiler.kernel_function.CompilationContext – unstable implementation context available during kernel compilation; it is not a top-level flydsl.compiler symbol

See also

Compiler and pipeline for the full compilation pipeline and pass details.

Contents