A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

arXiv:2608.13757v1 [cs.PL] 13 Aug 2026

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators Chungha Sung∗

Nikil V. Shyamsunder∗†

Hanliang Zhang

Amazon, Inc. USA chunghs@amazon.com

Cornell University USA nvs26@cornell.edu

Amazon, Inc. USA hlz@amazon.com

Daniel Kroening

Joonwon Choi

Amazon, Inc. USA dkr@amazon.com

Amazon, Inc. USA jwch@amazon.com

Abstract

1

Multi-engine AI accelerators such as AWS Trainium comprise specialized compute engines that execute in parallel, and the compiler must synchronize the data dependencies between them. For straight-line code this is simple: each dependency reduces to waiting for a threshold count of instruction completions, which the compiler computes statically. Loops admit no such static threshold; a simple solution inserts all-engine barriers at iteration boundaries, resetting synchronization state so each loop body can be treated as straight-line, at the cost of parallelism. We present a barrier-free synchronization algorithm that instead enforces each dependency precisely across structured control flow with arbitrarily nested, dynamically bounded loops. The key idea is to compute dynamic thresholds at runtime from tracked loop iteration counts. We implemented it as a compiler backend pass at the AWS Neuron ISA level. On a suite of ML kernels, it reduces latency 10–45% relative to the barrier-based baseline, achieves a 3.3× speedup on a synchronization-bound microbenchmark, and often matches or exceeds hand-tuned manual allocation. Issuing a consumer too early violates its dependency, while issuing too late unnecessarily stalls execution. We formally characterize the minimum synchronization required for correctness and verify in the Lean proof assistant, via bisimulation, that our algorithm meets this criterion.

AWS Trainium [4, 6, 9] is a family of AI accelerators comprising multiple specialized compute engines that share onchip buffers and execute concurrently [7]. Each engine runs its own instruction stream asynchronously, and within an engine instructions are deeply pipelined, so a later instruction may begin before an earlier one completes [7]. Data dependencies therefore arise both across engines (e.g., a DMA load must complete before the Tensor Engine consumes the data) and within an engine. AWS’s Neuron compiler [5] analyzes dependencies and enforces them through hardware semaphores: shared atomic counters incremented when an instruction completes and tested against a threshold before issue. For straight-line code, each engine’s semaphore counts instruction completions, and a dependent instruction waits for it to reach a statically computed threshold—its wait value. This is simple, efficient, and correct, but only applies to acyclic code. Loops break this static-threshold scheme: the same instruction completes many times, so a fixed ordinal threshold no longer identifies a single completion. Unrolling would restore it at the cost of code size, but this requires statically known trip counts, infeasible for modern ML workloads with dynamically bounded loops (e.g., variable-length attention, Mixture-of-Experts). A simple alternative synchronizes all engines at every loop entry and loop iteration boundary via all-engine barriers: each engine drains its in-flight instructions, performs a handshake, and resets all semaphores to zero before proceeding (§3.4). This is reliable, but barriers serialize execution even when no cross-iteration dependencies exist, preventing overlap across iterations and leaving hardware idle while faster engines wait for the slowest. We propose a barrier-free semaphore allocation algorithm that replaces these barriers with precise, per-dependency wait conditions. It operates on structured control-flow graphs (SCFGs) with arbitrarily nested, dynamically bounded loops and conditionals. When all engines traverse a common SCFG, each engine tracks the iteration counts of its enclosing loops in local registers and computes each wait value in closed form

CCS Concepts: • Software and its engineering → Compilers; • Theory of computation → Program verification; • Computer systems organization → Parallel architectures. Keywords: semaphore allocation, multi-engine synchronization, barrier elimination, bisimulation, formal verification, AI accelerator, Lean proof assistant

∗ Both authors contributed equally to this research. † Work done while at Amazon.

1

Introduction

Sung, Shyamsunder, Zhang, Kroening, and Choi

from those counts and the dependency structure. Each engine therefore stalls only for its actual dependencies—never for unrelated instructions on other engines. This performance boost comes with risk: an incorrect semaphore assignment or wait-value computation causes silent data corruption or deadlock, and such bugs are timingdependent, making them hard to reproduce. We therefore start by defining when a dependency is satisfied—the precise moment at which a consumer can safely issue without oversynchronizing. We claim our semaphore allocation targets exactly this moment, allowing each instruction to issue if and only if the dependency is satisfied. We formally prove this claim in the Lean proof assistant via bisimulation [13] in approximately 15K lines with zero admitted lemmas. The proof handles loop-carried dependencies and nested control flow, and is stated for our per-loop allocation, where multiple instructions in the same loop body share a single semaphore. We implement both a barrier-based baseline and our new barrier-free algorithm as compiler backend passes at the AWS Neuron ISA level, each taking a kernel plus a dependency graph and emitting a kernel with semaphores allocated. We work at the ISA level and hold the instruction schedule fixed, isolating the standalone effect of the allocation strategy. On a microbenchmark isolating the barrier’s cost, the barrierfree allocation achieves a 3.3× speedup over the barrierbased baseline, and across a suite of ML kernels compiled from production NKI kernels [8]—tiled matrix multiplication, RMSNorm, FlashAttention, AdamW, and Dropless MoE—we observe latency reductions of 10–45%, matching or exceeding hand-tuned manual allocation on most kernels. To the best of our knowledge, this work makes the following novel contributions:

instructions on different engines with a single inter-engine loop-carried dependency: 𝐼 2 (on Engine 2) in iteration 𝑖 waits for 𝐼 3 (on Engine 1) from iteration 𝑖−2. The timeline omits control-related instructions (e.g., branches) and shows one possible execution in which each engine executes one instruction at a time in program order. Using the barrier-based approach (Figure 1(a)), the loop entry and every iteration boundary have all-engine barriers, shown as dashed vertical lines. Only one dependency links two engines across an iteration distance of two, yet every iteration boundary forces all three engines to drain and handshake. The pre-loop and post-loop work on Engine 0 is pushed outside the loop region, and Engine 0 sits idle for the entire loop. Each iteration in the loop, Engine 2 waits for Engine 1 due to the barrier—even if Engine 2 is idle. The barrier-free allocation (Figure 1(b)) enforces only the single real dependency—the one red arrow, from the first completion of the producer 𝐼 3 to the third issue of the consumer 𝐼 2 . All other engine activity overlaps freely: the preloop and post-loop work on Engine 0 runs back-to-back while the loop executes on Engine 1 and Engine 2, and the loop iterations overlap across Engine 1 and Engine 2. The speedup grows with loop trip count, since every eliminated barrier is overhead the barrier-based approach paid unconditionally.

[Engine 0] 𝐼 0 [Engine 1] 𝐼 1 loop ...: [Engine 2]𝐼 2 [Engine 1]𝐼 3 [Engine 2] 𝐼 4 [Engine 0] 𝐼 5

• A precise characterization of dependency satisfaction on multi-engine architectures—the minimum synchronization required for correctness (§4). • A barrier-free semaphore allocation algorithm that achieves this minimum and supports arbitrarily nested, dynamically bounded loops (§5), together with an analysis of the resource tradeoffs across semaphoreallocation granularities (§6). • A machine-checked bisimulation proof in Lean for the allocation (∼15K lines, 0 admitted lemmas) (§7). • An evaluation on Trainium showing a 3.3× speedup on a microbenchmark isolating the barrier’s cost and 10–45% latency reductions across a suite of ML kernels relative to a barrier-based baseline (§8).

2

Engine 0 Engine 1 Engine 2 Loop-carry (2) Engine 0 Engine 1 Engine 2

𝐼0

𝐼5

𝐼1

𝐼3

𝐼3

𝐼3

𝐼2

𝐼2

𝐼2

𝐼4

(a) Barrier-based

𝐼0 𝐼1 𝐼2

𝐼5 𝐼3 𝐼2

𝐼3 𝐼2 𝐼4

𝐼3

(b) Barrier-free

Figure 1. Motivating example. A kernel with three engines and one cross-engine loop-carried dependency (offset 2). Dashed vertical lines in (a) mark all-engine barriers at every iteration boundary; in (b), a single red edge enforces the true dependency, letting all other engine work overlap.

3

Background

We now review the Trainium architecture: multiple compute engines, instruction pipelines, and semaphore primitives that both the barrier-based and barrier-free approaches build on.

Motivating Example 3.1

Figure 1 shows a small kernel that runs on three engines and highlights the cost of barrier-based synchronization. Each line of the code assigns an instruction to an engine— for example, 𝐼 0 runs on Engine 0. The loop body runs for three iterations in this example, and its body contains two

Trainium Multi-Engine Architecture

Figure 2 illustrates the NeuronCore, the fundamental compute unit of Trainium [7]. It comprises heterogeneous compute engines (e.g., tensor, vector, and scalar engines) interacting with the on-chip SRAM buffers (SBUF and PSUM) 2

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

and a DMA engine for accessing High Bandwidth Memory (HBM) [4, 6, 9]. The different engines can work asynchronously on independent parts of a kernel at once, provided the data dependencies between them are correctly enforced.

In the Trainium architecture [7], each engine’s instruction stream consists of two interleaved instruction types: datapath instructions and control instructions. datapath instructions are the computational and memory operations (tensor contractions, vector or scalar arithmetic, and DMA transfers). Each engine is deeply pipelined: multiple datapath instructions per engine can be in-flight at once and their executions can be overlapped. Datapath instructions issue to this pipeline and retire in-order. Control instructions handle control flow (branches, loops) and register manipulation (arithmetic on registers). Registers are engine-local, whereas semaphores (§3.3) are globally accessible across engines. They execute atomically and are interleaved with datapath instructions in the instruction stream. We make the precise ordering assumptions of this execution model explicit when we formalize it in §7.

