TIRx export and target realization

This page owns the shared structural export and target-specific realization contracts. Runtime integration owns shader creation, argument ABI and dispatch; planning owns legal candidate selection. Proposed scheduled forms must not be confused with the implemented bridge-local plans.

Compiler bridges and native backends

Boundary

TileIR is the source of truth. TVM is an optional lowering bridge used to bootstrap scheduling, legalization, and code generation. TVM types never appear in the C++ DSL API, TileIR core headers, cache ABI, or runtime ABI.

Interoperability code lives below luisa/tile/bridge/: bridge/tirx owns the TVM dependency and bridge/xir owns TileIR-to-Luisa-XIR translation. The first XIR realization targets independent CPU workers and connects to the existing SIMD backend; it is separate from both native MPP and TIRx. The TIRx bridge constructs tvm::tirx::PrimFunc, layouts, expressions, and statements directly through TVM’s public C++ API. TVMScript text is useful for debug printing and differential tests, but source generation or Python parsing is not a compiler boundary.

Native hardware lowering is not another bridge. DeviceInterface exposes an optional create_tile_kernel factory: TileIR and compile options enter the active backend, which owns the compiler-route choice and returns a normal ShaderCreationInfo plus launch/argument metadata. The default implementation reports unsupported. TVM types do not cross this ABI. A backend can use native lowering or an installed bridge to obtain its executable; a bridge is not a new hardware device. Validation wraps this factory like ordinary shader creation, and launches use existing shader handles and Stream commands.

Metal MPP and the optional TIRx device-artifact route are implemented through this factory, described in Independent Tile Runtime. tile::Lowering::NATIVE and tile::Lowering::TIRX select explicitly. TIRx retains its independent CPU/Metal TVM-runtime regression and comparison path; the new device artifact preserves a typed PrimFunc, argument binding map and static launch geometry, with unchanged generated Metal source. The CUDA backend adds the same direct-buffer reference route on Lowering::TIRX: compile_device emits a CUDA C source artifact (target kind cuda) that the standalone-NVRTC pipeline turns into PTX for a direct static launch; a nvptx target is accepted as direct PTX text. CUDA never requests cooperative-matrix/vector or indirect dispatch atoms — semantic Tile MMA, REDUCE, elementwise lift and tile map/extract keep the reference SIMT realization — and every Metal-only planning knob is rejected with an explicit diagnostic. Reference worker/elementwise domains on CUDA/NVPTX are warp aligned by the mapper: partial domains are rounded up to the 32-thread warp and unaligned per-block caps down. NVPTX PTX text is final output that the CUDA backend cannot recompile for another architecture; load failures therefore recommend the CUDA C source ("cuda") TIRx target. All factory routes return ordinary Runtime shaders. Neither silently falls back to TVM, MPP or MPS when its selected realization rejects a kernel. Intra-kernel Tile/SIMT DSL mixing is deliberately out of scope.

Layout bridge

TVM TIRx models a layout as named physical-axis coordinates with shard, replica, and offset components; see the official TIRx layout model. The TileIR embedding preserves that exact set-valued direction. For logical space X, replica space R, and named physical space P:

F = X x R
left(x, r)  = x
right(x, r) = D(x) + replica(r) + O

TIRx Layout X -> Set(P)
  = LayoutCorr<X, P; F>(left, right)

TIRx logical axes become X, its replica iters become R, named hardware or memory axes become P, its shard body becomes D, and its offset becomes Translate. This is an exact embedding, including replication.

A Triton-style distribution arrives in the other direction as a map PhysicalSlot -> Logical. TileIR views it as a correspondence and swaps the two legs when an exporter needs logical-to-physical placement. No unique inverse is assumed. Export to a TIRx TileLayout is structural when the reoriented correspondence factors as shard plus replica plus offset. Registered swizzles export as ComposeLayout when supported. Other layouts lower to explicit TIRx index computation or a legal materialization; they are never silently approximated.

Execution bridge

