PaperlayerView exact revision on arXiv

Source-linked AI summary

Schedules Are Solvable Symbols: Tuning-Free Compilation of Tile Programs on Dataflow Architectures

Heru Wang, Wei Li, Zhenyu Bai, Tulika Mitra

arXiv:2609.29219v1cs.PL

TL;DR

Tile-based SPMD compilers and vendor libraries do not provide portable, joint optimization of inter-core and intra-core dataflow, despite data movement being a central bottleneck for regular kernels. Loom uses hardware-derived symbolic optimization to jointly select dataflow, scheduling, and block sizes. It matches or exceeds TTNN across GEMM, Flash Attention, and Flash Decode on two Tenstorrent generations without per-shape profiling.

  • Problem

    For regular kernels, data movement is a central performance and energy bottleneck, but existing vendor libraries and tile-based DSLs do not provide portable, joint optimization of inter-core and intra-core dataflow.

    Source: §1 · p. 1

  • Method

    Loom derives symbolic legality constraints and latency expressions from hardware descriptions, then jointly optimizes inter-core dataflow, intra-core asynchronous scheduling, and block sizes at compile time.

    Source: §1 · p. 2

  • Results

    Loom matches or exceeds TTNN across GEMM, Flash Decode, and Flash Attention, with 2.15×–2.29× geometric-mean speedup on Flash Attention.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    On Wormhole and Blackhole, hardware-parameterized compilation produces out-of-the-box results across three operators without per-kernel profiling.

    Source: §1 · pp. 2–3

    The analytical model abstracts away fixed kernel overhead, pipeline warm-up, synchronization cost, imperfect utilization, and difficulty reaching peak throughput on small or skinny matrices, so it can overestimate underutilized configurations.

    Source: §3.4.3 · p. 12

Abstract

from arXiv · show

Modern AI and HPC accelerators increasingly expose dataflow features: software-visible mechanisms for data movement and overlap, such as inter-core communication through the on-chip network and intra-core asynchronous pipelining. These features shift scheduling responsibility from hardware to the compiler, and because placement, movement, and synchronization become software-visible, they also make the performance of static schedules predictable. Yet high performance on such hardware still relies on vendor-engineered kernel libraries or profile-based auto-tuning, whose embedded expert knowledge transfers poorly across architectures and algorithms. We present Loom, a tuning-free symbolic compiler framework for tile-based SPMD programs on spatial dataflow architectures. The central idea is to treat tile-based SPMD compilation as a hardware-explicit static optimization problem. Loom enumerates discrete spatial-mapping and communication candidates while keeping value parameters, such as tiling factors and pipeline knobs, symbolic within each candidate. From an explicit hardware description, it derives symbolic legality constraints and latency expressions, formulates one CP-SAT problem per schedule candidate, and jointly solves inter-core dataflow, intra-core asynchronous scheduling, and block sizes at compile time. On two Tenstorrent generations, Wormhole and Blackhole, Loom matches or exceeds the vendor-optimized TTNN library on GEMM, Flash Attention, and Flash Decode, out of the box and without per-shape profiling or profile-based platform-specific schedule tuning. These results suggest that hardware-derived symbolic compilation provides a retargetable alternative to profiling-based tuning for spatial dataflow architectures while remaining interpretable by keeping optimization decisions traceable to source-level symbols.

1 Introduction

Loom treats tile-based SPMD compilation as hardware-explicit symbolic optimization, jointly choosing inter-core dataflow, intra-core scheduling, and block sizes. It matches or exceeds TTNN across three operators on Wormhole and Blackhole without per-kernel profiling.

1 Introduction

  • For regular tensor and stencil kernels, data movement—not computation—is the central performance and energy bottleneck; spatial dataflow features expose communication and execution overlap to software.These features include inter-core forwarding and multicast, plus intra-core asynchronous transfers, double buffering, and producer–consumer synchronization.Source: §1 · p. 1
  • Vendor libraries encode platform-specific expert knowledge, while tile-based DSLs rely on profiling or heuristics and leave spatial-dataflow opportunities implicit.The result is that neither path provides portable, joint optimization of inter-core and intra-core dataflow.Source: §1 · p. 2
  • Tile-based SPMD structure can support planning both grid-level placement and forwarding and block-level asynchronous movement, buffering, and overlap.The regular grid and heterogeneous block program expose the structure for these two planning levels.Source: §1 · p. 2
  • Inter-core choices constrain local buffers and pipeline stages, while intra-core schedules change memory needs, overlap, and the best placement and forwarding pattern.Hardware capacity, topology, bandwidth, synchronization costs, and compute throughput determine which combinations are legal and profitable.Source: §1 · p. 2
  • Loom keeps block sizes and pipeline decisions symbolic, derives legality constraints and latency expressions from hardware descriptions, and jointly optimizes dataflow and scheduling without profiling in the loop.Its architecture-parameterized hardware knowledge is intended to make retargeting a new generation require an updated description rather than kernel re-engineering.Source: §1 · p. 2
  • Against TTNN, Loom's geometric-mean speedups are 0.99×/1.16× on GEMM and 1.01×/1.17× on Flash Decode (Wormhole/Blackhole), and 2.15×–2.29× on Flash Attention.Evaluated on Wormhole and Blackhole out of the box.Source: §1 · pp. 2–3