an update, which atomically modifies the counter when an instruction retires. How the underlying dependency graph is obtained is orthogonal to our work, which takes the graph as input and encodes it with semaphores. The conventional encoding scheme uses a wait condition (sem >= threshold) and an increment-at-retirement (sem += inc), so each semaphore acts as a monotone counter of retirements. The threshold can be an immediate constant or a register, enabling dynamically computed wait conditions. Encoding a dependency is simple for straight-line code— the consumer waits for the producer to have retired by setting the threshold to the producer’s ordinal position within its engine’s instruction stream. Figure 3 shows the code where one semaphore is allocated to each engine and each instruction increments its engine’s semaphore on retirement. Engine 0’s semaphore 𝑠 0 counts its retirements, so the dependency where 𝐼 2 depends on 𝐼 1 is enforced by the wait s0 ≥ 2, since 𝐼 1 is the second instruction in Engine 0’s instruction stream. Thus, 𝐼 2 issues only once 𝐼 1 retires. This scheme breaks down with loops: the wait value must account for which iteration the producer and consumer are in. We precisely specify dependency satisfaction in §4.

3.3

3.4

3.2

Instruction Types and Execution Pipelines

Semaphore-Based Synchronization

When a loop’s trip count is statically known it can be unrolled, restoring the straight-line scheme of §3.3; but when the trip count is determined at runtime, unrolling is no longer possible. One approach that handles both cases is to insert an all-engine barrier (SYNC) at every loop entry and iteration boundary. Each barrier drains every engine’s in-flight instructions, performs a cross-engine handshake, and resets all semaphores to zero, so each iteration again uses statically computed wait values (Figure 4).

Trainium provides a synchronization primitive that implements a semaphore as a mechanism to enforce dependencies. A semaphore is a shared 32-bit primitive that supports two operations: a wait, which blocks an instruction’s issue until the counter satisfies a comparison against a value, and

Partial Sum Buffer (PSUM)

Tensor

Vector

Scalar

Engine 0

GPSIMD

State Buffer (SBUF)

𝐼0 {𝑠 0 ++} SYNC loop A: 𝐼 1 {𝑠 0 ++} 𝐼 2 {𝑠 0 ++} SYNC

DMA

High Bandwidth Memory (HBM)

Figure 2. Trainium’s NeuronCore architecture. Four compute engines share on-chip buffers SBUF and PSUM, while DMA moves data between SBUF/PSUM and HBM. Engine 0 𝐼 0 {𝑠 0 ++} 𝐼 1 {𝑠 0 ++}

A Barrier-Based Approach

Engine 1 SYNC loop A: {𝑠 0 ≥ 2} 𝐼 3 {𝑠 1 ++} SYNC

Figure 4. Barrier-based synchronization for a sameiteration dependency where 𝐼 3 depends on 𝐼 2 . Because SYNC drains and resets all semaphores at every iteration boundary, the wait threshold 𝑠 0 ≥ 2 can be computed statically.

Engine 1 {𝑠 0 ≥ 2} 𝐼 2 {𝑠 1 ++} 𝐼 3 {𝑠 1 ++}

While easy to show correct, this approach has two fundamental limitations. First, it over-synchronizes: every engine waits at every barrier, even those with no dependency in that iteration, so no engine can run ahead into the next iteration or post-loop instructions. Second, the drain, handshake, and reset are themselves costly, and this cost is paid

Figure 3. Straight-line semaphore allocation. Each instruction increments its engine’s semaphore on retirement. 𝐼 2 depends on 𝐼 1 : the wait 𝑠 0 ≥ 2 blocks 𝐼 2 until 𝐼 1 retires. 3

Sung, Shyamsunder, Zhang, Kroening, and Choi

once per iteration. This motivates our barrier-free approach, which replaces these coarse-grained barriers with precise, per-dependency wait conditions that allow full engine parallelization.

4

The consumer may issue once producer has retired in every common-loop iteration up to the target 𝑡®. The target underflows, written ∃𝑗 : 𝑡 𝑗 ≤ 0, when some component is non-positive: it names an iteration before the first iteration of that loop, so no producer iteration is owed and the consumer issues with no wait. Otherwise it waits for the cumulative producer executions up to 𝑡®. The dependency is thus satisfied if and only if

Defining Dependency Satisfaction

We define when a dependency is satisfied: the earliest point at which the producer has completed enough times for the consumer to issue. This is the minimum synchronization required for correctness: issue earlier and the consumer may execute before the producer, violating the dependency; later and the consumer waits longer than necessary. 4.1

∃𝑗 : 𝑡 𝑗 ≤ 0 ∨ | {z } no-wait

∑︁

𝐻𝑒 (𝑖®, 𝑃)

.

(1)

𝑖® ⪯ 𝑡®

|

{z

}

cumulative producer retirements through 𝑡®

Here ⪯ orders common-loop iterations lexicographically from the outermost loop inward, and each 𝑖® contributes 𝐻𝑒 (𝑖®, 𝑃), the number of producer executions in that iteration. Figure 5 traces the condition on a loop-carried dependency. Producer and consumer both sit in the inner loop B of A (so 𝑃 = B), with offset 2 on B: the consumer at (𝑎, 𝑏) reads what the producer wrote at (𝑎, 𝑏−2), giving 𝛿 = (0, 2). B runs 3 times under A=1 and 4 under A=2, so its trip count is not fixed. Since the producer has no loop nested below B, it executes once per B-iteration, so 𝐻𝑒 (𝑖®, B) = 1 for every 𝑖® and the sum simply counts iterations up to 𝑡® (if the producer nested deeper, each term could exceed 1). For the consumer at (2, 4), the target is 𝑡® = (2, 2) and the sum materializes as

Program and Execution Model

We target programs with structured control flow, modeling every construct as a loop: an ordinary loop runs its trip count, conditional branches run 0 or 1 times, and the program is a top-level loop that runs once. This yields a tree of nested loops, where each non-top-level loop has a direct enclosing parent loop. Following classical loop-nest dependency theory [3], an instruction 𝑖 is indexed by its iteration vector 𝑣®𝑖 : the iteration counts of its enclosing loops, outermost first. 4.2

𝑅(producer) ≥

The Dependency-Satisfaction Condition

A dependency makes a consumer consumer wait for a producer producer. Following classical loop-nest theory [3], we describe the dependency between producer and consumer by a distance vector 𝛿 over their common loops: one component per common loop giving the iteration distance on that loop. If the consumer’s current iteration vector is 𝑣®consumer , the producer iteration it depends on is the target 𝑡® = 𝑣®consumer − 𝛿. 𝑡® = 𝑣®consumer −𝛿 is only well-defined when the producer and consumer are interpreted over a common structured controlflow graph: both engines run the same loop nest, so 𝛿’s components refer to loops shared by producer and consumer. The trip counts of those loops may still differ across engines. We therefore assume a common SCFG throughout the rest of the paper and formalize it in §7.1. Whether the dependency is satisfied is decided from the target 𝑡® and two quantities: the producer’s retirement count and the loop-entry counts.

∑︁ 𝑖® ⪯ (2,2)

𝐻𝑒 (𝑖®, B) = 𝐻𝑒 ((1, 1)) + 𝐻𝑒 ((1, 2)) + 𝐻𝑒 ((1, 3)) | {z } A=1: 3

+ 𝐻𝑒 ((2, 1)) + 𝐻𝑒 ((2, 2)) = 5, | {z } A=2 through 𝑏=2: 2

so the consumer at (2, 4) may issue once the producer has retired 5 times.

5

The Allocation Algorithm

We now compile the dependency-satisfaction condition in (1) into semaphore waits and updates, achieving the minimum synchronization for correctness. Two architecture facts shape the compilation: semaphores are globally accessible while registers are engine-local (§3), and each instruction carries at most one semaphore update, applied at retirement. We restrict to a practical subset of dependencies whose check reduces to a comparison against two registers and one shared semaphore, computed at issue time. We present the intuition for the resulting closed-form wait values here; their exact correspondence to the dependency-satisfaction condition is established in the machine-checked proof of §7.

Producer retirements. 𝑅(producer) is the total number of times producer has retired; it is monotone. Since producer appears once in the body of its parent loop 𝑃 (the loop immediately enclosing it) and retires once per issue (§7.1), 𝑅(producer) equals the number of completed entries of 𝑃.

5.1

Loop-entry counts. 𝐻𝑒 (𝑖®, ℓ) is the number of times engine 𝑒 enters loop ℓ during the common-loop iteration 𝑖®, where ℓ lies at or below the innermost loop 𝑖® names. This count is 1 when ℓ is the innermost common loop, and is equal to ℓ’s trip count when ℓ is nested below it.

From the Condition to Registers

The condition in (1) is a disjunction of two checks: a nowait check ∃𝑗 : 𝑡 𝑗 ≤ 0, and a cumulative-retirement check Í 𝑅(producer) ≥ 𝑖®⪯𝑡® 𝐻𝑒 (𝑖®, 𝑃). The producer’s retirement count 𝑅(producer) is the one cross-engine quantity in the condition, so we carry it in a semaphore, which is globally accessible 4

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

Engine 0

Engine 1

loop A: loop B: Producer # (a ,b)

loop A : loop B : Consumer # (a ,b -2)

𝑣®consumer (1, 1) (1, 2) (1, 3) (2, 1) (2, 2) (2, 3) (2, 4)

𝛿 (0, 2) (0, 2) (0, 2) (0, 2) (0, 2) (0, 2) (0, 2)

𝑡® (1, -1) (1, 0) (1, 1) (2, -1) (2, 0) (2, 1) (2, 2)

® 𝑖®⪯𝑡® 𝐻𝑒 (𝑖 , B) underflow underflow { (1,1) } = 1 underflow underflow { (1,1), (1,2), (1,3), (2,1) } = 4 { (1,1), (1,2), (1,3), (2,1), (2,2) } = 5 Í

Satisfaction no wait no wait 𝑅 (producer ) ≥ 1 no wait no wait 𝑅 (producer ) ≥ 4 𝑅 (producer ) ≥ 5

Figure 5. Loop-carried dependency, offset 2 on B (𝛿 = (0, 2), 𝑃 = B). B runs 3 times under A=1 and 4 under A=2. Each B-entry contributes 𝐻𝑒 (𝑖®, B) = 1, so the sum counts iterations 𝑖® ⪯ 𝑡®. A target with a non-positive component underflows, and the dependency is satisfied with no wait. Cumulative retirements via monotone register. The monotone register 𝑟 ℓ incremented at every entry of ℓ and never reset, counting the total entries of ℓ.

