Triton 3.9.0: Atomic operations, earlier compile-time error detection

1. Control memory reads and writes with atomic operations

This version adds atomic operations that read or write memory values in a single indivisible action. The new tl.atomic_load and tl.atomic_store are also available in Gluon and the interpreter.

tl.atomic_poll waits until the value at a specified address matches the expected value. It returns false if no match occurs within the time limit, providing a termination condition when waiting for a value to become ready.

Extending memory ordering, the order in which reads and writes are observed, across an entire block comes with added synchronization costs. Atomic operations with specified read/write ordering include barriers, synchronization points that pause and align execution within the same block. Developers can specify sem="relaxed" to avoid these barriers, weighing required ordering guarantees against synchronization overhead.

2. Detect invalid inputs before execution

A mismatch in the number of items unpacked from a tuple now fails at compile time. Invalid selection counts in tl.topk, which selects top values, and incompatible shapes in tl.cat, which joins arrays, also trigger errors at this stage.

Non-32-bit inputs to tl.histogram, which computes frequency distributions, also cause errors before execution. Moving input validation to compile time lets developers inspect inputs and call sites instead of looking for anomalies in results after running a kernel.

Compilation time has also improved. According to the release notes, compilation of deeply nested loops with a nesting depth of 8 is about 10x faster. This result applies to that specified nesting depth.

3. Align interpreter and compiled kernel results

The debugging interpreter executes code by interpreting it, and these fixes focus on aligning its results with compiled kernels. They correct integer overflow, where values exceed the representable range, and operations that compute the high bits of integer products.

Fused operations that combine multiplication and addition, along with histogram calculations, were also corrected for consistency. Developers checking these operations in the interpreter will encounter fewer result mismatches caused by execution differences. Alongside compile-time checks that catch errors early, these changes improve result verification.

4. Check type conversions in existing kernels

If operands, the input values to division or remainder operations, include floating-point values narrower than 32 bits, the result is float32. This also applies when used with Python scalars. To retain the original result type, cast the operation's result to the desired type.

Tensor math methods now undergo the same type checks as functions. For example, the exponential method fails to compile with 16-bit floating-point input. The release notes provide this example of converting to 32-bit floating point first:

x.to(tl.float32).exp()

These changes also create points to revise in existing code that handled types automatically. When upgrading, developers need to fix new compilation errors and check the types used to store operation results.

5. Prepare for the next version's fusion default

Implicit floating-point fusion, which automatically combines multiplication and addition, remains enabled in 3.9. However, compiling a kernel without specifying enable_fp_fusion issues a warning once per process. In 3.10, this fusion is scheduled to be disabled by default.

The release notes advise explicitly calling tl.fma in kernels whose numerical results depend on fusion. Expressing computational intent in code rather than relying on the current default prepares kernels for the next version's default change.

Triton 3.9.0 lets kernel authors directly manage atomic memory behavior, catch errors earlier, and review changed ordering and type rules ahead of the next version.