The open logical hierarchy remains in TileIR until target binding. Current TIRx execution scopes name a fixed set of GPU-like scopes, so they cannot be the canonical model for an open hierarchy.

After ExecBinding is concrete, the exporter maps target axes to TIRx scopes or TIRx loop/thread-binding constructs. Serial and vector axes remain explicit.

There is one shared structural exporter, not a complete lowerer per backend. The changing component is the execution schedule: a CPU plan may map a logical parallel prefix to a task loop and SIMD suffix, while a GPU plan may map the same prefix through an affine split to grid, threadgroup, subgroup, and worker coordinates. Target-specific resource selection and intrinsic dispatch happen after this binding. Schematically:

Candidate TileIR (logical execution tree)
                 |
          target/autotuned ExecBinding
                 v
Scheduled TileIR (physical scopes + index maps)
                 |
       shared structural TIRx exporter
                 v
       target-specific TVM code generation

A target realization may collapse a proved region into one opaque atom without mechanically preserving every loop as a machine loop. It must preserve the region’s observable values/effects, layout correspondence, alias contract and selected numerical policy. Current examples are a complete compact FP32 GEMM realized as one CBLAS call and shared exp/add/max/min regions realized through vForce/vDSP. Their versioned contracts come from typed TileIR dataflow and are revalidated against the transformed TIRx body. This is atom selection behind the execution model, not a GEMM/softmax frontend primitive or a memory-derived hierarchy.

In the complete planner, execution binding and atom selection are joint choices: a matrix atom may constrain participant/layout maps, while a whole CPU provider atom may consume an entire logical region and make an internal worker mapping irrelevant. The current CPU provider choices are explicit compile policies rather than an automatic cost-model decision; the portable reference path remains available for comparison and unsupported explicit requests fail closed.

The reference schedule leaves logical parallel as marked serial TIRx during structural export, including its optional execution-scope constraint. The target mapper resolves each constraint before consuming its annotation. The outermost unbound region maps to LLVM kParallel or to a Metal/CUDA-style blockIdx.x * threads + threadIdx.x grid with a tail predicate. Empty logical domains become no-ops, not invalid zero-sized GPU launches.

The currently implemented reference-mapper subset is listed below. The optional Metal row-reduction realization additionally admits proved subgroup/cooperating-group maps; it does not extend every reference operation to those scopes.

Target

Constraint

Realization

LLVM

outer exec::Scope::WORKER

host parallel loop

LLVM

exec::Scope::VECTOR root or worker suffix

vector loop with lane-private temporaries

Metal

outer exec::Scope::WORKER

one logical instance per GPU thread

Metal

outer exec::Scope::GROUP

one logical instance per threadgroup; independent elements or child workers cooperate

Outside a cooperative group, unbound nested parallel regions remain serial. Inside a group, the first child parallel or independent Tile-element domain uses the worker coordinate; deeper unbound regions remain serial per worker. Explicit nested worker/vector rebindings need a coordinate factorization that this reference planner does not yet implement, so they are rejected even through unbound/serial intermediate scopes. Device, subgroup, unknown scope names, CPU group bindings, and other unavailable bindings also fail closed. Disabling vectorization cannot silently override an explicit vector constraint. Automatic pointwise graphs can additionally use the fused grid described below. An automatic program containing reassociable MMA can try the composed cooperative family, including ordinary Tile phases and temporal carries. Unsupported candidates retain the reference worker mapping. These choices are not a public CPU_THREADS/GPU_GRID compile option.

Vector binding includes a separate resource transformation. A compiler-local temporary declared inside a vector instance has one independent copy per logical lane; a temporary in an ancestor scope does not. Before TIRx vectorization the bridge expands compact local storage by the address map:

Address(local_coord, lane) = flatten(local_coord) * vector_width + lane