and atomically incremented when the producer retires (§3.3). The thresholds it is compared against are loop-entry counts that every engine tracks locally, so we hold those in registers. Computing these thresholds at runtime is what costs registers: each is built from the per-iteration entry counts 𝐻𝑒 (𝑖®, ℓ), which an engine tracks by incrementing a register as it enters loops. One loop’s entry count is one register, but 𝐻 is indexed by the full iteration vector—an inner loop’s count depends on which iteration of every enclosing loop it falls in. This means that to support all possible dependencies for a loop nested under 𝑑 enclosing loops, we must track a separate entry count for each combination of the enclosing loops’ iterations, growing exponentially in 𝑑, so materializing 𝐻 in registers is infeasible. Instead, we restrict to two per-loop counters and confine the distance vector’s only nonzero component to the innermost common loop between producer and consumer. Reducing dependency structure to a scalar distance on a shared loop is common in loop optimizations: modulo scheduling, for instance, models each loop-carried dependency by a scalar distance on the scheduled loop [35]. Let 𝑆 be the innermost common loop of producer and consumer, and 𝑃 the producer’s parent loop—either 𝑆 itself or a loop nested inside it. We require the dependency to be carried at 𝑆: producer and consumer agree on every common loop strictly above 𝑆, so the distance vector has a single nonzero component there, the offset 𝑘 ≥ 0. The gap is now one scalar, and each check reduces to a register comparison. No-wait check via trip register. The trip register 𝑟ˆℓ incremented at every entry of ℓ but reset to 0 whenever ℓ’s parent is re-entered, so it counts the entries of ℓ in the current trip through its parent. The only level whose target can underflow is 𝑆, so the no-wait check becomes a single comparison 𝑟ˆ𝑆 ≤ 𝑘. It holds when at most 𝑘 iterations of 𝑆 have run in the current trip: the offset then points before this trip began, so no producer is owed. When every dependency carried at 𝑆 has offset 0, the no-wait check 𝑟ˆ𝑆 ≤ 𝑘 becomes 𝑟ˆ𝑆 ≤ 0, which is false whenever the consumer issues (the loop has been entered, so 𝑟ˆ𝑆 ≥ 1), so we omit 𝑟ˆ𝑆 for that loop.

5.2

The Allocatable Cases

Carried at 𝑆, the cumulative count depends only on where the producer sits relative to 𝑆. The dependency is forward [3] if consumer follows producer in 𝑆’s body and backward otherwise. If producer sits inside a loop 𝐵 nested in 𝑆, we use 𝐵’s position in 𝑆’s body as producer’s (likewise for consumer). Because producer and consumer may run on different engines, this order is forced only when a nested loop or block boundary separates them; when they share the same straight-line region of 𝑆 from different engines, we take the dependency to be forward. This leaves three cases, by where producer sits relative to 𝑆 and the order of the dependency. • Producer in 𝑆 (𝑃 = 𝑆), forward or backward: producer retires once per iteration of 𝑆, so the count is 𝑟 𝑆 − 𝑘 = 𝑟 𝑃 − 𝑘, for any offset 𝑘 ≥ 0 (Figure 5). • Producer nested, forward (offset 0): producer is in a loop 𝑃 nested inside 𝑆, before consumer in 𝑆’s body; when consumer issues, its engine has finished this iteration’s trip through 𝑃, so the count is 𝑟 𝑃 (Figure 6 left). • Producer nested, backward (offset 1): consumer precedes producer’s loop, so its engine has not entered 𝑃 this iteration; the count through 𝑆’s previous iteration is again 𝑟 𝑃 (Figure 6 right). The nested cases admit only these two offsets: offset 0 when forward and offset 1 when backward. A backward dependency at offset 0 is unsatisfiable, as the consumer would wait on a producer iteration that has not yet run. Any larger offset requires additional registers to recover how many times 𝑃 executed in earlier trips through 𝑆; the monotone counter 𝑟 𝑃 only records total entries, not the number of entries for a specific past iteration of 𝑆 (e.g., B’s total iteration count varies in different trips through A in Figure 5). Once the no-wait pre-check fails, the wait value is in every case a formula in the single monotone register 𝑟 𝑃 : we write expr for it, equal to 𝑟 𝑃 − 𝑘 for a producer in 𝑆 and 𝑟 𝑃 for the two nested cases. 5

Sung, Shyamsunder, Zhang, Kroening, and Choi Forward (offset 0)

Backward (offset 1)

loop A: loop B: Producer Consumer

loop A : Consumer loop B : Producer

Figure 6. Producer nested below 𝑆. Both share 𝑆 = 𝐴 and producer parent 𝑃 = 𝐵, with consumer directly in 𝐴’s body. Producer and consumer may run on different engines; only their relative nesting and order in 𝐴’s body matter here, so we use a single stream. Forward: 𝐵 precedes consumer, so 𝑟 𝐵 already counts 𝐵 through the current 𝐴-iteration. Backward: consumer precedes 𝐵, so 𝑟 𝐵 counts 𝐵 through 𝐴’s previous iteration. Either way the wait value is 𝑟 𝐵 .

A Single Branchless Wait Value

waitVal ≜ mask · expr.

Engine 0

Engine 1

[𝑟ˆ𝐴 := 0] loop A: [𝑟 𝐴 ++, 𝑟ˆ𝐴 ++] 𝐼 0 {𝑠 0 ++} 𝐼 1 {𝑠 0 ++}

[𝑟ˆ𝐴 := 0] loop A: [𝑟 𝐴 ++, 𝑟ˆ𝐴 ++] [𝑤 := isGT(𝑟ˆ𝐴 , 0) · (2 · 𝑟 𝐴 ) ] {𝑠 0 ≥ 𝑤 } 𝐼 2 {𝑠 1 ++}

Per-Loop Allocation

The per-loop allocation, our main algorithm (§7), assigns one semaphore per (loop, engine) pair, shared by all instructions in that loop body on that engine. A shared counter no longer equals a single producer’s retirement count, but we recover it in closed form, since all instructions sharing a parent loop execute the same number of times on every trace. Within a loop body holding 𝑁 instructions on a given engine, each increments the loop’s shared semaphore 𝑠𝑙𝑜𝑜𝑝 (in place of the per-producer 𝑠 producer above) by 1 at retirement. After the 𝐾-th of them (by position in the body) retires in the 𝑟 -th iteration, 𝑠𝑙𝑜𝑜𝑝 = 𝑁 · 𝑟 − (𝑁 − 𝐾), so a producer producer at position 𝐾 has retired at least 𝑟 times exactly when

(2)

Per-Instruction Allocation

sloop ≥ 𝑁 · 𝑟 − (𝑁 − 𝐾).

It remains to supply 𝑅(producer) through a semaphore. The simplest scheme is per-instruction: assign one semaphore 𝑠 producer to each producer, incremented when producer retires, so 𝑠 producer = 𝑅(producer) exactly. The dependency check is 𝑠 producer ≥ isGT(𝑟ˆ𝑆 , 𝑘) · expr,

[𝑟ˆ𝐴 := 0] loop A: [𝑟 𝐴 ++, 𝑟ˆ𝐴 ++] [𝑤 := isGT(𝑟ˆ𝐴 , 0) · 𝑟 𝐴 ] {𝑠𝐼1 ≥ 𝑤 } 𝐼 2

5.5

The mask is the complement of the no-wait condition 𝑟ˆ𝑆 ≤ 𝑘: it is 0 exactly when no producer is owed, forcing waitVal = 0 so any non-negative semaphore passes; otherwise it is 1 and the threshold is expr. The single check 𝑅(producer) ≥ waitVal is thus equivalent to the full disjunction in (1), emitted as a fixed three-operation sequence (compute mask, compute expr, multiply) before each consumer issue. 5.4

[𝑟ˆ𝐴 := 0] loop A: [𝑟 𝐴 ++, 𝑟ˆ𝐴 ++] 𝐼 0 {𝑠𝐼0 ++} 𝐼 1 {𝑠𝐼1 ++}

Figure 8. Per-loop allocation for the program of Figure 7. Engine 0’s two instructions now share one semaphore 𝑠 0 , so the wait value scales by 𝑁 =2 and offsets the producer’s position 𝐾.

The register-level check is a disjunction: the no-wait precheck 𝑟ˆ𝑆 ≤ 𝑘, or the cumulative check against the wait value expr of §5.2. A semaphore wait, however, takes a single threshold computed into a wait register before issue, not a branch. We collapse the disjunction with a multiplicative mask: a boolean materialized as a 0/1 integer and multiplied into the threshold. We use an isGT operation with isGT(𝑎, 𝑏) = 1 iff 𝑎 > 𝑏. We compute mask ≜ isGT(𝑟ˆ𝑆 , 𝑘),

Engine 1

Figure 7. Per-instruction allocation. Each instruction gets its own semaphore, incremented at retirement; the bracketed regOps maintain the monotone register 𝑟 𝐴 and trip register 𝑟ˆ𝐴 . Consumer 𝐼 2 waits on the producer’s semaphore 𝑠𝐼1 with the masked wait value.

Dependencies with a multi-level gap or any other nested offset fall back to the barrier-based scheme (§3.4); our evaluation (§8) shows the subset suffices for our target AI kernels. 5.3

Engine 0

(4)

This holds because in-order retirement (§7.1), combined with sequential issue within a body, keeps all retirement counts in the loop within 1 of each other—the first few at 𝑚+1, the rest at 𝑚. Substituting the masked wait value in (2) for 𝑟 in (4) gives the per-loop check:

(3)

with the consumer computing waitVal using three controlarithmetic operations before waiting on 𝑠 producer . Figure 7 shows an inter-engine dependency 𝐼 2 depending on 𝐼 1 , with registers and wait value computations in place to realize the per-instruction allocation. A production workload can comprise millions of instructions, so one semaphore per producer does not scale. This motivates our new allocation algorithm.

𝑠𝑙𝑜𝑜𝑝 ≥ isGT(𝑟ˆ𝑆 , 𝑘) · 𝑁 · expr − (𝑁 − 𝐾) ,

(5)

where 𝑠 is the semaphore of producer’s loop on producer’s engine, and 𝑁 , 𝐾 are static. This adds two operations to the per-instruction sequence—the 𝑁 · scaling and the −(𝑁 − 𝐾) offset—for five in total per issue. Figure 8 shows the allocation on the same program from Figure 7. 6

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

