GPU 模式讲座 1 的笔记

def time_pytorch_function(func, input):    # cuda is async so can't use python time module    start = torch.cuda.event(enable_timing=true)    end = torch.cuda.event(enable_timing=true)    # warmup    for _ in range(5):        func(input)    start.record()    func(input)    end.record()    torch.cuda.synchronize()    return start.elapsed_time(end)time_pytorch_function(torch.square, b)time_pytorch_function(square_2, b)time_pytorch_function(square_3, b)

profiling torch.square:self cpu time total: 10.577msself cuda time total: 3.266msprofiling a * a:self cpu time total: 5.417msself cuda time total: 3.276msprofiling a ** 2:self cpu time total: 6.183msself cuda time total: 3.274ms

from torch.utils.cpp_extension import load_inlinesquare_matrix_extension = load_inline(    name='square_matrix_extension',    cpp_sources=cpp_source,    cuda_sources=cuda_source,    functions=['square_matrix'],    with_cuda=true,    extra_cuda_cflags=["-o2"],    build_directory='./load_inline_cuda',    # extra_cuda_cflags=['--expt-relaxed-constexpr'])a = torch.tensor([[1., 2., 3.], [4., 5., 6.]], device='cuda')print(square_matrix_extension.square_matrix(a))

import torch# method 1: use `torch.mean()`def mean_all_by_torch(input_tensor):    return torch.mean(input_tensor)# method 2: use `mean()` of the tensordef mean_all_by_tensor(input_tensor):    return input_tensor.mean()# method 3: use `torch.sum()` and `tensor.numel()`def mean_all_by_combination(input_tensor):    return torch.sum(input_tensor) / input_tensor.numel()def time_pytorch_function(func, input_tensor, warmup=5, runs=100):    # warmup    for _ in range(warmup):      func(input_tensor)    times = []    start = torch.cuda.event(enable_timing=true)    end = torch.cuda.event(enable_timing=true)    for _ in range(runs):        start.record()        func(input_tensor)        end.record()        torch.cuda.synchronize()        times.append(start.elapsed_time(end))    return sum(times) / len(times)input_tensor = torch.randn(10000, 10000).cuda()print("torch.mean() time:", time_pytorch_function(mean_all_by_torch, input_tensor))print("tensor.mean() time:", time_pytorch_function(mean_all_by_tensor, input_tensor))print("manual mean time:", time_pytorch_function(mean_all_by_combination, input_tensor))with torch.profiler.profile() as prof:    mean_all_by_torch(input_tensor)print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))with torch.profiler.profile() as prof:    mean_all_by_tensor(input_tensor)print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))with torch.profiler.profile() as prof:    mean_all_by_combination(input_tensor)print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))

import torchfrom torch.profiler import profile, profileractivitywith profile(activities=[profileractivity.cpu, profileractivity.cuda]) as prof:    for _ in range(10):        mean_tensor = torch.mean(torch.randn(10000, 10000).cuda())prof.export_chrome_trace("mean_trace.json")

import tritonimport triton.language as tlimport torch@triton.jitdef mean_kernel(    x_ptr,          # pointer to input tensor    output_ptr,     # pointer to output tensor    n_elements,     # total number of elements    BLOCK_SIZE: tl.constexpr,  # number of elements per block):    pid = tl.program_id(0)    block_start = pid * BLOCK_SIZE    block_end = tl.minimum(block_start + BLOCK_SIZE, n_elements)    acc = 0.0    for idx in range(block_start, block_end):        x = tl.load(x_ptr + idx)        acc += x    block_mean = acc / n_elements    # Store result    tl.store(output_ptr + pid, block_mean)# Wrapper functiondef triton_mean(x: torch.Tensor) -> torch.Tensor:    x = x.contiguous().view(-1)    n_elements = x.numel()    BLOCK_SIZE = 1024    grid = (triton.cdiv(n_elements, BLOCK_SIZE),)    output = torch.empty(grid[0], device=x.device, dtype=x.dtype)    mean_kernel[grid](        x_ptr=x,        output_ptr=output,        n_elements=n_elements,        BLOCK_SIZE=BLOCK_SIZE,    )    return output.sum()# Example usage:if __name__ == "__main__":    # Create test tensor    x = torch.randn(1000000, device='cuda')    # Compare results    torch_mean = torch.mean(x)    triton_mean_result = triton_mean(x)    print(f"PyTorch mean: {torch_mean}")    print(f"Triton mean: {triton_mean_result}")    print(f"Difference: {abs(torch_mean - triton_mean_result)}")