The allocation is hoisted immediately outside the vector loop, and all of its loads/stores retain the lane coordinate. This is necessary because the current upstream TIRx vectorizer does not privatize AllocBuffer; merely changing the loop kind can incorrectly broadcast the last lane’s value to every lane. Parent Tiles, lane-local multidimensional Tiles, and simultaneous loop-carried updates have separate numerical regressions, alongside a generated LLVM vector-instruction check. Explicitly placed/laid-out resources are not rewritten by this compact-local-storage transformation.

Opt-in CPU automatic vectorization consumes the independent-element domain contract, without matching a particular operator or re-proving independence. For a supported rectangular domain it factors the innermost element axis as i = min + pack * width + lane, emits power-of-two vector packs of 4–16 logical lanes plus a scalar tail, and preserves the entire per-element body. In particular, an MMA’s K loop remains serial inside each output instance; no horizontal reduction or floating-point reassociation is introduced. Native TIRx ConvertSSA renews definitions copied into the full/tail bodies, and VectorizeLoop/LLVM perform SIMD legalization. A logical pack can become several native vectors, exposing independent accumulator chains without changing temporal order. There is no new per-MMA scratch allocation or DSL entity. Bodies containing allocations, while loops, or nested nonserial execution keep their reference mapping until the corresponding vectorization/storage support exists. CompileOptions::auto_vectorize defaults to false: measurements show both wins and substantial regressions, so this is a tunable candidate, not the default realization. Enabling it requires vectorize = true. Explicit vector execution bindings remain independent of this opt-in; disabling vectorize still rejects those explicit bindings. Tests cover scalar tails, transposes, ordered cancellation, and generated vector products.

Metal group binding similarly includes a resource transformation, not just a loop tag. The structural exporter marks independent element domains; contraction axes remain sequential in the reference realization. Temporal iterations retain their dependences even when a safe pipeline plan overlaps independent producer/consumer work. It retains each domain as a rectangular serial loop nest with a rank marker. Flattening is a cooperative-binding decision, not a shared-export decision; CPU worker/vector paths keep the individual axes available to their optimizer. The group mapper partitions the flattened independent domain as element = chunk * worker_count + worker, with a tail predicate. Compact compiler temporaries allocated in the group context become shared memory; those allocated inside a distributed worker region stay private. Multiple resources can therefore share the same execution level, and descendants can read ancestor Tiles without recreating a private copy:

parallel(..., GROUP)              one threadgroup per logical group
  |
  +-- load A -> temporary X -----+   group resource instance
  +-- map X  -> temporary Y -----+   another resource, same exec level
  |                             |
  +-- parallel(..., WORKER) -----+-> read X and Y; private local state
  |       each worker handles its own elements
  |
  +-- fence shared + device memory (all workers, including inactive tail)
  +-- next Tile operation        may consume a neighbor's result

The first implementation synchronizes conservatively between distributed operations. A scalar group effect executes once and publishes its result. The fence orders both shared resources and global views within that group; it is not device-wide synchronization between groups. Safe pipeline cuts may use the two-window software-prefetch plan; other pipelines stay ordered. Neither plan implies hardware-asynchronous copy. Reference MMA distributes output elements and retains each element’s serial contraction; parallel reduction trees remain separate planner work.

Automatic GPU pointwise graphs

The element-grid mapper consumes the product of a logical program domain and its independent Tile-element domain. It does not assign one entire Tile to each hardware thread. A single automatic root with same-domain pure SSA producers and one or more independent output domains can become:

logical program p × local element e
                  |
       linear = p * elements_per_program + e
                  |
       block = linear / threads; worker = linear % threads
                  |
       v = producer(input[p, e])       // one worker-local SSA value
       output_a[p, e] = f(v)
       w = another_producer(v)         // definition after an output is legal
       output_b[p, e] = g(v, w)
       output_c[p, e] = h(w)

This is an IR-structure rule, not recognition of sigmoid, GELU, a kernel name, or a particular tensor size. Each original domain is normalized to the same local coordinates, including nonzero loop minima. Interleaving preserves per-element statement order. An output may use a different coordinate permutation or predicate if its write map is independently proved injective. Shared compiler-owned Tiles become scalar definitions only after every load proves the same owner as its unique dominating producer; neither shape equality nor the parallel contract alone establishes that ownership.

