Paper Reading: TileLoom

This is the slide used in the paper reading circle. Here we focus on the paper TileLoom: Automatic Dataflow Planning for Tile-Based Languages on Spatial Dataflow Accelerators.

Transcript

1. TileLoom: Automatic Dataflow Planning for Tile-Based Languages on Spatial Dataflow Accelerators Reading Circle Yuri Takigawa September 16th, 2026 M1, Taura Laboratory, Creative Information, Graduate School of Information Science Technology, The University of Tokyo
2. Contents Contents 1. Introduction . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ⁠2 2. Background . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ⁠5 3. Framework . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ⁠11 4. Evaluation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . ⁠20 1 / 23
3. 1. Introduction
4. Paper Information • Link: https://www.usenix.org/system/files/osdi26-li-wei.pdf • Operating Systems Design and Implementation 2026 (OSDI’26) • Authors: Wei Li et.al, (National University of Singapore, Arizona State University, Google) • Source code: https://github.com/ecolab-nus/loom-dataflow 3 / 23
5. 40 seconds Summary Since spatial dataflow accelerators (SDA) resolves the traditional bottleneck of global memory access, it often exposes another demand on programs. The performance of program significantly depends on how workloads are mapped to hardware. This paper introduces TileLoom, an MLIR-based e2e framework that compiles tile-programs onto SDA. As a preliminary step to code optimization/generation, TileLoom determines spatiotemporal mapping of the program starting from dataflow-agnostic tile-based user language combined with hardware representation. Figure 1: TileLoom framework overview 4 / 23
6. 2. Background
7. Spatial Dataflow Architecture • 2D grid of cores • each typically a SIMD or a vector engine • sometimes with a matrix unit • plus a local scratchpad • connected by a packed-switched mesh NoC. Figure 2: https://dl.acm.org/doi/epdf/10.1145/3579371.3589348 While granularity and number of cores differ between architectures, they commonly aim to reduce reliance on large shared caches (i.e., incurs bandwidth contention) and highlatency off-chip memories. 6 / 23
8. Demand of another level of programmability (1) Performance depends on • which cores execute which parts of the program. • how data is partitioned across local memories. • how traffic is scheduled on the network. Naive mappings can cause • severe load imbalance • network congestion • underutilized cores • excessive off-chip memory traffic resulting in poor performance and efficiency (cost, energy)1. Figure 3: critical path spans cores in space and time. 1We often uses throughput per box (tok/s/box) as a efficiency metrics. 7 / 23
9. Programming model for parallel accelerators GPU • grid-block-thread model popularized by CUDA interface. • Programmer decomposes the problem into blocks and launches a grid of block instances across the device ‣ each block run the same program on a different region of the input or output tensor ‣ grid covers the full problem domain. • A hardware scheduler dynamically assigns blocks to cores (SM) • Compiler/programmer decides intra-block mapping of program/data onto threads, warps, local resource. Tile-based languages such as Triton let users write programs with tile operators that define the work of a single block, which is lowered onto intra-core/SM resource. 8 / 23
10. Demand of another level of programmability (2) Halide1, TVM2 treats tile shapes and launch configurations as schedule parameters and search over them to improve locality and parallelism. Architectures that distribute programs across large arrays of PEs connected by a highbandwidth on-chip network expose/provide richer inter-core programmability to software. 1. decide how to implement the tile program on a single core 2. how to place tile instances across cores and schedule them in time so that data can be forwarded or multicast efficiently over the NoC and memory system. • inter-core programmability creates an combinatorially large number of different valid mappings. • Different mapping = different load balance, critical path. Vendor offers oracle in compilers, but not generalized for new kernels, models, hardwares. 1https://dl.acm.org/doi/10.1145/2491956.2462176 2https://www.usenix.org/system/files/osdi18-chen.pdf 9 / 23
11. Hypothesis TileLoom starts with • Since hardware exposes explicit cores, networks and memories (reducing hardware arbitration influence), dataflow architectures become more deterministic and predictable. • A compiler can more easily reason about the perfomance of different schedules and deduce good static mappings. Thoughts about the motivation • Strongly Agree Due to inter-core NoC, monolith arbiter/memory is no longer reasonable choice. At the same time, each core is no longer homogeneous in terms of space. • On top of this understanding, specialization would be the key to optimize mapping from high-level program. 10 / 23
12. 3. Framework
13. Phase of compilation 1. Frontend • take tile-level kernel • generate block-shapes candidates. • generates dataflow-agnostic IR for each candidate. 2. Dataflow planning with Architecture representation • determines how this logical grid of tiles is realized on the target architecture. ‣ how (logical) tiles are packed into (physical) cores. – same core = reused over time via buffering – different core = shared over the NoC ‣ where each tile is allocated ‣ when the move occurs. • mapping and placement generates a design space of candidates Figure 4: TileLoom framework overview 12 / 23
14. Phase of compilation 2. Dataflow planning with Architecture representation • Given architecture representation and dataflow plan, the performance model estimates the cost of each candidate. • Filter out top-k candidates from design space. • At the last step of this phase, dataflow-aware IR is generated, which includes ‣ memory allocations, move operations, communication endpoints 3. Block-level optimization and code generation. 4. Profiles surviving top-k candidates on the hardware to choose the final floorplan. 13 / 23
15. How IR would look like (1) dataflow-agnostic IR with block-shaped Figure 5: Block size is embedded, grid size is just named.1 1Named dimension is a common expression for tensor-based IR. 14 / 23
16. How IR would look like (2) dataflow-aware IR with core index Figure 6: In the inner-most loop, per-core program is described using core spatial index. dataflow expression in the IR Figure 7: type="broadcast" annotation represents dataflow 15 / 23
17. Data reuse and memory allocation Interpretation of “data reuse”: A tile can be reused only temporally (loaded once per core and reused across iterations), only spatially (broadcast once and immediately consumed), or in both ways (broadcast once and then reused across several temporal iterations). In all cases, the decision can be viewed as picking when a tile is first loaded or received and how long it remains live in local storage under a fixed loop order. What it takes • Agressive “data reuse” means larger buffer footprints. • In the actual serving, cache footprint is random (depending on requests). Usually the buffer size for cache is statically given first (based on the service profile) at spatial architecture. Residual buffer will guide this reuse decisions.1 1If the total context length of requests exceeds the budget, one of them will be evicted to host. 16 / 23
18. Architecture representation Topology expression Intra-core specification Figure 8: Grid shape, link between cores are represented. Intra-core specification Figure 9: Matrix unit, vector unit has shape, tpt attributes. Figure 10: DRAM is represented like a “special core” without compute capability 17 / 23
19. Performance Modeling My take: Larger the scale of architecture gets, more critical the relative accuracy of PM becomes. Scope of PM in TileLoom: not attempt to exactly simulate the core’s instruction scheduler. Instead, it is calibrated to be accurate enough to distinguish compute-bound from memory-bound mappings and to reason about overlap between compute and data movement. 18 / 23
20. Coarse-grained Performance Modeling Figure 11: each iteration executes as a pipelined load-compute-store sequence with double buffering: while iteration i is computing, the stores for iteration 𝑖 − 1 and the loads for iteration 𝑖 + 1 proceed in parallel whenever possible. Two bucket (core compute, communicate) level-sort preserving memory dependencies.1 1static program order may be preserved in such architecture 19 / 23
21. 4. Evaluation
22. Evaluation Setup • Tenstorrent Wormhole and Blackhole devices with 16-core Intel Xeon Gold CPUs and 512 GB DRAM. Figure 12: Specifications of the Tenstorrent Wormhole and Blackhole. • The (small) scale of hardware could just enable this approach. (design space explosion) 21 / 23
23. E2E performance FlashAttention Figure 13: TileLoom consistently outperform TTNN with various batch size, seq length, n heads. FlashDecode Mostly underperforms TTNN (reference). Paper describes this as “specialized optimization on the reference side”. GEMM, Irregular-size matmul Similar or better performance against irregular size matmul, GEMM is similar level with existing framework. 22 / 23
24. Performance Modeling Accuracy In the exploration made in this paper, this level of miscorrelation doesn’t harm the quality of candidate selection. Figure 14: Perf-model accuracy. 23 / 23