tilelang-cuda-examples-torch

PyTorch + TileLang CUDA 示例代码

本 Skill 包含完整的可运行示例代码，展示如何在 PyTorch 中使用 TileLang CUDA 编写高性能 kernel。

示例列表

1. 矩阵乘法（GEMM）

算子类型: MatMul 关键点:

• 共享内存缓存输入块
• T.gemm 利用 Tensor Core
• 软件流水线 T.Pipelined
• 混合精度（float32 累加器）

import torch
import tilelang
import tilelang.language as T

@tilelang.jit(out_idx=[-1])
def matmul(M, N, K, block_M, block_N, block_K):
    @T.prim_func
    def main(A: T.Tensor((M, K), "float16"),
             B: T.Tensor((K, N), "float16"),
             C: T.Tensor((M, N), "float16")):
        
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            A_shared = T.alloc_shared((block_M, block_K), "float16")
            B_shared = T.alloc_shared((block_K, block_N), "float16")
            C_local = T.alloc_fragment((block_M, block_N), "float")
            
            T.clear(C_local)
            
            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                T.copy(A[by * block_M, ko * block_K], A_shared)
                T.copy(B[ko * block_K, bx * block_N], B_shared)
                T.gemm(A_shared, B_shared, C_local)
            
            T.copy(C_local, C[by * block_M, bx * block_N])
    
    return main

# 调用方式
def matmul_call(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    M, K = A.shape
    K2, N = B.shape
    block_M, block_N, block_K = 128, 128, 32
    kernel = matmul(M, N, K, block_M, block_N, block_K)
    C = kernel(A, B)  # out_idx=[-1]，只传输入
    return C

2. 矩阵乘法（float32，手动管理输出）

算子类型: MatMul 关键点:

• 不使用 out_idx，手动管理输出
• float32 数据类型
• 需要手动创建输出张量并一起传入

import torch
import tilelang
import tilelang.language as T

@tilelang.jit
def square_matrix_multiply(M, N, K, block_M, block_N, block_K):
    @T.prim_func
    def main(
            A: T.Tensor((M, K), "float32"),
            B: T.Tensor((K, N), "float32"),
            C: T.Tensor((M, N), "float32")):
        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
            A_shared = T.alloc_shared((block_M, block_K), "float32")
            B_shared = T.alloc_shared((block_K, block_N), "float32")
            C_local = T.alloc_fragment((block_M, block_N), "float")
            
            T.clear(C_local)

            for ko in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
                T.copy(A[by * block_M, ko * block_K], A_shared)
                T.copy(B[ko * block_K, bx * block_N], B_shared)
                T.gemm(A_shared, B_shared, C_local)

            T.copy(C_local, C[by * block_M, bx * block_N])

    return main

def square_matrix_multiply_call(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    N = A.size(0)
    block_M, block_N, block_K = 128, 128, 32
    # 不使用 out_idx 时，需要手动创建输出张量
    C = torch.empty_like(A)
    kernel = square_matrix_multiply(N, N, N, block_M, block_N, block_K)
    kernel(A, B, C)  # 传入所有参数包括输出
    return C