1818import re
1919import shutil
2020import pytest
21+ import torch
2122import triton
2223import triton .language as tl
24+ from triton import knobs
2325from triton .backends .compiler import GPUTarget
2426
2527tle_backend = pytest .importorskip (
@@ -44,9 +46,22 @@ def _ppu_sdk_available() -> bool:
4446_skip_no_sdk = pytest .mark .skipif (not _ppu_sdk_available (), reason = "PPU SDK not available" )
4547
4648
47- def _compile (kernel , signature , constexprs , target = None ):
49+ def _compile (kernel , signature , constexprs , target = None , options = None ):
4850 src = triton .compiler .ASTSource (fn = kernel , signature = signature , constexprs = constexprs )
49- return triton .compile (src , target = target or _GPU_TARGET )
51+ return triton .compile (src , target = target or _GPU_TARGET , options = options )
52+
53+
54+ def _compile_through_llir (kernel , signature , constexprs , target = None , options = None ):
55+ previous_hook = knobs .runtime .add_stages_inspection_hook
56+
57+ def stop_before_hgbin (_backend , stages , _options , _language , _capability ):
58+ stages ["hgbin" ] = lambda _src , _metadata : b""
59+
60+ knobs .runtime .add_stages_inspection_hook = stop_before_hgbin
61+ try :
62+ return _compile (kernel , signature , constexprs , target , options )
63+ finally :
64+ knobs .runtime .add_stages_inspection_hook = previous_hook
5065
5166
5267def _assert_stages_exist (compiled , stages = ("ttir" , "ttgir" , "llir" )):
@@ -60,6 +75,17 @@ def _assert_no_tle_residue(compiled):
6075 assert not leak , f"residual tle.* ops in LLIR:\n " + "\n " .join (leak [:5 ])
6176
6277
78+ def _assert_int8_uses_ppu_aiu_v1_b8 (compiled ):
79+ ttgir = compiled .asm ["ttgir" ]
80+ llir = compiled .asm ["llir" ]
81+ aiu_instructions = [line for line in llir .splitlines () if "ppu.cp.async.aiu" in line ]
82+ assert "versionMajor = 1" in ttgir
83+ assert aiu_instructions , "expected an async AIU copy instruction"
84+ assert all (".2d.b8" in line for line in aiu_instructions )
85+ assert all (".b8" in line for line in aiu_instructions )
86+ assert all (".b16" not in line for line in aiu_instructions )
87+
88+
6389# ---------------------------------------------------------------------------
6490# 1. Basic promotion & fallback
6591# ---------------------------------------------------------------------------
@@ -144,6 +170,37 @@ def test_aiu_promotion_dtypes(dtype):
144170 _assert_stages_exist (compiled )
145171 assert "aiu_load" in compiled .asm ["ttir" ]
146172 _assert_no_tle_residue (compiled )
173+ if dtype == "fp32" :
174+ aiu_instructions = [
175+ line for line in compiled .asm ["llir" ].splitlines ()
176+ if "ppu.cp.async.aiu" in line
177+ ]
178+ assert aiu_instructions
179+ assert all (".b16" in line for line in aiu_instructions )
180+ assert all (".b8" not in line for line in aiu_instructions )
181+
182+
183+ @_skip_no_sdk
184+ def test_int8_aiu_load_uses_v1_b8_instruction ():
185+ compiled = _compile_through_llir (_typed_aiu_load , {"a_ptr" : "*i8" , "c_ptr" : "*i8" },
186+ {"M" : 256 , "K" : 256 , "BLOCK_M" : 64 , "BLOCK_K" : 64 })
187+ _assert_stages_exist (compiled )
188+ assert "aiu_load" in compiled .asm ["ttir" ]
189+ _assert_no_tle_residue (compiled )
190+ _assert_int8_uses_ppu_aiu_v1_b8 (compiled )
191+
192+
193+ @_skip_no_sdk
194+ @pytest .mark .skipif (not torch .cuda .is_available (), reason = "PPU device not available" )
195+ def test_int8_aiu_load_device_correctness ():
196+ expected = torch .arange (32 * 32 , device = "cuda" , dtype = torch .int32 )
197+ expected = expected .remainder (127 ).to (torch .int8 ).reshape (32 , 32 )
198+ actual = torch .empty_like (expected )
199+ _typed_aiu_load [(1 ,)](
200+ expected , actual , 32 , 32 , 32 , 32 , num_warps = 4 , num_stages = 1
201+ )
202+ torch .cuda .synchronize ()
203+ torch .testing .assert_close (actual .cpu (), expected .cpu (), rtol = 0 , atol = 0 )
147204
148205
149206# ---------------------------------------------------------------------------
@@ -231,6 +288,94 @@ def _gemm_aiu_kernel(
231288_GEMM_CONSTEXPRS = {"M" : 512 , "N" : 512 , "K" : 256 , "BLOCK_M" : 64 , "BLOCK_N" : 64 , "BLOCK_K" : 64 }
232289
233290
291+ @triton .jit (do_not_specialize_on_alignment = ["a_ptr" , "b_ptr" , "c_ptr" ])
292+ def _int8_gemm_aiu_kernel (
293+ a_ptr ,
294+ b_ptr ,
295+ c_ptr ,
296+ M : tl .constexpr ,
297+ N : tl .constexpr ,
298+ K : tl .constexpr ,
299+ BLOCK_M : tl .constexpr ,
300+ BLOCK_N : tl .constexpr ,
301+ BLOCK_K : tl .constexpr ,
302+ ):
303+ pid_m = tl .program_id (0 )
304+ pid_n = tl .program_id (1 )
305+ a_bp = tl .make_block_ptr (a_ptr , shape = (M , K ), strides = (K , 1 ),
306+ offsets = (pid_m * BLOCK_M , 0 ),
307+ block_shape = (BLOCK_M , BLOCK_K ), order = (1 , 0 ))
308+ # b_ptr is physically contiguous [N, K], logically column-major [K, N].
309+ b_bp = tl .make_block_ptr (b_ptr , shape = (K , N ), strides = (1 , K ),
310+ offsets = (0 , pid_n * BLOCK_N ),
311+ block_shape = (BLOCK_K , BLOCK_N ), order = (0 , 1 ))
312+ acc = tl .zeros ((BLOCK_M , BLOCK_N ), dtype = tl .int32 )
313+ for _ in range (0 , K , BLOCK_K ):
314+ a = tle .load (a_bp , is_async = True )
315+ b = tle .load (b_bp , is_async = True )
316+ acc += tl .dot (a , b )
317+ a_bp = tl .advance (a_bp , (0 , BLOCK_K ))
318+ b_bp = tl .advance (b_bp , (BLOCK_K , 0 ))
319+ offs_m = pid_m * BLOCK_M + tl .arange (0 , BLOCK_M )
320+ offs_n = pid_n * BLOCK_N + tl .arange (0 , BLOCK_N )
321+ tl .store (c_ptr + N * offs_m [:, None ] + offs_n [None , :], acc ,
322+ mask = (offs_m [:, None ] < M ) & (offs_n [None , :] < N ))
323+
324+
325+ @_skip_no_sdk
326+ def test_int8_gemm_uses_v1_b8_aiu_load_and_int8_mma ():
327+ compiled = _compile_through_llir (
328+ _int8_gemm_aiu_kernel ,
329+ {"a_ptr" : "*i8" , "b_ptr" : "*i8" , "c_ptr" : "*i32" },
330+ _GEMM_CONSTEXPRS ,
331+ options = {"num_stages" : 2 },
332+ )
333+ llir = compiled .asm ["llir" ]
334+ aiu_instructions = [line for line in llir .splitlines () if "ppu.cp.async.aiu" in line ]
335+ assert len (aiu_instructions ) >= 2
336+ assert all (".2d.b8" in line for line in aiu_instructions )
337+ assert llir .count ("llvm.ppu.tsm.ld.swizzle.b32x4.p3i8" ) >= 2
338+ assert "ldmatrix.sync.aligned.m8n8.x4.swzl.shared.b8" not in llir
339+ assert "ldmatrix.sync.aligned.m16n16.x1.swzl.trans.shared.b8" not in llir
340+ assert "llvm.ppu.mat.trans.b8" not in llir
341+ assert "ppu.mma.sync.aligned.m16n16k32.row.col.satfinite.s32.s8.s8.s32" in llir
342+
343+
344+ @_skip_no_sdk
345+ @pytest .mark .skipif (not torch .cuda .is_available (), reason = "PPU device not available" )
346+ @pytest .mark .parametrize (
347+ "m,n,k,bm,bn,bk,num_warps" ,
348+ [(16 , 16 , 32 , 16 , 16 , 32 , 1 ),
349+ (16 , 32 , 64 , 16 , 32 , 64 , 1 ),
350+ (64 , 64 , 64 , 64 , 64 , 64 , 4 ),
351+ (128 , 128 , 128 , 64 , 64 , 64 , 4 )],
352+ )
353+ def test_int8_gemm_aiu_device_correctness (m , n , k , bm , bn , bk , num_warps ):
354+ torch .manual_seed (123 )
355+ a = torch .randint (- 8 , 9 , (m , k ), device = "cuda" , dtype = torch .int8 )
356+ b = torch .randint (- 8 , 9 , (n , k ), device = "cuda" , dtype = torch .int8 )
357+ actual = torch .empty ((m , n ), device = "cuda" , dtype = torch .int32 )
358+ grid = (triton .cdiv (m , bm ), triton .cdiv (n , bn ))
359+ _int8_gemm_aiu_kernel [grid ](a , b , actual , m , n , k , bm , bn , bk ,
360+ num_warps = num_warps , num_stages = 1 )
361+ torch .cuda .synchronize ()
362+ expected = a .cpu ().to (torch .int32 ) @ b .cpu ().to (torch .int32 ).T
363+ torch .testing .assert_close (actual .cpu (), expected , rtol = 0 , atol = 0 )
364+
365+
366+ @_skip_no_sdk
367+ def test_global_store_uses_sdk_opcode_spelling ():
368+ compiled = _compile_through_llir (
369+ _int8_gemm_aiu_kernel ,
370+ {"a_ptr" : "*i8" , "b_ptr" : "*i8" , "c_ptr" : "*i32" },
371+ _GEMM_CONSTEXPRS ,
372+ options = {"num_stages" : 2 },
373+ )
374+ llir = compiled .asm ["llir" ]
375+ assert "st.global" in llir
376+ assert "ppu.st.global" not in llir
377+
378+
234379@_skip_no_sdk
235380def test_gemm_aiu_both_operands_promoted ():
236381 """Both A and B matrix loads should be promoted to AIULoadOp."""
@@ -300,10 +445,38 @@ def kernel(a_ptr, c_ptr, M: tl.constexpr, K: tl.constexpr, BLOCK_M: tl.constexpr
300445 mask = (offs_m [:, None ] < M ) & (offs_k [None , :] < K ))
301446
302447 compiled = _compile (kernel , {"a_ptr" : "*fp16" , "c_ptr" : "*fp16" },
303- {"M" : 512 , "K" : 512 , "BLOCK_M" : 64 , "BLOCK_K" : 64 })
448+ {"M" : 512 , "K" : 512 , "BLOCK_M" : 64 , "BLOCK_K" : 64 },
449+ options = {"num_stages" : num_stages })
450+ _assert_stages_exist (compiled )
451+ assert "aiu_load" in compiled .asm ["ttir" ]
452+ _assert_no_tle_residue (compiled )
453+
454+
455+ @_skip_no_sdk
456+ def test_int8_aiu_pipelined_loop_uses_v1_b8_instruction ():
457+
458+ @triton .jit
459+ def kernel (a_ptr , c_ptr , M : tl .constexpr , K : tl .constexpr , BLOCK_M : tl .constexpr , BLOCK_K : tl .constexpr ):
460+ pid = tl .program_id (0 )
461+ a_bp = tl .make_block_ptr (a_ptr , shape = (M , K ), strides = (K , 1 ), offsets = (pid * BLOCK_M , 0 ),
462+ block_shape = (BLOCK_M , BLOCK_K ), order = (1 , 0 ))
463+ acc = tl .zeros ((BLOCK_M , BLOCK_K ), dtype = tl .int32 )
464+ for _ in range (0 , K , BLOCK_K ):
465+ a = tle .load (a_bp , is_async = True )
466+ acc += a .to (tl .int32 )
467+ a_bp = tl .advance (a_bp , (0 , BLOCK_K ))
468+ offs_m = pid * BLOCK_M + tl .arange (0 , BLOCK_M )
469+ offs_k = tl .arange (0 , BLOCK_K )
470+ tl .store (c_ptr + K * offs_m [:, None ] + offs_k [None , :], acc ,
471+ mask = (offs_m [:, None ] < M ) & (offs_k [None , :] < K ))
472+
473+ compiled = _compile_through_llir (kernel , {"a_ptr" : "*i8" , "c_ptr" : "*i32" },
474+ {"M" : 512 , "K" : 512 , "BLOCK_M" : 64 , "BLOCK_K" : 64 },
475+ options = {"num_stages" : 2 })
304476 _assert_stages_exist (compiled )
305477 assert "aiu_load" in compiled .asm ["ttir" ]
306478 _assert_no_tle_residue (compiled )
479+ _assert_int8_uses_ppu_aiu_v1_b8 (compiled )
307480
308481
309482# ---------------------------------------------------------------------------
0 commit comments