Rosetta Experimental Learning Gpu Operator 251230

🧪 GPU 算子底层原理学习 - 从 PyTorch 到 GPU 硬件

技能: rosetta-experimental-learning 目标系统: PyTorch GPU 算子执行流程 学习模式: 完整实验（complete） 学习目标:

• ✅ 理解编译流程（Python → PTX → SASS）
• ✅ 性能优化（内存带宽、计算吞吐、occupancy）
• ✅ GPU 硬件原理（SM、Warp、Thread、Memory Hierarchy）
• ✅ 自定义算子开发（Triton/CUDA） 硬件: NVIDIA GPU 生成时间: 2025-12-30

🎯 为什么要用罗塞塔方法？

GPU 算子学习的现实痛点：

📐 第一步：定层级

对于 PyTorch GPU 算子，完整的数据流转层级：

核心公式：

Python 代码 → Triton/CUDA Kernel → PTX 指令 → SASS 机器码 → GPU 执行

🚧 第二步：定关卡

🔗 关卡因果链路（横向展示）

Python 代码执行
    ↓
torch.matmul(a, b)
    ↓
检查 tensor 设备、dtype
    ↓
┌─ CUDA Tensor ─────────────┐   ┌─ CPU Tensor ────┐
↓                              ↓                    ↓
触发 CUDA Kernel               CPU 执行
    ↓                              ↓
Dispatcher 分发
    ↓
选择实现（CUDA/HIP/MPS）
    ↓
┌─ 有 Triton/CUDA 实现 ──────┐   ┌─ fallback ───────┐
↓                              ↓                    ↓
编译 Kernel（首次）           使用通用实现
    ↓                              ↓
缓存 Kernel 对象              可能较慢
    ↓
配置 Grid/Block
    ↓
启动 Kernel <<<Grid, Block>>>()
    ↓
GPU 接收命令
    ↓
Warp 调度
    ↓
分配 SM
    ↓
┌─ Memory 访问 ────┐   ┌─ 计算 ──────┐
↓                    ↓                ↓
Cache Hit?         完成           访问 Global Memory
↓                                      ↓
┌─ 是 ─┐       ┌─ 否 ─┐                 ↓
↓       ↓       ↓       ↓           增加 latency
返回   继续   数据   访问 Global
            Memory   Memory

🛠️ 第三步：架工具

关卡1：PyTorch Dispatch 层工具

关卡2：Triton/CUDA 编译工具

关卡3：PTX/SASS 查看工具

关卡4：Kernel Launch 工具

关卡5：GPU 执行分析工具

关卡6：Memory 分析工具

💉 第四步：投示踪（按层级追踪）

阶段1：准备示踪剂

目标代码：矩阵乘法

# tracer_matmul.py
import torch
import torch.cuda.nvtx as nvtx

def run_matmul():
    M, N, K = 1024, 1024, 1024

    # 示踪剂1: 特殊大小的矩阵
    a = torch.randn(M, K, device='cuda', dtype=torch.float16)
    b = torch.randn(K, N, device='cuda', dtype=torch.float16)

    # 示踪剂2: NVTX 标记
    nvtx.range_push("matmul_start")
    c = torch.matmul(a, b)
    nvtx.range_pop()

    return c

# 预热（触发编译）
for _ in range(3):
    run_matmul()

result = run_matmul()
torch.cuda.synchronize()
print("✓ MatMul completed")

阶段2：按层级追踪

🔷 层级1：Python 层

观测点：函数调用、张量对象

操作：

# 基础运行
python tracer_matmul.py

# 查看张量信息
python -c "
import torch
a = torch.randn(1024, 1024, device='cuda', dtype=torch.float16)
print(f'Device: {a.device}')
print(f'Dtype: {a.dtype}')
print(f'Shape: {a.shape}')
"

预期输出：

Device: cuda:0
Dtype: torch.float16
Shape: torch.Size([1024, 1024])

🔗 认知映射

Python 表现              实际含义
torch.matmul(a, b)    →  触发 C++ aten::mm
        ↓
  理解：Python 只是接口，真实执行在 C++ 层

🔷 层级2：ATen/Dispatch 层

观测点：算子分发、backend 选择

操作：

TORCH_LOGS="dispatch" python tracer_matmul.py 2>&1 | grep -A 5 "matmul"

预期输出：

