[BACKEND][PPU] Apply formatting to INT8 AIU changes

Author: July-h5kf3

third_party/ppu/lib/TritonPPUGPUToLLVM/AIUUtility.cpp

@@ -257,17 +257,15 @@ DenseMap<unsigned, Value> getPPUAIUV1SwizzledSharedPtrs(
     Value sliceStartPtr =
         b.gep(dstPtrTy, resElemTy, dstPtrBase, sliceStartOffset);
 
-    Value idxColInnerSlice =
-        b.urem(idxColInnerCube, b.i32_val(sliceElems));
+    Value idxColInnerSlice = b.urem(idxColInnerCube, b.i32_val(sliceElems));
     // new swizzled row index inside slice, swizzled slice shape is (cubeW/4,
     // 64)
     Value rowSwizzleID = b.udiv(idxRowInnerCube, b.i32_val(4));
     // new linear slice index inside slice, slice shape is (cubeW/4, 64)
     Value idxColSlicelinear = b.urem(
         b.add(b.mul(idxRowInnerCube, b.i32_val(sliceElems)), idxColInnerSlice),
         b.i32_val(128 / elemBytes));
-    Value colSliceID =
-        b.udiv(idxColSlicelinear, b.i32_val(16 / elemBytes));
+    Value colSliceID = b.udiv(idxColSlicelinear, b.i32_val(16 / elemBytes));
 
     // rotated length: (((sliceID>1)|(sliceID<1))&0x3) << 1
     // sliceID 0, 1, 2, 3 ---> rotated length: 0, 4, 2, 6
@@ -284,13 +282,12 @@ DenseMap<unsigned, Value> getPPUAIUV1SwizzledSharedPtrs(
     Value colRotID = b.sub(b.i32_val(7), colRotBitPos);
     Value colSwizzleID = b.xor_(colRotID, b.urem(rowSwizzleID, b.i32_val(2)));
 
-    Value swizzleOffset = b.add(
-        b.mul(rowSwizzleID, b.i32_val(128 / elemBytes)),
-        b.mul(colSwizzleID, b.i32_val(16 / elemBytes)));
+    Value swizzleOffset = b.add(b.mul(rowSwizzleID, b.i32_val(128 / elemBytes)),
+                                b.mul(colSwizzleID, b.i32_val(16 / elemBytes)));
 
     // for minVec is not equal to outVec
-    swizzleOffset = b.or_(
-        swizzleOffset, b.urem(idxCol, b.i32_val(16 / elemBytes)));
+    swizzleOffset =
+        b.or_(swizzleOffset, b.urem(idxCol, b.i32_val(16 / elemBytes)));
     ret[elemIdx] = b.gep(dstPtrTy, resElemTy, sliceStartPtr, swizzleOffset);
   }
 

third_party/ppu/lib/TritonPPUGPUToLLVM/ConvertLayoutOpToLLVM/SharedToDotOperandPPUAIUV1.cpp

@@ -54,20 +54,19 @@ loadX4B8(ConversionPatternRewriter &rewriter, Location loc, Value smemBase,
     llvm::report_fatal_error(
         "PPU0010 B8 AIU dot requires native A-row/B-col layout");
   Value bytePtr = b.bitcast(smemBase, ptr_ty(rewriter.getContext(), 3));
-  Value topLeftIndex =
-      b.add(b.mul(startCoordY, blockLineStride), startCoordX);
+  Value topLeftIndex = b.add(b.mul(startCoordY, blockLineStride), startCoordX);
   Value sliceId = b.udiv(channelOffset, b.i32_val(32));
   Value sliceOffset = b.mul(sliceId, b.mul(blockLineStride, b.i32_val(32)));
-  bytePtr = b.gep(ptr_ty(rewriter.getContext(), 3), i8_ty,
-                  bytePtr, sliceOffset);
+  bytePtr =
+      b.gep(ptr_ty(rewriter.getContext(), 3), i8_ty, bytePtr, sliceOffset);
   Value sBase = b.or_(
       b.shl(b.and_(sliceId, b.i32_val(3)), b.i32_val(27)),
       b.or_(b.shl(b.and_(blockLineStride, b.i32_val(0x7ff)), b.i32_val(16)),
             b.and_(topLeftIndex, b.i32_val(0xffff))));
   auto resultTy = vec_ty(i32_ty, 4);
-  Value loaded = callIntrinsic(rewriter, loc,
-                               "llvm.ppu.tsm.ld.swizzle.b32x4.p3i8",
-                               resultTy, {bytePtr, b.i32_val(1), sBase});
+  Value loaded =
+      callIntrinsic(rewriter, loc, "llvm.ppu.tsm.ld.swizzle.b32x4.p3i8",
+                    resultTy, {bytePtr, b.i32_val(1), sBase});
   Value r0 = b.extract_element(i32_ty, loaded, b.i32_val(0));
   Value r1 = b.extract_element(i32_ty, loaded, b.i32_val(1));
   Value r2 = b.extract_element(i32_ty, loaded, b.i32_val(2));

third_party/ppu/python/test/unit/tle/test_tle_aiu_async_load.py

@@ -171,10 +171,7 @@ def test_aiu_promotion_dtypes(dtype):
     assert "aiu_load" in compiled.asm["ttir"]
     _assert_no_tle_residue(compiled)
     if dtype == "fp32":
-        aiu_instructions = [
-            line for line in compiled.asm["llir"].splitlines()
-            if "ppu.cp.async.aiu" in line
-        ]
+        aiu_instructions = [line for line in compiled.asm["llir"].splitlines() if "ppu.cp.async.aiu" in line]
         assert aiu_instructions
         assert all(".b16" in line for line in aiu_instructions)
         assert all(".b8" not in line for line in aiu_instructions)
@@ -196,9 +193,7 @@ def test_int8_aiu_load_device_correctness():
     expected = torch.arange(32 * 32, device="cuda", dtype=torch.int32)
     expected = expected.remainder(127).to(torch.int8).reshape(32, 32)
     actual = torch.empty_like(expected)
-    _typed_aiu_load[(1,)](
-        expected, actual, 32, 32, 32, 32, num_warps=4, num_stages=1
-    )
+    _typed_aiu_load[(1, )](expected, actual, 32, 32, 32, 32, num_warps=4, num_stages=1)
     torch.cuda.synchronize()
     torch.testing.assert_close(actual.cpu(), expected.cpu(), rtol=0, atol=0)
 
@@ -302,12 +297,10 @@ def _int8_gemm_aiu_kernel(
 ):
     pid_m = tl.program_id(0)
     pid_n = tl.program_id(1)
-    a_bp = tl.make_block_ptr(a_ptr, shape=(M, K), strides=(K, 1),
-                             offsets=(pid_m * BLOCK_M, 0),
+    a_bp = tl.make_block_ptr(a_ptr, shape=(M, K), strides=(K, 1), offsets=(pid_m * BLOCK_M, 0),
                              block_shape=(BLOCK_M, BLOCK_K), order=(1, 0))
     # b_ptr is physically contiguous [N, K], logically column-major [K, N].
-    b_bp = tl.make_block_ptr(b_ptr, shape=(K, N), strides=(1, K),
-                             offsets=(0, pid_n * BLOCK_N),
+    b_bp = tl.make_block_ptr(b_ptr, shape=(K, N), strides=(1, K), offsets=(0, pid_n * BLOCK_N),
                              block_shape=(BLOCK_K, BLOCK_N), order=(0, 1))
     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.int32)
     for _ in range(0, K, BLOCK_K):
@@ -318,8 +311,7 @@ def _int8_gemm_aiu_kernel(
         b_bp = tl.advance(b_bp, (BLOCK_K, 0))
     offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
     offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
-    tl.store(c_ptr + N * offs_m[:, None] + offs_n[None, :], acc,
-             mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))
+    tl.store(c_ptr + N * offs_m[:, None] + offs_n[None, :], acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))
 
 
 @_skip_no_sdk
@@ -345,9 +337,7 @@ def test_int8_gemm_uses_v1_b8_aiu_load_and_int8_mma():
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="PPU device not available")
 @pytest.mark.parametrize(
     "m,n,k,bm,bn,bk,num_warps",
-    [(16, 16, 32, 16, 16, 32, 1),
-     (16, 32, 64, 16, 32, 64, 1),
-     (64, 64, 64, 64, 64, 64, 4),
+    [(16, 16, 32, 16, 16, 32, 1), (16, 32, 64, 16, 32, 64, 1), (64, 64, 64, 64, 64, 64, 4),
      (128, 128, 128, 64, 64, 64, 4)],
 )
 def test_int8_gemm_aiu_device_correctness(m, n, k, bm, bn, bk, num_warps):
@@ -356,8 +346,7 @@ def test_int8_gemm_aiu_device_correctness(m, n, k, bm, bn, bk, num_warps):
     b = torch.randint(-8, 9, (n, k), device="cuda", dtype=torch.int8)
     actual = torch.empty((m, n), device="cuda", dtype=torch.int32)
     grid = (triton.cdiv(m, bm), triton.cdiv(n, bn))
-    _int8_gemm_aiu_kernel[grid](a, b, actual, m, n, k, bm, bn, bk,
-                                num_warps=num_warps, num_stages=1)
+    _int8_gemm_aiu_kernel[grid](a, b, actual, m, n, k, bm, bn, bk, num_warps=num_warps, num_stages=1)
     torch.cuda.synchronize()
     expected = a.cpu().to(torch.int32) @ b.cpu().to(torch.int32).T
     torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
@@ -445,8 +434,7 @@ def kernel(a_ptr, c_ptr, M: tl.constexpr, K: tl.constexpr, BLOCK_M: tl.constexpr
                  mask=(offs_m[:, None] < M) & (offs_k[None, :] < K))
 
     compiled = _compile(kernel, {"a_ptr": "*fp16", "c_ptr": "*fp16"},
-                        {"M": 512, "K": 512, "BLOCK_M": 64, "BLOCK_K": 64},
-                        options={"num_stages": num_stages})
+                        {"M": 512, "K": 512, "BLOCK_M": 64, "BLOCK_K": 64}, options={"num_stages": num_stages})
     _assert_stages_exist(compiled)
     assert "aiu_load" in compiled.asm["ttir"]
     _assert_no_tle_residue(compiled)
@@ -467,12 +455,10 @@ def kernel(a_ptr, c_ptr, M: tl.constexpr, K: tl.constexpr, BLOCK_M: tl.constexpr
             a_bp = tl.advance(a_bp, (0, BLOCK_K))
         offs_m = pid * BLOCK_M + tl.arange(0, BLOCK_M)
         offs_k = tl.arange(0, BLOCK_K)
-        tl.store(c_ptr + K * offs_m[:, None] + offs_k[None, :], acc,
-                 mask=(offs_m[:, None] < M) & (offs_k[None, :] < K))
+        tl.store(c_ptr + K * offs_m[:, None] + offs_k[None, :], acc, mask=(offs_m[:, None] < M) & (offs_k[None, :] < K))
 
     compiled = _compile_through_llir(kernel, {"a_ptr": "*i8", "c_ptr": "*i32"},
-                                     {"M": 512, "K": 512, "BLOCK_M": 64, "BLOCK_K": 64},
-                                     options={"num_stages": 2})
+                                     {"M": 512, "K": 512, "BLOCK_M": 64, "BLOCK_K": 64}, options={"num_stages": 2})
     _assert_stages_exist(compiled)
     assert "aiu_load" in compiled.asm["ttir"]
     _assert_no_tle_residue(compiled)