The current admission proof is deliberately bounded:

  • Parameters have the checked noalias contract; each written global buffer is never read or escaped by this graph.

  • Output buffers have supported compact address maps; producer and consumer domains have the same static rectangular rank/extents. Conditional output stores are allowed, conditional shared producers are not.

  • Shared definitions are compiler-owned pure Tiles. Unknown effects, manual resource/layout constraints, a different domain, pipeline boundaries and explicit execution bindings retain their original realization.

  • Multiple stores to one compact buffer require pairwise separation of their linear address ranges. Each proof uses two independent local-coordinate tuples and the same logical program coordinate. Checking only equal local coordinates would miss overlapping, shifted output loops. In-bounds static coordinates bound linearization in int64; predicate-sensitive bounds and unknown separation retain the original loops. Either source order is legal. Cross-domain RAW/WAR/WAW dependencies cannot be erased by assigning the same shape to both domains. This admits partitioned outputs such as split-half rotations without recognizing RoPE or changing the DSL.

The mapper emits worker-local scalar storage and a bijective flat launch; the existing bounded thread policy supplies its width. Plans record elementwise_elements_per_program and elementwise_scalar_temporaries. PlannerOptions::fuse_gpu_elementwise = false retains the same-binary reference control. This extension adds a legal realization family, not a newly fitted cost model or a general fusion-partition solver. CPU vector mapping and reduction/matrix realization are separate. Profitability for many-output, register-heavy graphs and more general affine layouts remains an open planning problem; passing the proof is not a cross-device speed claim.

Guarded native Metal matrix realization

An optional native atom selector consumes the same MMA operation; no extra frontend hierarchy or memory object is required. CompileOptions has an explicit cooperative_matrix capability contract (default off). For Metal, the caller must target a device supporting native FP32 SIMD-group matrices and specify thread_warp_size=32; merely naming metal does not prove that capability. This keeps cross-compilation independent of the host GPU. The initial atom follows Apple’s SIMD-group matrix model introduced with Apple GPU family 7; the target contract must be established before selecting it.

TileIR mma(A, B, C), with typed arithmetic policy
  |
  +-- reference TIRx contraction + MMA provenance marker
  |
  +-- capability + numerical policy + actual body/operand-map proof
       |
       +-- proven: complete 32-worker SIMD groups, 8x8 FP32 matrix atoms
       |     native TIRx matrix load / multiply[/accumulate] / store
       |
       +-- otherwise: the existing checked contraction loops

The first selector handles rank-two FP32 operands/results, constant or independent Tile accumulators, positive strided row-major/transposed shared projections, and M/N/K extents divisible by eight. Global ragged edges are already handled by bounded Tile loads/stores; two-window pipeline slots remain uniform outer address coordinates. The matcher rechecks the actual typed load/multiply/add/store body, not just its marker, so a transform cannot leave stale metadata that changes semantics. Mixed conversions, opaque placed layouts, explicit worker-local execution, insufficient participants, and an ordered MMA policy keep the reference path.

Native TIRx metal.simdgroup allocations and registered matrix operations go through TVMx’s C++ compiler. Each matrix instruction runs on a complete SIMD group; job-tail predicates are uniform within that group, and the following group-wide fence remains outside those predicates. The selector never treats an individual worker’s Tile as a cooperative fragment. Matrix selection is a legality decision, not a profitability claim or a change to default mapping. test_tile_tirx_matrix checks generated instructions and full numerical output on CPU and physical Metal, including transposes, ragged shapes, pipeline versions, numerical/capability fallbacks, stale markers, and zero contraction.