6

Table 1. Resource cost across the three allocations, before the ALU reductions of §6.2. 𝑁 is the number of datapath instructions, 𝐿 the number of loops, 𝐸 the number of engines, and 𝐷 the number of loops on the path to the producer.

Resource Tradeoffs

Although the wait-value computation of §5 is fixed, we can trade the number of semaphores against per-issue ALU work by choosing how many producers share a semaphore. A more widely shared semaphore must be decoded to recover a single producer’s retirement count 𝑅(producer). We make this tradeoff explicit, adding a third per-engine point to the two allocations of §5, and then show how standard compiler passes cut the ALU cost. The per-engine allocation and the ALU-reduction passes are unverified extensions on top of the per-loop algorithm we verify in §7. 6.1

Semaphores

Registers/engine

ALU/issue

Per-instruction Per-loop Per-engine

𝑂 (𝑁 ) 𝑂 (𝐿 · 𝐸) 𝑂 (𝐸)

𝑂 (𝐿) 𝑂 (𝐿) 𝑂 (𝐿)

3 5 𝑂 (𝐷)

consumer’s computation across the others amortizes the cost.

The Semaphore–ALU Tradeoff

Each allocation maps a producer to one semaphore and recovers 𝑅(producer) from its value by a decoding whose cost grows as the semaphore is shared more widely. §5 gave two: per-instruction (3), with one semaphore per producer (𝑠 producer = 𝑅(producer), no decoding, 𝑂 (𝑁 ) semaphores), and per-loop (5), with one semaphore per (loop, engine) pair, decoded by 𝑁 · expr − (𝑁 − 𝐾) (𝑂 (𝐿 + ·𝐸) semaphores). We now add the coarsest point. Per-engine. Since in-order retirement holds across all of an engine’s instructions (§7), one counter per engine tracks its total retirements—folding in everything that retires before producer on that engine, not just producer’s own loop. Recovering 𝑅(producer) therefore means accounting for every retirement before producer on its engine, summing the loop-entry counts 𝐻 over all loops on the path from program entry to producer. The per-loop closed form 𝑁 ·expr −(𝑁 −𝐾) in (5) applies to each such loop 𝑙, reducing to the full body count 𝑁𝑙 (𝑟𝑙 − 1) for the other loops that executed before producer and to the precise 𝑁𝑃 ·expr − (𝑁𝑃 −𝐾) for producer’s parent loop 𝑃; summing over the path gives 𝑅(producer). The no-wait condition is handled as before, by the same mask isGT(𝑟ˆ𝑆 , 𝑘) in (2) applied to the entire summed wait value. This uses only 𝑂 (𝐸) semaphores but 𝑂 (𝐷) ALU operations per issue, where 𝐷 is the number of loops on that path. Each step down in semaphore count thus pushes more of the recovery of 𝑅(producer) onto runtime ALU work. 6.2

Granularity

6.3 Resource Summary Table 1 summarizes the three allocations, differing only in resource cost. A compiler can mix them across a program— per-loop where the dependency footprint is small, per-engine for deeply nested regions under a tight semaphore budget— trading semaphores, registers, and ALU work according to their relative cost on the target.

7

Verified Correctness

We formalize the multi-engine execution model of §4 as a specification transition system and the register-and-counter execution model of §5 as an implementation transition system. The implementation abstracts the functionality of a semaphore as a shared counter: a global counter that any engine may increment and read. The results therefore apply to any architecture providing such counters. We then prove via bisimulation [13] that, whenever an implementation program realizes a specification program via our per-loop allocation of §5.5, every wait fires if and only if its dependency is satisfied without under- or over-synchronization. Bridging the two requires showing 𝑟 ℓ counts the total entries of loop ℓ, 𝑟ˆℓ its entries in the current trip, and the register check in (5) coincides with the dependency satisfaction condition in (1). This holds for programs related by our allocation, over our restricted set of dependencies. The proof is mechanized in Lean in approximately 15K lines with zero admitted lemmas.

Reducing Per-Issue ALU Cost

Standard compiler passes reduce the per-issue ALU cost without affecting correctness. Loop-invariant code motion. Each loop’s contribution to the sum changes only when that loop is re-entered, so it is invariant to any loop nested deeper. Hoisting it out of those deeper loops leaves, in steady state, only the innermost contribution plus a read of a cached outer partial sum, at the cost of a few registers per engine to hold the cached sums. Common-subexpression elimination. Consumers sharing enclosing loops can share most of the wait-value computation—the register reads, the masking isGT, and the partial sums up to their closest common ancestor—so reusing one

7.1

Specification

We first formalize the semantic dependency-satisfaction model of §4 as a small-step transition system. Abstracting away the ISA. We aim to reason about execution traces and the possible relative orderings of instructions; concrete instruction semantics are irrelevant to the proof. We therefore abstract datapath instructions opaquely and keep only their effect on a shared datapath state, consisting of HBM, SBUF, PSUM, and the datapath registers. Engines are likewise treated abstractly, while the control 7

Sung, Shyamsunder, Zhang, Kroening, and Choi

State. The implementation state 𝜎𝑖 is defined similarly to 𝜎𝑠 as a tuple 𝜎𝑖 = (ds, cs, 𝜅, ifl, sema, regs). Note that the shared counter state and the engine-local register files are in place of the maps 𝐻 and 𝑅.

state is engine-local and consists only of the control registers used to steer structured execution. Program. We inductively define statements as: Stmt 𝑆 ::= block(𝑓 ) | loop(𝑔, 𝑢, 𝑆) | cond(𝑔, 𝑢𝑇 , 𝑢 𝐹 , 𝑆𝑇 , 𝑆 𝐹 )

Transition system. The implementation transition system augments the specification’s with regOp steps, checks of the shared counter at issue time, and counter updates at retirement. The rules are given in §A.5.

A specification program 𝑃𝑠 consists of a 𝑆𝑡𝑚𝑡 list and a dependency graph. The grammar keeps a separate cond form for faithfulness to the definition of SCFG, but as in §4.1 a conditional is just a loop of trip count 0 or 1, so we refer only to loops hereafter. Here, 𝑓 maps each engine to its datapath instruction list for that block, 𝑔 is a guard function from engines’ control state to a Boolean, the 𝑢-functions update control state upon exiting loops, and 𝑆 is a nested statement list. The program’s 𝑆𝑡𝑚𝑡 structure defines the SCFG common across engines, though the instruction streams and dynamic control decisions may differ by engine. The proof therefore assumes no agreement on control decisions or trip counts across engines; §7.5 explains why the bisimulation holds in such cases.

7.3

State. A specification state is a tuple 𝜎𝑠 = (ds, cs, 𝜅, ifl, 𝐻, 𝑅) consisting of the shared datapath state ds, a map cs from engines to their individual control states, a per-engine continuation stack 𝜅 (a path from the root down to the currently executing loop and the progress made within each loop), a per-engine in-flight list ifl of instructions in the pipeline, the loop-entry counts 𝐻 of §4.2, and the retirement counts 𝑅. The precise state typing is given in §A.1.

7.4

Proof of Correctness

Our correctness criterion is bisimulation [13] between the specification and implementation: a simulation relation (∼) on states that holds at corresponding initial states and is preserved by every step on either side, where the responding side may take zero or more steps of its own to match.

Transition system. The step relation fires on a single non-deterministically chosen engine per step, so the specification admits every engine interleaving consistent with the dependency-satisfaction condition in (1). Datapath instructions progress through an issue/commit/retire lifecycle: they issue only when their dependency is satisfied and retire in order, incrementing 𝑅 at retirement. We model overlap among in-flight instructions by allowing a non-deterministically chosen instruction in the pipeline to commit its update to the shared datapath state. Separate control-flow steps evaluate guards, perform loop entry and exit, update the continuation stack, and maintain 𝐻 for use in the dependency predicate. The detailed step rules are given in §A.2. 7.2

Allocation Relation

To prove the correctness of the algorithm independent of a particular compiler implementation, inspired by CompCert [31], we define an allocation relation R alloc : a declarative predicate on (𝑃𝑠 , 𝑃𝑖 ) program pairs stating that 𝑃𝑖 realizes 𝑃𝑠 via our per-loop allocation. The relation says that the two programs agree on the opaque instruction semantics and engine set, and that the implementation has the same statement-tree structure as the specification save for inserted regOp statements in block streams. In particular, the inserted regOp statements properly maintain 𝑟 and 𝑟ˆ for each loop, shared counters are injectively assigned per (loop, engine) pair, and wait/update annotations are consistent with the per-loop decoding of §5.5.

Simulation relation. A simulation relation is a binary relation between specification and implementation states. In our proof, the relation is defined by 𝜎𝑠 ∼ 𝜎𝑖 ≜ MatchStates(𝜎𝑠 , 𝜎𝑖 ) ∧ SpecInv(𝜎𝑠 ) ∧ ImplInv(𝜎𝑖 ). MatchStates(𝜎𝑠 , 𝜎𝑖 ) asserts the cross-system invariants: datapath-state equality, per-engine in-flight queue equality, control-state equality, and continuation-stack correspondence modulo inserted regOp statements. It also relates the registers to 𝐻 : for every loop ℓ and engine 𝑒, 𝑟 ℓ holds the total entries of ℓ and 𝑟ˆℓ its entries in the current trip. SpecInv and ImplInv are the side-specific invariants the proof maintains on each system. We prove bisimulation under the following hypotheses, which we unpack before stating the final theorem. • R alloc (𝑃𝑠 , 𝑃𝑖 ) of §7.3, stating that 𝑃𝑖 realizes 𝑃𝑠 via our per-loop allocation. • Allocatable(𝑃𝑠 )—every dependency in 𝑃𝑠 falls into one of the allocatable cases of §5.2. • UniqueLoopIds(𝑃𝑠 ) and UniqueInstrIds(𝑃𝑠 ), so that identifiers assigned to instructions and loops are unambiguous keys into the 𝐻 and 𝑅 maps.

Implementation