[debug] torch._C._EngineBase::_operator_call: aten::matmul
[info] aten::matmul: dispatching to CUDA backend

🔗 认知映射

Python 层              C++ 层
torch.matmul()    →  aten::matmul
                          ↓
                    Dispatcher
                          ↓
                    CUDA Kernel

🔷 层级3：Triton/CUDA Kernel

观测点：查看调用的是哪个 kernel

操作：

import torch.autograd.profiler as profiler

M, N, K = 1024, 1024, 1024
a = torch.randn(M, K, device='cuda', dtype=torch.float16)
b = torch.randn(K, N, device='cuda', dtype=torch.float16)

with profiler.profile(
    activities=[profiler.ProfilerActivity.CUDA],
    record_shapes=True
) as p:
    for _ in range(10):
        c = torch.matmul(a, b)

print(profiler.tensorboard_profiler.Table(p.key_averages().table(
    sort_by="self_cuda_time_total", row_limit=10
)))

预期输出：

Name                                  Self CUDA    Calls
ampere_sgemm_128x128_tn               1.150ms      10

发现：PyTorch 调用了 ampere_sgemm_128x128_tn kernel（针对 Ampere 架构优化）

🔗 认知映射

Python                 Kernel
torch.matmul()    →  ampere_sgemm_128x128_tn
                              ↓
  理解：PyTorch 根据架构选择最优 kernel
        128x128 = tile size
        tn = transpose(normal)

🔷 层级4：PTX 层

观测点：查看 PTX 指令（虚拟 ISA）

操作：

# 反汇编 PyTorch 库
cuobjdump --ptx /path/to/libtorch_cuda.so | grep -A 100 "sgemm_128x128" > sgemm.ptx

# 或直接查看 PTX
nvdisasm -c -g /path/to/libtorch_cuda.so | grep -A 50 "ampere_sgemm"

预期 PTX 示例：

// 加载数据
ld.global.nc.f32  {r0, r1, r2, r3}, [x_ptr];
ld.global.nc.f32  {r4, r5, r6, r7}, [y_ptr];

// Tensor Core 操作
mma.sync.aligned.m16n8k8.row.col.f32.f16.f32.f32
    {r0, r1, r2, r3},
    {b0, b1, b2, b3},
    {b4, b5, b6, b7},
    {r8, r9, r10, r11};

// 存回结果
st.global.f32  [z_ptr], r0;

🔗 认知映射

高操作              PTX 指令
矩阵乘法       →  mma.sync (Tensor Core)
加载全局内存    →  ld.global.nc
存储结果        →  st.global
        ↓
  理解：PTX 是 NVIDIA 的虚拟汇编
        mma.sync = Matrix Multiply-Accumulate Sync
        nc = Non-Coherent

🔷 层级5：SASS 层

观测点：查看真实执行的机器码

操作：

cuobjdump -sass /path/to/libtorch_cuda.so | grep -A 100 "ampere_sgemm_128x128_tn" > sgemm.sass

预期 SASS 示例（Ampere 架构）：

/*0000*/                   MOV R1, c[0x0][0x148];
/*0008*/                   S2R R0, SR_TID.X;
/*0018*/                   IMAD.U32.U32 R4, R3, param[0x2], R0;
/*0028*/                   LDG.E64.ULK128 R8, [R8];
/*0038*/                   HMMA.16832.F32.F32.F32.F32 R8, R4, R12, R8;
/*0040*/                   STS.E64 [R20+0x200], R8;
/*0048*/                   LDS.E64.ULK128 R8, [R20+0x200];
/*0050*/                   STG.E64 [R16+0x100], R8;

关键指令：

🔗 认知映射

PTX 虚拟指令              SASS 真实指令
mma.sync           →  HMMA.16832.F32.F32.F32.F32
ld.global          →  LDG.E64.ULK128
st.global          →  STG.E64
        ↓
  理解：SASS 是 GPU 真正执行的机器码
        HMMA = Tensor Core 硬件指令
        E64 = 64-bit 访问

🔷 层级6：GPU 硬件层

观测点：SM、Warp、Memory 性能

操作：

# 完整分析
ncu --set full --target-processes all python tracer_matmul.py

# 查看 Occupancy
ncu --metrics sm__warps_active.avg.pct_of_peak python tracer_matmul.py