The Metal matrix path now has a bounded-family execution planner: it searches complete-subgroup thread counts and exact two-dimensional subgroup/fragment factorizations. A subgroup can retain several output fragments and reuse A/B fragments between them. A proved closed accumulator recurrence can also stay in native fragments across temporal iterations. Other observers of the carried state prevent that promotion. The plan is a side-table result, not another frontend hierarchy or a serialized program.

The cost model ranks legal plans by relative issue work and fragment pressure; its initial coefficients are priors, not measured nanoseconds or hardware occupancy. Pareto dynamic programming preserves shared-memory/cost tradeoffs between multiple matrix operations. The exact solver is exact only for this implemented finite family and this model, not for all GPU programs. The execution planner design specifies the constraints, layout correspondence, objective, solver, calibration, and extension boundary.

Shared-memory budgeting remains the conservative sum of compact group allocations minus the eliminated result buffers of proved resident recurrences, checked against target capacity. General lifetime-based reuse and selective materialization remain planner work. Unsupported placed layouts, opaque buffer escapes, noncanonical distributed loop steps, and captures of host-local materialized storage are rejected. Unsupported descendant scopes are still diagnosed inside empty groups, which otherwise lower to no-ops. Regressions cover generated shared declarations and full fences, multiple resources, cross-worker and cross-phase communication, ragged shapes, softmax/GEMM numerics, resource exhaustion, and illegal scope mappings.

Implemented native software-prefetch path

The current capture implementation retains cursor cuts as mutable STAGE boundary operations in its single pipeline body. Native export consumes those boundaries into labeled TIRx statement segments; it does not discard them or interpret a source ordinal as a hardware cycle. Ordered child regions remain the canonical TileIR representation planned above, not a claim that this capture representation has already been migrated.

The CPU and Metal bridge currently chooses a safe cut into two scheduling phases. Several source segments can belong to one phase. For example, the load / score / update source stages of attention can become an early load phase and a late score/update phase without moving the recurrence forward.

Source iterations                 Scheduled execution (two-phase example)

i=0: load(0) -> compute(0)         prologue:  load(0)
i=1: load(1) -> compute(1)         steady:    load(1) -> compute(0)
i=2: load(2) -> compute(2)                    load(2) -> compute(1)
                                  epilogue:               compute(2)

cross-phase temporary: slot(i) = i % 2
recurrence acc:         updated only by compute, in original iteration order

This is software prefetching, not a promise of an asynchronous DMA engine or warp-specialized execution. The bridge invokes TVMx’s native C++ InjectSoftwarePipeline pass through a temporary opaque-SBlock adapter, then restores ordinary TIRx AllocBuffer/For/IfThenElse statements before target execution mapping. The public bridge boundary remains native TIRx; there is no Python source, parser round-trip, or MLIR dependency.

Legality is intentionally conservative:

  • Only iteration-local, statically sized, unplaced storage is versioned. Allocation/resource annotations are never lost to a buffer-rebuilding pass.

  • Cross-phase access to an outer resource with any write prevents that cut. This includes loop-carried Scalars/Tiles after materialization and explicit Memory. Carries confined to the consumer phase preserve their original order.

  • External buffers may alias unless CompileOptions::noalias explicitly promises otherwise. Read-only aliasing does not prevent prefetching, but cross-phase read/write or write/write alias hazards do.

  • Buffer live ranges include writes as well as reads, so a late dead write cannot overwrite a newer iteration’s live slot. Yield snapshots still update multiple carried values simultaneously.

  • Extra versions must fit a conservative shared-memory capacity bound before Metal group mapping. If versioning would overflow it, the legal ordered implementation is retained.

PipelinePolicy::window bounds in-flight iterations (0 lets the planner choose; 1 disables iteration overlap), not a source-stage count. This initial planner uses at most two in-flight iterations. It requires PipelinePolicy::interval == 1; other positive intervals retain ordered reference execution until a target latency/issue model is available. Zero intervals and invalid IR policy payloads are rejected. Pipelines without cuts, opaque effects, unsupported placement, or unproved dependencies likewise keep their ordered semantics. General modulo schedules, inferred cuts, and hardware async events are not implemented by this path.