We next formalize a concrete counter-and-register execution model. This model is expressive enough to realize the masked wait-value computation and counter updates of §§5.2, 5.3 needed to instantiate a per-loop allocation. Program. An implementation program 𝑃𝑖 has the same statement grammar as a specification program, but block streams may additionally contain inserted regOp statements and there is no dependency graph. Instead, registers compute wait values while the program also records explicit wait and signal counter annotations for each datapath instruction. 8

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

Theorem 7.1 (Bisimulation). Let 𝑃𝑠 and 𝑃𝑖 satisfy the four hypotheses above. Then from any pair of common initial states 𝜎𝑠0 and 𝜎𝑖 0 (§A.6), the specification 𝑃𝑠 and implementation 𝑃𝑖 are bisimilar under ∼. That is,

hold (Allocatable, UniqueLoopIds, UniqueInstrIds). Finally, the Lean kernel is trusted, as is standard.

1. 𝜎𝑠0 ∼ 𝜎𝑖 0 ; 2. for every 𝜎𝑠 ∼ 𝜎𝑖 and every specification state 𝜎𝑠′ reached by one step from 𝜎𝑠 , the implementation can take steps from 𝜎𝑖 to some state 𝜎𝑖′ such that 𝜎𝑠′ ∼ 𝜎𝑖′ ; 3. for every 𝜎𝑠 ∼ 𝜎𝑖 and every implementation state 𝜎𝑖′ reached by one step from 𝜎𝑖 , the specification can take steps from 𝜎𝑠 to some state 𝜎𝑠′ such that 𝜎𝑠′ ∼ 𝜎𝑖′ .

We evaluate the barrier-free semaphore allocation algorithm on Trainium hardware around four questions: RQ1 (Latency vs. barriers): how much latency does barrier-free synchronization save over barrier-based? RQ2 (vs. unrolling): how does it compare to loop unrolling in latency and code size? RQ3 (vs. manual): how competitive is the barrier-free allocation with hand-tuned expert placement? RQ4 (Resource usage): what is the allocation’s register and semaphore overhead? Does it match the bounds of §6?

7.5

8

Bisimulation Under Control-Flow Divergence

Under the common SCFG assumption, loop trip counts may still differ across engines (§7.1), so a target 𝑡® can name a producer iteration that never runs—leaving the dependency permanently blocked. Concretely, consider a dependency with offset 0 from a producer on one engine to a consumer on another, both in a loop that runs zero times on the producer’s engine but once on the consumer’s. This does not break the bisimulation: on the specification side the producer’s retirement count never reaches 1, and on the implementation side the producer’s shared counter never reaches the expected wait value of 1. Thus, neither the specification nor the implementation can take the issue step for the consumer. Under the hypotheses of Theorem 7.1, the two systems block together in such cases, so their states remain related. Liveness is therefore relative to the specification induced by the input dependency graph: the implementation introduces no additional blocking, but faithfully reproduces any deadlock already present in the specification. 7.6

8.1

Experimental Setup

Hardware. We run all experiments on an AWS Trainium 2 chip through an Amazon EC2 trn2.48xlarge instance. Each kernel runs on a single Neuron core. Benchmarks are written in the AWS Neuron ISA for Trainium translated from production NKI kernels [8]. Variants. We compare four variants of each kernel that enforce the same data dependencies but differ in how they synchronize them; the dependencies are manually extracted from the bare, unsynchronized ISA program: • Barrier-Based — the baseline: at every loop iteration boundary, all engines drain their in-flight instructions, perform an all-engine handshake, reset all semaphores, and handshake again before proceeding. To isolate the effect of the synchronization strategy, we implement this scheme directly in our AWS Neuron ISA backend rather than compare against the production compiler. • Barrier-Free — our verified per-loop allocation (§5.5). • Manual — semaphore allocation placed manually by a developer familiar with the AWS Neuron ISA. • Unrolled — the loop fully unrolled with statically computed semaphore waits; unavailable for Dropless MoE which has a dynamically bounded loop.

Consequences

Because ∼ is preserved on both sides and includes equality of the shared datapath state at every related pair, Theorem 7.1 yields trace and reachable-state equivalence: the specification and implementation admit exactly the same set of datapath traces. It also yields safety (the implementation never issues before a dependency is satisfied) and liveness (when the specification can progress, so can the implementation), so it never over-synchronizes relative to the input dependency graph. 7.7

Evaluation

Benchmarks. We evaluate one microbenchmark and a suite of ML kernels, chosen to span the patterns the algorithm targets: static and dynamically bounded loops, single/nested loops and a range of dependency structures. All dependency offsets fall within the allocatable cases of §5.2. • AddManyLoop — our microbenchmark: a 104 -iteration loop of element-wise tensor addition with no inter-iteration dependencies, isolating the pure barrier overhead. • TiledMatmul — a tiled 20483 matrix multiply over M, N, and K, with the DMA, Tensor, and Vector engines cooperating per tile. Its innermost K loop is a tight, computebound body of a weight load and a matmul. We evaluate it with all three loops as hardware loops, and in a K-unrolled variant whose innermost loop is statically unrolled.

Trusted Computing Base

Theorem 7.1 holds modulo trusted assumptions. First, the specification and implementation (§§7.1, 7.2) faithfully model multi-engine execution on real hardware; in particular, the datapath, control, and allocation registers are disjoint—so computation never clobbers the counters and wait values the regOps maintain—and the compiler is trusted to allocate the three separately. Second, R alloc faithfully reflects the allocation of §5.5. Third, the syntactic hypotheses on 𝑃𝑠 9

Sung, Shyamsunder, Zhang, Kroening, and Choi

• RMSNorm — RMS normalization with FP8 quantization over a 32768 × 8192 input; software-pipelined, doublebuffering loads on one engine against stores on another. • FlashAttention — tiled attention (81922 × 128) with two nested loops over query groups and key/value sections; heavy DMA-prefetch/Tensor-compute overlap. • AdamW — the Adam optimizer with weight decay over 1M (220 ) parameters; a loop over parameter chunks with several engines cooperating per chunk. • Dropless MoE — a Mixture-of-Experts MLP iterating over the active expert blocks, a count set by the runtime token-to-expert routing not the input shape. We evaluate three configurations of increasing iteration count, varying the tokens 𝑇 and experts 𝐸 with top-𝑘 routing fixed at 𝑘=2: (𝑇 =384, 𝐸=4), (𝑇 =2048, 𝐸=8), and (𝑇 =4096, 𝐸=16). 8.2

lands directly on the critical path and leaves it far behind unrolling (0.90 vs. 0.39). Statically unrolling just that innermost loop moves the counter updates out to once per tile, where they hide behind the matmul, and barrier-free returns to par with unrolling (0.55 for both). Barrier-free thus composes well with selective unrolling and delivers a larger gain where its wait-value ALU cost can be hidden behind other engines’ work. Note Dropless MoE cannot be unrolled at all—its loop runs over the active expert blocks, set by the runtime token-to-expert routing rather than the input shape. RQ3: versus manual expert placement. On most kernels the barrier-free allocation is comparable to manual allocation, and sometimes beats it by precisely encoding the dependencies. An expert may add conservative oversynchronization for safety. The gap appears when the waitvalue ALU cannot hide behind the engines’ other work: sharing one semaphore per loop, the barrier-free algorithm decodes the semaphore to recover the producer’s retirement count, and also maintains monotone and trip counters. An expert can omit the decode step by giving the producer its own semaphore (i.e., selectively applying the perinstruction allocation) or comparing against a register the kernel already maintains (e.g., the native loop counter). On TiledMatmul the decode and counter updates fall on the critical path (RQ2); the expert replaces them with direct register comparisons, using even fewer semaphores than barrierfree (4 vs. 7). On Dropless MoE the expert instead takes the first route, spending more semaphores—21 vs. 13—to read each count directly. It also dedicates a register to each wait threshold and updates them once per iteration at the loop tail, sharing them across all waits in the body, whereas the barrier-free allocation reuses a temporary and so recomputes a decode before every wait. Making these choices well—which producers earn a dedicated semaphore or register, and how many the budget can afford—takes repeated profiling, critical-path inspection, and re-validation of the hand-placed semaphores on every change. This is feasible for small kernels, but scales poorly to larger code while risking subtle correctness bugs.

Results

Table 2 reports four metrics per variant. Lat is end-to-end latency normalized to the barrier-based baseline (lower is better), and Code is the code size. Sema is the number of semaphores the variant uses. Regs is the number of registers the allocation adds per engine for synchronization, beyond those the original kernel uses for its own computation: the loop counters 𝑟 ℓ , 𝑟ˆℓ , the masks, and the temporaries that compute wait values. Since the allocation inserts the same register families on every engine (§5.1), this number of registers is identical across engines. First, note the robustness of the per-loop allocation with the allocatable subset of §5.2, successfully synchronizing every benchmark spanning dense nested loops (TiledMatmul, FlashAttention), software-pipelined loops (RMSNorm, AdamW), and dynamically routed workloads (Dropless MoE). RQ1: latency versus barriers. The barrier-free allocation improves over the barrier-based baseline on every benchmark. The barrier-free allocation reduces latency by 70% on AddManyLoop, whose independent iterations make every barrier pure overhead, and by 10–45% on the real kernels, since the algorithm stalls only on each true dependency rather than on all engines. RQ2: versus unrolling. Unrolling a loop removes its periteration synchronization and branching, but produces code size proportional to the iteration count—up to three orders of magnitude larger here (RMSNorm 6.1 KB vs. 508.6 KB, FlashAttention 37.1 KB vs. 1.2 MB). On most kernels barrierfree allocation comes within a few percent of the unrolled latency while keeping code compact (RMSNorm 0.57 vs. 0.56, AdamW 0.74 vs. 0.76). The two TiledMatmul variants show where this breaks down. TiledMatmul is compute-bound, and its innermost K loop is a tight body of two datapath instructions (a weight load and a matmul). With all three loops kept, barrier-free must, every K iteration and on the same Tensor engine that runs the matmul, take the loop branch and update its monotone and trip counters—overhead that

