[TLE][PPU] Add INT8 AIU support to PPU0010 V1 lowering

[July-h5kf3]

Background

TLE promotes tle.load(block_ptr, is_async=True) to PPU AIU asynchronous copies. The FP16/BF16 path already lowers through PPU0010 lowering V1, but the INT8 path did not have a complete B8 AIU data path.

This is not a PPU0010 hardware limitation. The public ACTLIZE implementation provides direct evidence:

• ACTLIZE 1.0.0 README lists INT8 GEMM, AIU load, and TSM swizzle load support for PPU 1.0, and includes ZW810E in the supported hardware list.
• copy_ppu0010_aiu.hpp implements both non-transposed and transposed 8-bit linear/swizzled AIU loads for PPU0010 and emits ppu.cp.async.aiu.bulk.tensor.shared.global.padz.swzl.2d.b8.
• default_gemm_configuration.hpp configures PPU0010 INT8 GEMM with MainloopPPUAiu<3>, DefaultGemm_AIU_Operand<int8_t, ...>, and S8 x S8 -> S32 tensor-core MMA.

Moving PPU0010 to lowering V2 is not a valid workaround because V2 emits the unsupported 2d.tile instruction form. This PR keeps the existing V1 dispatch and completes B8 support in V1.

Root cause

The PPU0010 V1 lowering contained several related FP16 assumptions:

1. V1 async-copy lowering hard-coded the .b16 suffix instead of selecting the instruction from the source element width.
2. getPPUAIUV1SwizzledSharedPtrs hard-coded FP16 geometry for the 32-byte slice, 16-byte vector, and 128-byte swizzle period.
3. Shared-to-dot lowering only used the B16 ldmatrix form and did not emit the native B8 TSM swizzle load.
4. For K >= 64, the shared pointer for the second K32 fragment omitted slice_id * cubeW * 32 INT8 elements.
5. End-to-end GEMM exposed ppu.st.global, while the accepted global-store mnemonic is st.global. Shared stores still require ppu.st.shared.

Minimal reproducer

The following kernel loads both INT8 operands through TLE async block-pointer loads and executes S8 x S8 -> S32 GEMM:

import torch
import triton
import triton.language as tl
import triton.experimental.tle.language as tle


@triton.jit
def int8_aiu_gemm(
    a_ptr,
    b_ptr,
    c_ptr,
    M: tl.constexpr,
    N: tl.constexpr,
    K: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    a_block = tl.make_block_ptr(
        base=a_ptr,
        shape=(M, K),
        strides=(K, 1),
        offsets=(pid_m * 16, 0),
        block_shape=(16, 32),
        order=(1, 0),
    )
    # b_ptr is physically [N, K] and viewed logically as [K, N].
    b_block = tl.make_block_ptr(
        base=b_ptr,
        shape=(K, N),
        strides=(1, K),
        offsets=(0, pid_n * 16),
        block_shape=(32, 16),
        order=(0, 1),
    )

    a = tle.load(a_block, is_async=True)
    b = tle.load(b_block, is_async=True)
    acc = tl.dot(a, b, out_dtype=tl.int32)

    offs_m = pid_m * 16 + tl.arange(0, 16)
    offs_n = pid_n * 16 + tl.arange(0, 16)
    tl.store(
        c_ptr + offs_m[:, None] * N + offs_n[None, :],
        acc,
        mask=(offs_m[:, None] < M) & (offs_n[None, :] < N),
    )


M, N, K = 16, 16, 32
a = torch.randint(-8, 8, (M, K), device="cuda", dtype=torch.int8)
b = torch.randint(-8, 8, (N, K), device="cuda", dtype=torch.int8)
out = torch.empty((M, N), device="cuda", dtype=torch.int32)

int8_aiu_gemm[(1, 1)](a, b, out, M, N, K, num_warps=1)
torch.testing.assert_close(
    out.cpu(), a.cpu().int() @ b.cpu().int().T, rtol=0, atol=0
)

Before this change, the INT8 path could not lower into a valid end-to-end B8 AIU GEMM. After the change it emits:

ppu.cp.async.aiu...2d.b8
llvm.ppu.tsm.ld.swizzle.b32x4.p3i8
ppu.mma.sync.aligned.m16n16k32...s8.s8.s32

There is no BF16 fallback and no fallback to ordinary INT8 pointer loads.

Solution

• Select .b8 or .b16 for PPU0010 V1 async AIU copies from the operand element width.

• Make V1 shared-memory swizzle addressing element-width-aware:

  sliceElems = 32 / elemBytes
  vectorElems = 16 / elemBytes
  swizzlePeriodElems = 128 / elemBytes
  

• Add the native llvm.ppu.tsm.ld.swizzle.b32x4.p3i8 shared-to-dot load for B8 operands.

• Use the INT8 K32 fragment width and apply the required shared-memory slice offset for later fragments.

• Emit st.global for global stores while preserving ppu.st.shared for shared-memory stores.

• Add focused compile-time and device tests for B8 AIU load, native B8 shared-to-dot load, INT8 MMA, K32/K64/K128 paths, global-store spelling, and pipelined loads.

Validation

Code generation

The repaired path emits:

ppu.cp.async.aiu.bulk.tensor.shared.global.padz.swzl.zfill.2d.b8
llvm.ppu.tsm.ld.swizzle.b32x4.p3i8
ppu.mma.sync.aligned.m16n16k32.row.col.satfinite.s32.s8.s8.s32

The tests also verify that:

• both GEMM operands use B8 AIU copies;
• no .b16 AIU copy is emitted for INT8;
• no B8 ldmatrix fallback is emitted;
• PPU0010 continues to use lowering V1;
• the generated code contains st.global and not ppu.st.global.

Tests

Test scope	Result
Full TLE AIU async-load test file	33 passed
Focused B8/INT8 AIU regression set	9 passed
Downstream MM AIU tests	24 passed
Final downstream MM/BMM regression	93 passed
git diff --check	passed

Device correctness is bit-exact against the Torch INT32 reference for:

16 x 16 x 32
16 x 32 x 64
64 x 64 x 64
128 x 128 x 128

The 128 x 128 x 128 case was also bit-exact across 10 consecutive runs.

Performance

The downstream operators consume pre-quantized INT8 + FP32 scale inputs. Quantization, graph capture, and output allocation are excluded from kernel timing. Measurements use CUDAGraph with 25 warmup replays and 100 measured replays. Speedups are geometric means across valid shapes.

Downstream operator	Shapes	Correctness	vs previous INT8/TLE	vs Torch BF16	vs FlagGems BF16
MM	99	99/99	1.570x, 99/99 wins	1.756x, 99/99 wins	7.082x, 95/96 wins
BMM	9	9/9	1.176x, 8/9 wins	1.598x, 9/9 wins	1.598x, 8/9 wins
Block-wise BMM	29	29/29	1.367x, 24/29 wins	1.285x, 27/29 wins	2.041x, 27/27 wins

The following FlagGems BF16 baselines did not produce valid results and are excluded from the corresponding geometric means:

MM:
(8, 2048, 4096)
(8, 12288, 2048)
(8, 9216, 2048)

Block-wise BMM:
(1, 8, 1024, 4096)
(1, 8, 1024, 7168)

The MM comparison against FlagGems BF16 is amplified by anomalously slow small-N BF16 cases. The comparisons against the previous INT8/TLE path and Torch BF16 are more representative.

The downstream MM/BMM kernels and benchmark artifacts are performance evidence only and are not part of this PR.

[CLAassistant]

All committers have signed the CLA.

[sunnycase]

LGTM

[sunnycase]

LGTM