2 Framework

Loom’s framework enumerates hardware-aware schedule candidates for tile-based SPMD programs, keeping value parameters symbolic while modeling legality and latency from reusable hardware properties. It selects the lowest-latency feasible schedule, making dataflow, resource use, and scheduling decisions traceable to source-level symbols.

2.1 Overview

  • Loom lowers tile-based SPMD code through dataflow-aware representations, enumerates hardware-aware schedule candidates, and selects the feasible candidate with the lowest modeled latency.It derives symbolic latency expressions and legality constraints for each candidate, solves them with CP-SAT, then substitutes the chosen values and generates target code.Source: §2.1 · p. 3

2.2 Symbolic MLIR

  • Loom’s symbolic MLIR extension carries named compile-time unknowns through the IR until CP-SAT assigns their values.Symbols can appear in shapes, loop bounds, indices, memory footprints, and performance-model inputs; source-level labels such as tile_m and tile_n remain traceable in the solver and compiled variant.Source: §2.2 · p. 4

2.3 Reifying Tensor Values into Synchronized Buffers loom.vb

  • Loom reifies tensor dependences as virtual buffers with explicit readable and writable states, enabling asynchronous producer-consumer handoffs while deferring physical storage and synchronization choices.Consumers acquire and release readable data; producers reserve writable capacity and commit produced data for downstream consumers.Source: §2.3 · pp. 4–5

2.4 Reuse-Aware Virtual Buffer Concretization

  • Loom conserves local-memory resources by merging virtual buffers across exclusive handoffs and reusing physical storage when lifetimes permit.It groups shape-compatible buffers and colors an interference graph to share storage; minimizing allocations matters because local-memory capacity constrains feasible block sizes.Source: §2.4 · p. 5

2.5 Topology-Aware Intermediate Transformation

  • Loom maps program axes onto topology-aware core views, prunes dominated mappings, and rewrites eligible cross-core reuse into NoC-aware transfers.Its dominance-based axis ordering retains alternatives under parallel-extent and reuse metrics while eliminating up to 93% of legal candidate space for GEMM mappings.Source: §2.5 · p. 6

2.6 Latency Modeling

  • Loom uses a Hierarchical Execution Graph to compose symbolic compute and memory latencies, with solver-controlled pipeline boundaries representing potential overlap.The model targets regular SPMD grids with shared execution structure; irregular or asymmetric grids require regularization or position-dependent latency modeling.Source: §2.6 · p. 7
  • For each retained schedule candidate, Loom builds a symbolic latency expression, enforces architecture and program legality, solves the constrained optimization, and selects the feasible candidate with lowest modeled latency.Legality includes resource limits and candidate-specific requirements.Source: §2.6 · pp. 7–8
  • Legality predicates cover local-memory capacity and pipeline-buffer footprint, alignment, divisibility, primitive validity, valid communication patterns, and tiling bounds.These combine architecture constraints with program and candidate-specific requirements.Source: §2.6 · p. 8

2.7 Architecture Description

  • Kernel developers specify operators without placement, communication, or pipelining; hardware structure comes from specifications, while performance parameters come from isolated microbenchmarks reused across operators and input shapes.These architecture parameters are collected once per architecture, outside the per-kernel optimization loop.Source: §2.7 · p. 8
  • Loom targets dataflow architectures with sufficiently explicit, predictable placement, communication, synchronization, and pipelining, and its resource-based model can extend to heterogeneous units, hierarchical memories, and multi-level interconnects.Less deterministic behavior can use conservative or offline-calibrated resource models without per-kernel profiling.Source: §2.7 · p. 8

3 Evaluation

Across GEMM, Flash Attention, and Flash Decode on Wormhole and Blackhole, Loom outperforms TileLoom and matches or exceeds TTNN. Its solver usually selects near-best spatial mappings and block sizes, despite modestly higher compilation cost than TileLoom.

3.1 Experimental Setup

  • The evaluation uses Wormhole and Blackhole servers and covers GEMM, fused Flash Attention, and memory-dominated Flash Decode.These workloads represent compute-bound, fused-attention, and memory-dominated regimes.Source: §3.1 · pp. 8–9