RQ4: resource usage. The Regs and Sema columns empirically track the bounds of Table 1. Registers stay within 𝑂 (𝐿): 4 per engine for single-loop kernels, rising to 8 for the two-loop kernels and 11 for the three-loop kernel. The ALU-reduction passes (§6.2) and the no-wait trip-register elision (§5.1) keep these counts low. Semaphore counts scale with the loops and engines a dependency spans, ranging from 4–13, up to 19 on FlashAttention. The barrier-based baseline adds no registers, trading all its synchronization cost into per-iteration barrier latency; the Manual variant’s budget is comparable (0–7 per engine), trading semaphores against the runtime decode as RQ3 details. 10

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

Table 2. Four synchronization variants on Trainium 2 per benchmark. 𝐿 is the number of hardware loops; Offsets are the dependency offsets present in each benchmark; Norm. Lat is end-to-end latency normalized to the barrier-based baseline, so lower is better and 0.50 means half the baseline latency; Code is code size; Regs is registers added per engine for synchronization; Sema is the number of semaphores. Unrolled is unavailable (N/A) for Dropless MoE kernels due to dynamically bounded loops. Barrier-Based

Barrier-Free

Manual

Benchmark

𝐿

Offsets

Norm. Lat

Code

Regs

Sema

Norm. Lat

Code

Regs

Sema

AddManyLoop AdamW 1M flash_attn 81922 rmsnorm 32768x8192 tiled_matmul 20483 tiled_matmul 20483 (K-unrolled) MoE 384 (E=4, k=2) MoE 2048 (E=8, k=2) MoE 4096 (E=16, k=2)

1 1 2 1 3 2 1 1 1

0 0, 2 0, 1, 2 0, 2 0, 1, 2 0, 1, 2 0, 1 0, 1 0, 1

1.00 1.00 1.00 1.00 1.00 1.00 1.00 1.00 1.00

2.5 KB 5.3 KB 22.5 KB 4.0 KB 4.7 KB 6.9 KB 10.3 KB 10.3 KB 10.3 KB

0 0 0 0 0 0 0 0 0

3 4 8 4 5 5 21 21 21

0.30 0.74 0.86 0.57 0.90 0.55 0.87 0.88 0.88

1.7 KB 7.3 KB 37.1 KB 6.1 KB 6.3 KB 9.3 KB 9.3 KB 9.3 KB 9.3 KB

4 4 8 4 11 8 4 4 4

4 10 19 6 7 6 13 13 13

Summary. The barrier-free allocation beats the barrierbased baseline on every benchmark and, though fully automated, matches hand-tuned manual and unrolled performance on most kernels. Unlike unrolling, it handles dynamically bounded loops and keeps code compact. Its wait-value ALU mostly hides behind other engines’ compute; the exception is a tight, compute-bound innermost loop (TiledMatmul), resolved by unrolling it while keeping the outer loops.

9

Norm. Lat

Unrolled

Code Regs Sema

0.22 1.8 KB 0.76 5.3 KB 0.88 23.5 KB 0.56 3.9 KB 0.71 3.7 KB 0.61 6.1 KB 0.81 9.8 KB 0.79 6.1 KB 0.79 6.1 KB

0 5 0 5 7 7 4 4 4

4 4 8 4 4 4 21 21 21

Norm. Lat

Code Regs

0.04 1.2 MB 0.76 11.8 KB 0.70 1.2 MB 0.56 508.6 KB 0.39 177.5 KB 0.55 177.5 KB N/A N/A N/A N/A N/A N/A

0 0 0 0 0 0 N/A N/A N/A

Sema 4 4 8 4 4 4 N/A N/A N/A

across hardware generations and configurations, but the algorithm applies to any such multi-engine system. Verification scope. We verify the per-loop allocation under the execution model of §7, modulo the trusted computing base (§7.7). The optimizations of §6 are unverified.

10

Related Work

Synchronization primitives in parallel architectures. Synchronization primitives in parallel systems range from coarse-grained barriers [24] to fine-grained wait-free algorithms [25]. GPU programming relies on thread-block barriers and warp-level primitives in the SIMT model [34]. Among ASIC accelerators, Groq’s TSP is fully deterministic and software-scheduled by design, with the compiler planning all instruction timing statically [1, 2, 21]; Google’s TPUs [20, 27, 28] exhibit heterogeneous compute units, but their published sources do not describe on-chip inter-unit synchronization. By contrast, Trainium exposes a multiengine model with explicit semaphore-based synchronization and dynamic control flow, so the compiler must decide when each engine may proceed. DeNovo [17, 41] provides synchronization for disciplined parallelism at the hardwareprotocol level, whereas our work operates at compiler-time allocation.

Threats to Validity

Coverage. The study spans a microbenchmark and a suite of ML kernels—enough to show the dependencies we target arise in practice, but not an exhaustive sample of workloads. Kernels with different engine usage or memory behavior may shift the magnitude of the speedups. The manual baseline is one reasonable expert allocation, not a latency lower bound. Isolation. We evaluate synchronization at the AWS Neuron ISA level, holding the instruction schedule fixed. Results report the standalone effect of the synchronization strategy rather than end-to-end compiler speedups. Our contribution is orthogonal to instruction-scheduling techniques (e.g., software pipelining), which speed up the baseline and our allocation alike. Our work targets a different axis: ensuring synchronization does not block available cross-engine parallelism.

Synchronization reduction. Reducing synchronization overhead is a long-studied problem. In the HPC literature, Tseng [44] introduced the closest ancestor to our approach— compiler analysis that eliminates SPMD barriers in favor of lighter-weight synchronization—and subsequent work extended barrier elimination via access-dependency analysis in OpenMP [47], dynamic runtime elision in production codes [15], and point-to-point producer–consumer synchronization that signals specific consumers rather than all threads [14]. Xiao et al. [45] proposed lock-free inter-block GPU synchronization. These techniques target homogeneous SPMD or threaded programs on shared memory, whereas we eliminate barriers across heterogeneous engines that share

Dependency graphs. The dependencies driving the barrierfree and manual variants are extracted by hand. Correctness is validated against a NumPy reference model. The residual risk is the opposite of corruption: a conservative dependency graph leads both variants to over-synchronize. Generality. The gains require that the target is a multiengine accelerator that synchronizes through semaphorelike, shared counters, and require workloads with recoverable cross-iteration overlap; on inherently serial loops, or hardware where ALU work is costly relative to global synchronization, the advantage shrinks. Results may also shift 11

Sung, Shyamsunder, Zhang, Kroening, and Choi

no common barrier primitive, replacing them with point-topoint waits derived per engine.

per-loop allocation correct in Lean via bisimulation, and implemented it as a compiler backend pass at the Neuron ISA level. On real kernels, it reduces latency by 10–45% relative to the barrier-based baseline (and achieves a 3.3× speedup on a synchronization-bound microbenchmark), remaining competitive with expert hand-tuning without the code-size blowup of unrolling. The result is a fully automated synchronization scheme that is precise, mechanically verified, and free of both barriers and over-synchronization.

Compiler frameworks and scheduling for AI accelerators. To achieve high compute utilization for tensor workloads, many compiler frameworks separate the algorithm from its schedule [16, 36], or the data placement from the computation placement for dataflow accelerators [26]. Others provide multi-level IRs [29] or polyhedral loop optimization [11], including polyhedral kernel generation for NPUs [48]. A closer line of work schedules data movement and execution overlap directly: scheduling languages make GPU memory-hierarchy mapping explicit [22, 23], and recent systems overlap computation across distributed kernels [49] or jointly schedule warp specialization for tensor-core kernels [39]. These techniques (e.g., software pipelining) structure execution to expose parallelism. Our work is orthogonal: when parallelism is present, we generate fine-grained synchronization to achieve it.

References [1] Dennis Abts, Garrin Kimmell, Andrew C. Ling, John Kim, Matthew Boyd, Andrew Bitar, Sahil Parmar, Ibrahim Ahmed, Roberto DiCecco, David Han, John Thompson, Michael Bye, Jennifer Hwang, Jeremy Fowers, Peter Lillian, Ashwin Murthy, Elyas Mehtabuddin, Chetan Tekur, Thomas Sohmers, Kris Kang, Stephen Maresh, and Jonathan Ross. 2022. A Software-defined Tensor Streaming Multiprocessor for Large-scale Machine Learning. In Proceedings of the 49th Annual International Symposium on Computer Architecture. Association for Computing Machinery, New York, NY, USA, 567–580. doi:10.1145/ 3470496.3527405 [2] Dennis Abts, Jonathan Ross, Jonathan Sparling, Mark Wong-VanHaren, Max Baker, Tom Hawkins, Andrew Bell, John Thompson, Temesghen Kahsai, Garrin Kimmell, Jennifer Hwang, Rebekah Leslie-Hurd, Michael Bye, E. R. Creswick, Matthew Boyd, Mahitha Venigalla, Evan Laforge, Jon Purdy, Purushotham Kamath, Dinesh Maheshwari, Michael Beidler, Geert Rosseel, Omar Ahmad, Gleb Gagarin, Richard Czekalski, Ashay Rane, Sahil Parmar, Jeff Werner, Jim Sproch, Adrian Macias, and Brian Kurtz. 2020. Think Fast: A Tensor Streaming Processor (TSP) for Accelerating Deep Learning Workloads. In 2020 ACM/IEEE 47th Annual International Symposium on Computer Architecture (ISCA). IEEE, 145–158. doi:10.1109/ISCA45697.2020.00023 [3] Randy Allen and Ken Kennedy. 2001. Optimizing Compilers for Modern Architectures: A Dependence-Based Approach. Morgan Kaufmann Publishers Inc. [4] Amazon Web Services. 2022. AWS Trainium. https://aws.amazon. com/machine-learning/trainium/. [5] Amazon Web Services. 2024. AWS Neuron Compiler. https://awsdocsneuron.readthedocs-hosted.com/en/latest/compiler/. [6] Amazon Web Services. 2024. AWS Trainium2. https://aws.amazon. com/machine-learning/trainium/trainium2/. [7] Amazon Web Services. 2024. NeuronCore-v2 Architecture Guide. https://awsdocs-neuron.readthedocs-hosted.com/en/latest/ nki/guides/architecture/trainium_inferentia2_arch.html. [8] Amazon Web Services. 2024. NKI Library Documentation. https: //awsdocs-neuron.readthedocs-hosted.com/en/latest/nki/library. [9] Amazon Web Services. 2025. AWS Trainium3. https://aws.amazon. com/machine-learning/trainium/trainium3/. [10] Jai Arora, Sirui Lu, Devansh Jain, Tianfan Xu, Farzin Houshmand, Phitchaya Mangpo Phothilimthana, Mohsen Lesani, Praveen Narayanan, Karthik Srinivasa Murthy, Rastislav Bodík, Amit Sabne, and Charith Mendis. 2025. TensorRight: Automated Verification of Tensor Graph Rewrites. Proceedings of the ACM on Programming Languages 9, POPL, Article 29 (2025). doi:10.1145/3704865 [11] Uday Bondhugula, Albert Hartono, Jagannathan Ramanujam, and Ponnuswamy Sadayappan. 2008. A practical automatic polyhedral parallelizer and locality optimizer. In Proceedings of the 29th ACM SIGPLAN Conference on Programming Language Design and Implementation. 101–113. [12] Timothy Bourke, Lélio Brun, Pierre-Évariste Dagand, Xavier Leroy, Marc Pouzet, and Lionel Rieg. 2017. A Formally Verified Compiler for Lustre. In Proceedings of the 38th ACM SIGPLAN Conference on