test_tile_tirx_pipeline checks zero/one/short loops, ragged shapes, multiaxis iteration order, multiple carries, aliasing, local/outer Memory, late writes, and capacity fallback on CPU and physical Metal. It also checks preserved native segments and actual doubled Metal storage, so numerically correct but still-serial lowering cannot satisfy the positive pipeline case.

Pipeline and memory bridge

Scheduled TileIR exports:

TileIR

TVM destination

memory declaration

alloc_buffer or target memory object

logical/physical correspondence

TIRx TileLayout when factorizable

view/address map

buffer layout or explicit TIRx index map

execution binding

TIRx scope IDs or TIRx thread/loop bindings

pipeline schedule

software-pipeline annotations or explicit staged loops

async copy/token

target-supported async operation and dependence

semantic tile op

TIRx tile op when available, otherwise decomposed scalar/vector TIRx

The initial bridge is one-way. Round-trip import is not required and must not weaken TileIR invariants.

The bridge’s native compiler driver mirrors the TIRx pass pipeline in C++ and dispatches target.build.<kind> through TVM’s C++ registry. It partitions host and device PrimFuncs by their bound target, finalizes each partition, and imports generated device modules into the host runtime module. A scalar PrimFunc compile-and-execute test is the minimum ABI smoke test; Python is not loaded even for pass orchestration.

The native implementation lowers static-JIT-specialized view arguments, Scalar and Tile constants, typed elementwise operations with named-dimension broadcasting, bounded subtile loads/stores, pure Tile extraction/map, semantic MMA, and parallel/serial/pipeline/reduce with inferred Scalar or Tile carried state. Loads preserve snapshots across later writes; structured yields snapshot every carried value before updating any carry storage.

Pure elementwise expressions can fuse into their consumer. The initial reference schedule materializes loaded Tiles, map results, and MMA results into compiler-owned storage. Worker mapping keeps this private; explicit Metal group mapping shares group-level allocations as described above. MMA retains a checked contraction fallback, with the opt-in Metal matrix matrix realization. Explicit Memory resources preserve MemoryState ordering and old-value snapshots; provable strided/composed layouts map to native address expressions, while unsupported placements fail closed. Safe pipelines use software prefetching, not hardware-asynchronous transfers. General distribution planning, storage reuse, and additional atom families remain work in progress. The compiler distinguishes ordinary TIRx statements (STANDARD) from programs containing native TIRx TilePrimitive calls (TILE), because only the latter require LowerTIRx. Both paths run LowerTIRxOpaque before host/device splitting so thread-binding loops become device regions. Buffer noalias is an explicit caller contract and defaults off until TileIR carries enough alias metadata to prove it.

Bootstrap lowering path

C++ capture
  -> Candidate TileIR
  -> value-to-SSA and structural verification
  -> layout/distribution inference
  -> execution binding and guarded variant selection
  -> pipeline/resource planning
  -> Scheduled TileIR
  -> native C++ TVM TIRx
  -> existing TVM target lowering

Later native passes can replace any segment:

Scheduled TileIR
  -> target atom legalization
  -> Machine TileIR
  -> Luisa XIR / LLVM / native backend IR

The bridge is therefore scaffolding, not an architectural dependency trap.

Target catalog

A target plugin provides data, not frontend syntax:

  • target execution axes and legal bindings;

  • the target-scope containment poset and parent projections;

  • legal region anchors, execution frontiers, and convergence rules;

  • resource kinds, instance topology, capacity, bank geometry, visibility, and coherence;

  • copy, MMA, vector, reduction, and synchronization atoms;

  • atom operand/result layouts and type constraints;

  • engines, events, barriers, and pipeline capabilities;

  • cost hooks and legality predicates.

The same logical hierarchy can target GPU blocks and warps, CPU thread teams and vectors, an accelerator core hierarchy, or a simulator. The verifier rejects a schedule whose resource-instance or visibility maps are illegal.