3.2 Kernel Performance

  • Across GEMM, Flash Attention, and Flash Decode, Loom outperforms TileLoom and matches or exceeds TTNN on Wormhole and Blackhole.Table 2 reports geometric-mean speedups over TTNN.Source: Table 2 · p. 9, §3.2 · p. 9, §3.2.1 · p. 9
  • GEMM reaches 0.989× and 1.160× TTNN performance on Wormhole and Blackhole; Flash Attention reaches 2.29× and 2.15×, respectively.For GEMM, the Blackhole result is 1.160× relative to TTNN; the Wormhole result is 0.989×.Source: Table 2 · p. 9, §3.2.1 · p. 9, §3.2 · p. 9
  • Flash Decode achieves 1.01× and 1.17× geometric-mean speedups over TTNN on Wormhole and Blackhole across all evaluated shapes.TileLoom does not support all evaluated shapes; on the common subset, Loom improves performance from 0.84× to 0.98× on Wormhole and 0.87× to 0.99× on Blackhole.Source: Table 2 · p. 9, §3.2 · p. 9

3.3 Compilation Cost

  • Loom's top-1/top-5 compilation takes 6.74 s/34.21 s versus TileLoom's 5.75 s/27.66 s, while average speedups rise from 0.948×/1.029× to 0.989×/1.089×.The paper characterizes Loom's compilation cost as 17.2%–23.7% higher but on the same scale as TileLoom's.Source: Table 3 · p. 11, §3.3 · p. 11

3.4 Ablation Studies

  • For most shapes, Loom ranks the best measured mapping first among the top five: the interval’s upper end coincides with 1.0; largest deviations occur on three leftmost, highly skewed shapes.Figure 10’s measured performance range is normalized to the solver-ranked top-1 candidate.Source: §3.4.1 · p. 11, Figure 10 · p. 11
  • Solver-selected block sizes are locally optimal or nearly so for almost all GEMM shapes; only a few neighboring configurations are slightly faster.Under a fixed spatial mapping, the model captures the trend that larger feasible blocks often improve utilization, while capacity and alignment constrain choices.Source: Figure 11 · p. 12, §3.4.2 · p. 12
  • The model preserves major throughput regimes but can overestimate underutilized configurations because it abstracts fixed overhead, warm-up, synchronization, and imperfect utilization.Its purpose is to rank and prune candidates rather than predict cycle-accurate throughput.Source: Figure 12 · p. 12, §3.4.3 · pp. 12–13

4 Related Work

Related work spans analytical hardware–software modeling, spatial and tile-level compilation, and symbolic schedule optimization. Loom combines these lines at the SPMD-block level for executable spatial-dataflow programs.

Analytical modeling for hardware–software co-design

  • Spatial and accelerator-modeling systems make data movement, reuse, and resource limits explicit, while Loom derives symbolic legality constraints and latency terms for executable tile programs.The cited systems primarily support hardware generation, design-space exploration, or mapping evaluation.Source: §4 · p. 13

Compilation for spatial and tile-level programs

  • CGRA, spatial-dataflow, and tile-level compilers address placement and communication, but Loom jointly reasons about symbolic block sizes, topology-aware mapping, NoC communication, buffers, and intra-core overlap [48] [50].The passage characterizes many tile-level compilers as relying on tuning, learned cost models, or user-provided schedules.Source: §4 · p. 13

Symbolic and solver-based schedule optimization

  • Analytical tile-size selection, polyhedral optimization, and solver-based mapping represent schedule choices mathematically; Loom applies symbolic tiling and pipeline variables at the SPMD-block level.Its hardware-parameterized latency model makes the solver an interpretable bridge between compiler IR and hardware-derived performance modeling.Source: §4 · p. 13

5 Conclusion

Loom makes hardware constraints explicit in symbolic compilation of tile-based SPMD programs. On Wormhole and Blackhole, it generates implementations that match and often outperform vendor-library kernels across three workloads.

Conclusion

  • Loom exposes placement, NoC communication, local-buffer capacity, and intra-core pipelining to an analytical model and solver instead of relying only on fixed templates, heuristics, or profile-based tuning.The compiler keeps part of schedule decisions symbolic for tile-based SPMD programs on spatial dataflow architectures.Source: §5 · p. 13
  • On Wormhole and Blackhole, Loom matches and often outperforms vendor-library kernels across GEMM, Flash Attention, and Flash Decode.The results suggest hardware-derived symbolic compilation can provide an interpretable, retargetable path toward high-performance code generation for accelerators with software-controlled dataflow mechanisms.Source: §5 · p. 13
Loading 2609.29219v1…