Verified compiler passes. Our bisimulation proof follows the simulation-based semantic-preservation methodology developed in CompCert [31], which is used broadly in verified compilers, including CakeML [42]. Schneider et al. [37] prove simulation-based compiler correctness by induction on program structure. Individual passes have also been verified, including instruction and superblock scheduling [38, 46] and register allocation [30, 43]—a similar kind of resourceallocation problem to our semaphore allocation. The verified Lustre compiler [12], also built on CompCert, instead targets a synchronous-reactive dataflow language whose logical lock-step clocks differ fundamentally from the asynchronous, semaphore-coordinated engines we compile for. Like us, Burrow reasons about the space of allowed concurrent executions, but its framework targets weak-memory consistency mappings (e.g., x86 to Arm), whereas we prove preservation of synchronization across asynchronous engines [40]. Formal methods for AI accelerator compilation. Formal methods are increasingly applied across the layers of the AI-compilation stack: graph-rewrite verification in XLA over unbounded tensor ranks [10], verified functional-tensor compilation [32], verification dialects embedded in MLIR [19], translation validation for Halide schedules [18], and equivalence checking of distributed training plans [33]. These efforts verify tensor-level computation, such as graph rewrites, schedules, and training plans, whereas we verify the hardwarespecific code generation beneath them: the synchronization and parallel execution of an accelerator’s engines.

11

Conclusion

We identified the exact synchronization condition for dependencies on a multi-engine AI accelerator, and compiled it into barrier-free, closed-form semaphore waits for arbitrarily nested loops with dynamic bounds. We proved the 12

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

Programming Language Design and Implementation. 586–601. [13] Stephen D. Brookes and William C. Rounds. 1983. Behavioural equivalence relations induced by programming logics. In Automata, Languages and Programming, Josep Diaz (Ed.). Springer Berlin Heidelberg, Berlin, Heidelberg, 97–108. [14] J Mark Bull and Carwyn Ball. 2003. Point-to-point synchronisation on shared memory architectures. In Proceedings of the Fifth European Workshop on OpenMP-EWOMP, Vol. 3. [15] Milind Chabbi, Wim Lavrijsen, Wibe de Jong, Koushik Sen, John Mellor-Crummey, and Costin Iancu. 2015. Barrier elision for production parallel programs. ACM SIGPLAN Notices 50, 8 (2015), 109–119. [16] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Meghan Cowan, Haichen Shen, Leyuan Wang, Yuwei Hu, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. TVM: An Automated End-to-End Optimizing Compiler for Deep Learning. In OSDI. 578–594. [17] Byn Choi, Rakesh Komuravelli, Hyojin Sung, Robert Smolinski, Nima Honarmand, Sarita V. Adve, Vikram S. Adve, Nicholas P. Carter, and Ching-Tsun Chou. 2011. DeNovo: Rethinking the Memory Hierarchy for Disciplined Parallelism. In Proceedings of the 2011 International Conference on Parallel Architectures and Compilation Techniques (PACT ’11). IEEE Computer Society, USA, 155–166. doi:10.1109/PACT.2011.21 [18] Basile Clément and Albert Cohen. 2022. End-to-End Translation Validation for the Halide Language. Proc. ACM Program. Lang. 6, OOPSLA1, Article 84 (April 2022), 30 pages. doi:10.1145/3527328 [19] Mathieu Fehr, Yuyou Fan, Hugo Pompougnac, John Regehr, and Tobias Grosser. 2025. First-Class Verification Dialects for MLIR. Proc. ACM Program. Lang. 9, PLDI, Article 206 (June 2025), 25 pages. doi:10.1145/ 3729309 [20] Google Cloud. 2025. TPU System Architecture. https://docs.cloud. google.com/tpu/docs/system-architecture-tpu-vm. [21] Groq. 2024. Determinism and the Tensor Streaming Processor. https: //groq.com/GroqDocs/TechDoc_Predictability.pdf Technical documentation, accessed 2026-06-05. [22] Bastian Hagedorn, Archibald Samuel Elliott, Henrik Barthels, Rastislav Bodik, and Vinod Grover. 2020. Fireiron: A Data-Movement-Aware Scheduling Language for GPUs. In Proceedings of the ACM International Conference on Parallel Architectures and Compilation Techniques (Virtual Event, GA, USA) (PACT ’20). Association for Computing Machinery, New York, NY, USA, 71–82. doi:10.1145/3410463.3414632 [23] Bastian Hagedorn, Bin Fan, Hanfeng Chen, Cris Cecka, Michael Garland, and Vinod Grover. 2023. Graphene: An IR for Optimized Tensor Computations on GPUs. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3 (Vancouver, BC, Canada) (ASPLOS 2023). Association for Computing Machinery, New York, NY, USA, 302–313. doi:10.1145/3582016.3582018 [24] Debra Hensgen, Raphael Finkel, and Udi Manber. 1988. Two algorithms for barrier synchronization. International Journal of Parallel Programming 17, 1 (1988), 1–17. [25] Maurice Herlihy. 1991. Wait-free synchronization. ACM Transactions on Programming Languages and Systems (TOPLAS) 13, 1 (1991), 124– 149. [26] Olivia Hsu, Alexander Rucker, Tian Zhao, Varun Desai, Kunle Olukotun, and Fredrik Kjolstad. 2025. Stardust: Compiling Sparse Tensor Algebra to a Reconfigurable Dataflow Architecture. In Proceedings of the 23rd ACM/IEEE International Symposium on Code Generation and Optimization (Las Vegas, NV, USA) (CGO ’25). Association for Computing Machinery, New York, NY, USA, 628–643. doi:10.1145/3696443.3708918 [27] Norman P. Jouppi, George Kurian, Sheng Li, Peter Ma, Rahul Nagarajan, Lifeng Nai, Nishant Patil, Suvinay Subramanian, Andy Swing, Brian Towles, Cliff Young, Xiang Zhou, Zongwei Zhou, and David Patterson. 2023. TPU v4: An Optically Reconfigurable Supercomputer

for Machine Learning with Hardware Support for Embeddings. In Proceedings of the 50th Annual International Symposium on Computer Architecture (ISCA ’23). doi:10.1145/3579371.3589350 [28] Norman P. Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, Rick Boyle, Pierre-luc Cantin, Clifford Chao, Chris Clark, Jeremy Coriell, Mike Daley, Matt Dau, Jeffrey Dean, Ben Gelb, Tara Vazir Ghaemmaghami, Rajendra Gottipati, William Gulland, Robert Hagmann, C. Richard Ho, Doug Hogberg, John Hu, Robert Hundt, Dan Hurt, Julian Ibarz, et al. 2017. In-Datacenter Performance Analysis of a Tensor Processing Unit. In Proceedings of the 44th Annual International Symposium on Computer Architecture (ISCA ’17). 1–12. doi:10.1145/3079856.3080246 [29] Chris Lattner, Mehdi Amini, Uday Bondhugula, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, and Oleksandr Zinenko. 2021. MLIR: Scaling Compiler Infrastructure for Domain Specific Computation. In CGO. 2–14. [30] Kent D Lee. 2003. A formally verified register allocation framework. Electronic Notes in Theoretical Computer Science 82, 3 (2003), 515–531. [31] Xavier Leroy. 2009. A Formally Verified Compiler Back-end. J. Autom. Reason. 43, 4 (Dec. 2009), 363–446. doi:10.1007/s10817-009-9155-4 [32] Amanda Liu, Gilbert Bernstein, Adam Chlipala, and Jonathan RaganKelley. 2024. A Verified Compiler for a Functional Tensor Language. Proceedings of the ACM on Programming Languages 8, PLDI, Article 160 (2024). [33] Yunchi Lu, Youshan Miao, Cheng Tan, Peng Huang, Yi Zhu, Xian Zhang, and Fan Yang. 2025. TrainVerify: Equivalence-Based Verification for Distributed LLM Training. In Proceedings of the ACM SIGOPS 31st Symposium on Operating Systems Principles (Lotte Hotel World, Seoul, Republic of Korea) (SOSP ’25). Association for Computing Machinery, New York, NY, USA, 237–253. doi:10.1145/3731569.3764850 [34] NVIDIA. 2024. CUDA C++ Programming Guide. https://docs.nvidia. com/cuda/cuda-c-programming-guide/. [35] Julian Oppermann, Melanie Reuter-Oppermann, Lukas Sommer, Andreas Koch, and Oliver Sinnen. 2019. Exact and Practical Modulo Scheduling for High-Level Synthesis. ACM Trans. Reconfigurable Technol. Syst. 12, 2, Article 8 (May 2019), 26 pages. doi:10.1145/3317670 [36] Jonathan Ragan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. 2013. Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines. In PLDI. 519–530. [37] Sigurd Schneider, Gert Smolka, and Sebastian Hack. 2016. An Inductive Proof Method for Simulation-based Compiler Correctness. arXiv:1611.09606 [cs.PL] https://arxiv.org/abs/1611.09606 [38] Cyril Six, Léo Gourdin, Sylvain Boulmé, David Monniaux, Justus Fasse, and Nicolas Nardino. 2022. Formally Verified Superblock Scheduling. In Proceedings of the 11th ACM SIGPLAN International Conference on Certified Programs and Proofs (Philadelphia, PA, USA) (CPP 2022). Association for Computing Machinery, New York, NY, USA, 40–54. doi:10.1145/3497775.3503679 [39] Rupanshu Soi, Rohan Yadav, Fredrik Kjolstad, Alex Aiken, Maryam Mehri Dehnavi, Michael Garland, and Michael Bauer. 2025. Optimal Software Pipelining and Warp Specialization for Tensor Core GPUs. arXiv:2512.18134 [cs.PL] https://arxiv.org/abs/2512.18134 [40] Dennis Sprokholt and Soham Chakraborty. 2026. Burrow: A Proof Framework for Weak Memory. In CAV. [41] Hyojin Sung and Sarita V. Adve. 2015. DeNovoSync: Efficient Support for Arbitrary Synchronization without Writer-Initiated Invalidations. SIGPLAN Not. 50, 4 (March 2015), 545–559. doi:10.1145/2775054. 2694356 [42] Yong Kiam Tan, Magnus O Myreen, Ramana Kumar, Anthony Fox, Scott Owens, and Michael Norrish. 2019. The Verified CakeML Compiler Backend. Journal of Functional Programming 29 (2019), e2.

