Submission 03: Matrix Multiplication with cuTile
The file assignments/03_assignment/src/__main__.py contains the main function that runs all the tasks for this assignment. Each task is implemented in a separate file in the same directory. The results of each task are printed to the console when the main function is executed.
Task 1: FP32 vs FP16 Performance
Output:
runtime torch.float16: 0.03595617373028527
runtime torch.float32: 1.84091664514234
speedup: 51.198902835197046
import cuda.tile as ct
import cupy as cp
from numpy.strings import index
import torch
import triton
def main():
times = []
for dt in [torch.float16, torch.float32]:
run(dt)
t = benchmark(dt)
print(f"runtime {dt}: {t}")
times.append(t)
print(f"speedup: {times[1]/times[0]}")
def setup(dt):
inner_size = 4096
A = torch.randn((64, inner_size), device='cuda', dtype=dt)
B = torch.randn((inner_size, 64), device='cuda', dtype=dt)
C = torch.empty((64, 64), device='cuda', dtype=torch.float32)
grid = (1, )
torch.cuda.init()
return A, B, C, grid
def run(dt):
A, B, C, grid = setup(dt)
ct.launch(torch.cuda.current_stream(), grid, kernel, (A, B, C))
torch.cuda.synchronize()
expected = torch.empty((64, 64), device='cuda', dtype=dt)
torch.matmul(A, B, out=expected)
expected = expected.to(torch.float32) # Convert to float32 for comparison
assert torch.allclose(C, expected, atol=1e-1), "The result is incorrect!"
def benchmark(dt):
A, B, C, grid = setup(dt)
def run_kernel():
ct.launch(torch.cuda.current_stream(), grid, kernel, (A, B, C))
torch.cuda.synchronize()
t = triton.testing.do_bench(run_kernel)
return t
@ct.kernel
def kernel(A, B, C):
m_tile=64
n_tile=64
k_tile=64
result = ct.load(C, index=(0, 0), shape=(m_tile, n_tile))
for i in range(0, A.shape[0] // m_tile):
for j in range(0, B.shape[1] // n_tile):
for k in range(0, A.shape[1] // k_tile):
A_block = ct.load(A, index=(i, k), shape=(m_tile, k_tile))
B_block = ct.load(B, index=(k, j), shape=(k_tile, n_tile))
result = ct.mma(A_block, B_block, acc=result)
ct.store(C, index=(0, 0), tile=result)
if __name__ == "__main__":
main()
Task 2: Simple Matrix Multiplication Kernel
import math
import cuda.tile as ct
import cupy as cp
from numpy.strings import index
import torch
import triton
def main():
M = 321
N = 123
K = 23
m_tile=64
n_tile=32
k_tile=128
M_padded = int(2**math.ceil(math.log2(max(M, m_tile))))
N_padded = int(2**math.ceil(math.log2(max(N, n_tile))))
K_padded = int(2**math.ceil(math.log2(max(K, k_tile))))
A = torch.randn((M, K), device='cuda', dtype=torch.float16)
B = torch.randn((K, N), device='cuda', dtype=torch.float16)
C = torch.empty((M, N), device='cuda', dtype=torch.float32)
grid = (math.ceil(M_padded / m_tile) * math.ceil(N_padded / n_tile), )
torch.cuda.init()
ct.launch(torch.cuda.current_stream(), grid, kernel_fp16, (A, B, C, m_tile, n_tile, k_tile, M_padded, N_padded, K_padded))
torch.cuda.synchronize()
expected = torch.empty((M, N), device='cuda', dtype=torch.float16)
torch.matmul(A, B, out=expected)
expected = expected.to(torch.float32) # Convert to float32 for comparison
assert torch.allclose(C, expected, atol=1e-1), "The result is incorrect!"
@ct.kernel
def kernel_fp16(A, B, C, m_tile: ct.Constant[int], n_tile: ct.Constant[int], k_tile: ct.Constant[int], M_padded: ct.Constant[int], N_padded: ct.Constant[int], K_padded: ct.Constant[int]):
bid = ct.bid(0)
bid_x = bid % (N_padded // n_tile)
bid_y = bid // (N_padded // n_tile)
result = ct.zeros((m_tile, n_tile), dtype=torch.float32)
for k in range(0, K_padded // k_tile):
A_block = ct.load(A, index=(bid_y , k), shape=(m_tile, k_tile), padding_mode=ct.PaddingMode.ZERO)
B_block = ct.load(B, index=(k, bid_x), shape=(k_tile, n_tile), padding_mode=ct.PaddingMode.ZERO)
result = ct.mma(A_block, B_block, acc=result)
# print("Result in kernel:\n", result)
ct.store(C, index=(bid_y, bid_x), tile=result)
if __name__ == "__main__":
main()
Task 3: Benchmarking the Matrix Multiplication Kernel
a) Benchmark your kernel with tile shapes (64, 64, 64) for square matrix multiplications of sizes:
b) Fix the matrix size at 2048 × 2048 × 2048, as well as 512 × 512 × 512, and benchmark all tile shape combinations (27 total):
Output:
-> BEST tile shape for 512x512x512 is (128, 64, 128) achieving 11.52 TFLOPS
-> BEST tile shape for 2048x2048x2048 is (128, 128, 64) achieving 54.06 TFLOPS
Task 4: L2 Cache Optimization via Block Swizzling
@ct.kernel
def kernel_matmul_swizzle(A, B, C, tm: ct.Constant[int], tn: ct.Constant[int], tk: ct.Constant[int], grid_x, grid_y):
swizzle_size = 8
pid = ct.bid(0)
index_m, index_n = calc_position(pid, swizzle_size, grid_y, grid_x)
num_tiles_k = ct.num_tiles(A, axis=1, shape=(tm, tk))
accumulator = ct.full((tm, tn), 0, dtype=ct.float32)
for k in range(num_tiles_k):
a = ct.load(A, index=(index_m, k), shape=(tm, tk), padding_mode=ct.PaddingMode.ZERO)
b = ct.load(B, index=(k, index_n), shape=(tk, tn), padding_mode=ct.PaddingMode.ZERO)
accumulator = ct.mma(a, b, accumulator)
ct.store(C, index=(index_m, index_n), tile=accumulator)
def calc_position(pid, swizzle_size, grid_y, grid_x):
num_pid_in_stripe = swizzle_size * grid_x
stripe_index = pid // num_pid_in_stripe
begin_n = stripe_index * swizzle_size
stripe_height = swizzle_size
if (begin_n + swizzle_size) > grid_y:
stripe_height = grid_y - begin_n
index_n_temp = pid % stripe_height
index_m_temp = pid // stripe_height
index_m = index_m_temp % grid_x
index_n = begin_n + index_n_temp
return index_m, index_n
PIDs are mapped into horizontal ‘stripes’ across the output matrix. Each stripe consists of 8 rows. Within a stripe, the PIDs traverse the tiles column by column: the first 8 PIDs compute a vertical column of 8 tiles downwards. When the stripe is finished. The next stripe is computed, starting at row index 8. At the last stripe the remaining heiht of the stripe (the rows) are calculated dynamically, to prevent out-of-bounds memory accesses.
Benchmarks
Output:
-> BEST tile shape for 512x512x512 is (64, 128, 64) achieving 11.17 TFLOPS
-> BEST tile shape for 2048x2048x2048 is (128, 128, 64) achieving 54.86 TFLOPS
Task 4b comparison at M=8192, N=8192, K=4096, tile=(128, 128, 64):
-> Task 2 kernel (row-major BIDs): 26.35 TFLOPS
-> Task 4 kernel (swizzled BIDs): 69.99 TFLOPS
-> Speedup from swizzling: 2.66x