# 查看 Memory Throughput
ncu --metrics dram__throughput.avg.pct_of_peak python tracer_matmul.py

# 查看 Warp Efficiency
ncu --metrics sm__warps_active.avg.per_cycle_active python tracer_matmul.py

关键指标：

🔗 认知映射（硬件层）

软件配置                      硬件行为
BlockDim(128)          →  128 threads / block
GridDim(256)           →  256 blocks → 2.4 SM (A100)
                              ↓
每 SM 执行              →  307 threads
                              ↓
分成 Warp               →  307 / 32 = 9.6 warps
                              ↓
Warp Efficiency         →  9.6 / 32 = 30% (偏低)
        ↓
  理解：配置未充分利用 SM，需优化

阶段3：性能瓶颈分析

决策树：

Nsight Compute 分析结果
    ↓
┌─ Memory Throughput >70% peak ─────────┐
│                                      │
是 (Memory Bound)                    否 (Compute Bound)
    ↓                                  ↓
检查 Cache Hit Rate                  检查 Warp Efficiency
    ↓                                  ↓
  <80% 需优化                          <60% 需优化
    ↓                                  ↓
优化策略：                           优化策略：
• 使用 Shared Memory                • 增加并行度
• 合并内存访问                      • 减少分支分歧
• 向量化读取                        • 增加 occupancy

🎯 实验场景

实验1：理解 Tile Size 对性能的影响

代码：

import torch
import torch.utils.benchmark as benchmark

def bench_matmul(M, N, K):
    a = torch.randn(M, K, device='cuda', dtype=torch.float16)
    b = torch.randn(K, N, device='cuda', dtype=torch.float16)

    timer = benchmark.Timer(
        stmt="torch.matmul(a, b)",
        globals={'a': a, 'b': b},
    )
    return timer.blocked_autorange(min_run_time=1)

sizes = [(512, 512, 512), (1024, 1024, 1024), (2048, 2048, 2048)]

for M, N, K in sizes:
    time = bench_matmul(M, N, K)
    gflops = (2 * M * N * K) / (time.mean * 1e9)
    print(f"{M}x{N}x{K}: {time.mean*1000:.2f} ms, {gflops:.1f} GFLOPS")

预期发现：

实验2：理解 Tensor Core

对比 FP32 vs FP16：

import torch

M, N, K = 2048, 2048, 2048

# FP32 (无 Tensor Core)
a_fp32 = torch.randn(M, K, device='cuda', dtype=torch.float32)
b_fp32 = torch.randn(K, N, device='cuda', dtype=torch.float32)

# FP16 (有 Tensor Core)
a_fp16 = torch.randn(M, K, device='cuda', dtype=torch.float16)
b_fp16 = torch.randn(K, N, device='cuda', dtype=torch.float16)

# 测量时间
# ...

# 结果
# FP32: 50 GFLOPS
# FP16: 1100 GFLOPS (22x 加速！)

Nsight Compute 分析：

FP32 版本：
  smsp__pipe_tensor_cycles_active.avg.pct_of_peak: 0% (未使用)

FP16 版本：
  smsp__pipe_tensor_cycles_active.avg.pct_of_peak: 90% (充分使用)

认知映射：

数据类型              硬件单元
FP32              →  CUDA Core (标量/向量)
FP16/BF16         →  Tensor Core (4x4x4 矩阵)
INT8              →  Tensor Core (8x8x8 矩阵)

💡 下一步学习路径

Week 1: 观测工具

• 使用 torch.profiler 分析 PyTorch 代码
• 使用 Nsight Compute 分析 kernel 性能
• 理解 Occupancy、Memory Throughput、Warp Efficiency

Week 2: Triton 编程

• 实现简单的向量加法
• 实现 Softmax
• 实现矩阵乘法

Week 3: CUDA 编程

• 理解 Thread Hierarchy（Grid/Block/Warp）
• 使用 Shared Memory 优化矩阵乘法
• 理解 Memory Coalescing

Week 4: 性能优化

• 分析实际项目的瓶颈
• 优化内存访问模式
• 调整 Block/Grid 配置

记住：GPU 性能优化的核心是隐藏延迟：

1. Memory Latency → 用 Cache、Shared Memory
2. Compute Latency → 用 Tensor Core
3. Warp Divergence → 减少分支
4. Synchronization → 异步执行