13

Sung, Shyamsunder, Zhang, Kroening, and Choi

[43] Jean-Baptiste Tristan and Xavier Leroy. 2008. Formal Verification of Translation Validators: A Case Study on Instruction Scheduling Optimizations. SIGPLAN Not. 43, 1 (Jan. 2008), 17–27. doi:10.1145/ 1328897.1328444 [44] Chau-Wen Tseng. 1995. Compiler optimizations for eliminating barrier synchronization. In Proceedings of the Fifth ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (Santa Barbara, California, USA) (PPOPP ’95). Association for Computing Machinery, New York, NY, USA, 144–155. doi:10.1145/209936.209952 [45] Shucai Xiao and Wu-chun Feng. 2010. Inter-block GPU Communication via Fast Barrier Synchronization. In 2010 IEEE International Symposium on Parallel & Distributed Processing (IPDPS). IEEE, 1–12. [46] Ziteng Yang, Jun Shirako, and Vivek Sarkar. 2024. Fully Verified Instruction Scheduling. Proc. ACM Program. Lang. 8, OOPSLA2, Article 299 (Oct. 2024), 26 pages. doi:10.1145/3689739 [47] Naoki Yonezawa, Koichi Wada, and Takahiro Aida. 2006. Barrier Elimination based on Access Dependency Analysis for OpenMP. In Proceedings of the 4th International Conference on Parallel and Distributed Processing and Applications (Sorrento, Italy) (ISPA’06). Springer-Verlag, Berlin, Heidelberg, 362–373. doi:10.1007/11946441_36 [48] Jie Zhao, Bojie Li, Wang Nie, Zhen Geng, Renwei Zhang, Xiong Gao, Bin Cheng, Chen Wu, Yun Cheng, Zheng Li, Peng Di, Kun Zhang, and Xuefeng Jin. 2021. AKG: Automatic Kernel Generation for Neural Processing Units Using Polyhedral Transformations. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation (Virtual, Canada) (PLDI 2021). Association for Computing Machinery, New York, NY, USA, 1233–1248. doi:10.1145/3453483.3454106 [49] Size Zheng, Wenlei Bao, Qi Hou, Xuegui Zheng, Jin Fang, Chenhui Huang, Tianqi Li, Haojie Duanmu, Renze Chen, Ruifan Xu, Yifan Guo, Ningxin Zheng, Ziheng Jiang, Xinyi Di, Dongyang Wang, Jianxi Ye, Haibin Lin, Li-Wen Chang, Liqiang Lu, Yun Liang, Jidong Zhai, and Xin Liu. 2025. Triton-distributed: Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler. arXiv:2504.19442 [cs.DC] https://arxiv.org/abs/2504.19442

14

A Barrier-Free Synchronization Algorithm for Multi-Engine AI Accelerators

A

Verification

(1) holds. The step appends the instruction to the inflight list as issued and advances the block position. 2. Commit. A non-deterministically chosen issued inflight instruction applies its opaque datapath semantics to the shared datapath state and is marked committed. This is how the model represents overlap among in-flight instructions. 3. Retire. The head of the in-flight list, if already committed, is removed and the corresponding producer retirement count in 𝑅 is incremented. Retirement is strictly in issue order. 4. Loop entry. A loop guard is evaluated. On loop entry, a new continuation-stack frame is pushed and the loop-entry histories/counts used by the dependency predicate are updated. 5. Loop exit / back edge. When execution reaches the end of a loop body, the corresponding frame is popped; loops then return to their header for guard re-evaluation. 6. Block completion. When all instructions of a block have issued, the continuation stack advances to the next statement.

The mechanization uses a map-based representation: every loop and every datapath instruction is assigned a unique identifier, written LoopId and DataPathInstrId, which serve as the keys for the internal maps below. This is the content of the UniqueLoopIds and UniqueInstrIds hypotheses in §7.4. A.1

Specification State

The specification state is 𝜎𝑠 = (ds, cs, 𝜅, ifl, 𝐻, 𝑅). Here ds is the shared datapath state. The remaining components are maps: cs : EngineId → ControlState, 𝜅 : EngineId → ContStack, ifl : EngineId → InFlightList, 𝐻 : EngineId → LoopId → N → LoopId → N, 𝑅 : DataPathInstrId → N. In the proof code, dependencies are represented only as triples (producer, consumer, 𝑘) where the offset 𝑘 is on the shared loop 𝑆, so we index 𝐻 as the number of entries of loop 𝑃 during the 𝑖𝑆 -th entry of 𝑆, where 𝑖𝑆 is 𝑆’s monotone entry count. This is equivalent to the iteration-vector presentation in the paper body for the shared-loop case, because once 𝑆 is fixed, each monotone entry of 𝑆 bijects with exactly one common-loop iteration vector at 𝑆. The specification program components used by these rules are likewise treated abstractly but have fixed shape. For a block statement block(𝑓 ), the block stream is

A.3

An implementation program again consists of a structured statement tree, but each block stream may interleave datapath instructions with inserted RegOp instructions: 𝑓𝑖 : EngineId → (DataPathInstrId + RegOp) list. It also carries two explicit semaphore-annotation maps on datapath instructions, waitOf : DataPathInstrId → SemaId,

𝑓 : EngineId → DataPathInstrId list,

updateOf : DataPathInstrId → SemaId,

while the shared functions appearing in loop/conditional constructors and the opaque instruction semantics are typed by

identifying, respectively, which semaphore is checked before the Issue step and which semaphore is incremented at Retire.

instrOp : DataPathInstrId → DataPathState

A.4

→ DataPathState,

Implementation State

The implementation state modifies the specification state with shared semaphore state and engine-local register files. Concretely, the tuple is

guard : EngineId → LoopId → ControlState → Bool, controlOp : EngineId → LoopId → ControlState → ControlState.

𝜎𝑖 = (ds, cs, 𝜅, ifl, sema, regs),

The dependency graph is likewise keyed by instruction identifiers, and this is the type-level information suppressed in the main text. A.2

Implementation Program Components

where the new map components are sema : SemaId → N, regs : EngineId → RegId → N.

Specification Step Rules

Under the semaphore allocation, an implementation maintains one distinguished register on each engine serving as the wait register, holding the computed waitVal read by issuetime semaphore checks. The information 𝑅 and 𝐻 carry semantically is instead recovered operationally, through RegOp counter maintenance and Issue-/Retire-time semaphore effects.

The specification small-step relation chooses one engine nondeterministically at each step. The detailed rule families are below. We write only loops here, though the proof code has separate entry and exit rules for loops and conditionals. 1. Issue. The next instruction in the current block may issue only if its dependency-satisfaction predicate in 15

Sung, Shyamsunder, Zhang, Kroening, and Choi

A.5

Implementation Step Rules

A.6

The implementation follows the same structured controlflow skeleton as the specification, but with three additional operational features. First, each inserted RegOp executes atomically as a control-side step: it reads the current enginelocal register file, performs the prescribed ALU computation, and writes back the updated register values in one transition. The mechanization supports the operations relevant to our proofs: addition, subtraction, and multiplication, each in immediate and register forms, together with isGT; all are modeled with natural-number arithmetic. Second, Issue no longer consults the semantic dependency predicate directly; instead, for a datapath instruction 𝐼 , the implementation uses waitOf (𝐼 ) to select the semaphore associated with 𝐼 and checks that this semaphore is at least the value currently stored in the distinguished wait register. Third, Retire uses updateOf (𝐼 ) to determine which semaphore to increment when a producer instruction retires. These are the only new effects beyond the specification’s behavior.

Initial States

The bisimulation theorem starts the specification and implementation from common initial states: both systems are built from the same shared datapath state value ds0 and the same per-engine control states cs0 , with all engines positioned at the first statement of the top-level loop. Concretely, let 𝜅 0 be the initial continuation-stack map pointing each engine to the top-level statement list at its first statement, with initial block offset 0. Let [ ] ifl map each engine to the empty in-flight list, and let 0𝑅 , 0𝐻 , 0sema , and 0regs denote the corresponding constant-zero functions. The specification initial state is 𝜎𝑠0 = (ds0, cs0, 𝜅 0, [ ] ifl , 0𝐻 , 0𝑅 ), and the implementation initial state is 𝜎𝑖 0 = (ds0, cs0, 𝜅 0, [ ] ifl , 0sema, 0regs ). Here, the empty in-flight lists and the zero maps indicate that no instruction has issued yet, no loop-entry or retirement counts have been accumulated, and all registers and semaphores start at zero.

16

Related documents

• Bridging the Deployment Gap: Integrating AI into Accelerator Control Systems#403519
• MATCHA: Efficient Deployment of Deep Neural Networks on Multi-Accelerator Heterogeneous Edge SoCs#5935
• Introduction to hardware accelerator systems for artificial intelligence and machine learning#10760
• AccelForge: Comprehensive Modeling and Co-Design Framework for AI Accelerators#999584
• A Multi-Engine Dataflow for MoE Decoding on Scratchpad-Based Tensor Accelerators#1050941