Matrix_Multiplication_Triton.py

import torch
import triton
import triton.language as tl

@triton.jit
def matmul_triton_kernel(
    A, B, C,
    M, N, K,
    stride_A_row, stride_A_col,
    stride_B_col, stride_B_k,
    stride_C_row, stride_C_col,
    TILE_M: tl.constexpr,
    TILE_N: tl.constexpr,
    TILE_K: tl.constexpr
):
    pid_row = tl.program_id(0)
    pid_col = tl.program_id(1)

    offs_row = pid_row * TILE_M + tl.arange(0, TILE_M)
    offs_col = pid_col * TILE_K + tl.arange(0, TILE_K)

    out_ptrs = C + (offs_row[:, None] * stride_C_row + offs_col[None, :] * stride_C_col)
    acc = tl.zeros((TILE_M, TILE_K), dtype=tl.float32)

    for n in range(0, N, TILE_N):
        offs_mid = n + tl.arange(0, TILE_N)
        A_tile = A + (offs_row[:, None] * stride_A_row + offs_mid[None, :] * stride_A_col)
        B_tile = B + (offs_mid[:, None] * stride_B_col + offs_col[None, :] * stride_B_k)

        A_mask = (offs_row[:, None] < M) & (offs_mid[None, :] < N)
        B_mask = (offs_mid[:, None] < N) & (offs_col[None, :] < K)

        a_block = tl.load(A_tile, mask=A_mask, other=0.0)
        b_block = tl.load(B_tile, mask=B_mask, other=0.0)
        acc += tl.dot(a_block, b_block)

    out_mask = (offs_row[:, None] < M) & (offs_col[None, :] < K)
    tl.store(out_ptrs, acc, mask=out_mask)

def run_matmul(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int):
    sA_row, sA_col = A.stride()
    sB_col, sB_k = B.stride()
    sC_row, sC_col = C.stride()

    TILE_M, TILE_N, TILE_K = 64, 32, 64
    grid = (triton.cdiv(M, TILE_M), triton.cdiv(K, TILE_K))

    matmul_triton_kernel[grid](
        A, B, C,
        M, N, K,
        sA_row, sA_col,
        sB_col, sB_k,
        sC_row, sC_col,
        TILE_M, TILE_N, TILE_K,
        num_warps=4, num_stages=2
    )