深入PyTorch架构：从张量到神经网络的完整设计剖析

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

┌─────────────────────────────────────────────────────────────┐
│                    Python API Layer                        │  
│  torch.nn • torch.optim • torch.distributed • torch.jit   │
├─────────────────────────────────────────────────────────────┤
│                   PyTorch Python Core                      │
│        torch.Tensor • torch.autograd • torch.fx           │
├─────────────────────────────────────────────────────────────┤
│                    C++ Extension API                       │
│            Python C Extension • pybind11 Binding          │
├─────────────────────────────────────────────────────────────┤
│                      ATen Library                          │
│          Tensor Operations • Backend Dispatch             │
├─────────────────────────────────────────────────────────────┤
│                       C10 Library                          │
│        Core Tensor • Memory Management • Device API       │
├─────────────────────────────────────────────────────────────┤
│                    Hardware Backends                       │
│       CPU • CUDA • MPS • XPU • Custom Devices            │
└─────────────────────────────────────────────────────────────┘

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

// 核心张量实现
class TensorImpl {
    Storage storage_;           // 数据存储
    SymIntArrayRef sizes_;      // 张量形状
    SymIntArrayRef strides_;    // 步长信息
    DispatchKeySet key_set_;    // 分发键集合
    ScalarType dtype_;          // 数据类型
    Device device_;             // 设备信息
    // ...其他元数据
};

// 存储管理
class Storage {
    DataPtr data_ptr_;          // 数据指针
    SymInt size_bytes_;         // 字节大小
    Allocator* allocator_;      // 内存分配器
    bool resizable_;            // 是否可调整大小
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

// 算子注册示例
TORCH_LIBRARY_IMPL(aten, CPU, m) {
  m.impl("add.Tensor", &cpu_add);
  m.impl("mul.Tensor", &cpu_mul);
}

TORCH_LIBRARY_IMPL(aten, CUDA, m) {
  m.impl("add.Tensor", &cuda_add);
  m.impl("mul.Tensor", &cuda_mul);
}

1
2
3
4
5
6

class Function:
    def forward(ctx, *args):    # 前向传播
        pass
    
    def backward(ctx, *grad_outputs):  # 反向传播
        pass

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

// CUDA缓存分配器核心逻辑
class CUDACachingAllocator {
    // 内存块管理
    std::map<size_t, std::set<Block*>> cached_blocks;
    std::set<Block*> active_blocks;
    
    // 分配策略
    DataPtr allocate(size_t size) {
        // 1. 尝试从缓存获取
        // 2. 无缓存则从GPU申请
        // 3. 失败时触发垃圾回收
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

enum class DispatchKey : uint16_t {
  Undefined = 0,
  CPU,                    // CPU后端
  CUDA,                   // CUDA后端  
  XLA,                    // XLA编译
  Lazy,                   // 延迟计算
  Meta,                   // 元张量
  Autograd,               // 自动微分
  Profiler,               // 性能分析
  Tracer,                 // 图追踪
  // ... 更多分发键
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

// 算子声明
TORCH_LIBRARY(aten, m) {
  m.def("add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor");
}

// 具体实现注册
TORCH_LIBRARY_IMPL(aten, CPU, m) {
  m.impl("add.Tensor", TORCH_FN(cpu_add_tensor));
}

TORCH_LIBRARY_IMPL(aten, Autograd, m) {
  m.impl("add.Tensor", TORCH_FN(autograd_add_tensor));
}

1
2
3
4
5
6
7
8

class CUDAStream {
    cudaStream_t stream_;       // CUDA流句柄
    int device_index_;          // 设备索引
    StreamPriority priority_;   // 流优先级
    
    void synchronize();         // 同步等待
    void record_event(CUDAEvent& event);  // 记录事件
};

1
2
3
4
5
6
7
8

# Python自定义算子
@torch.library.custom_op("mylib::my_op", mutates_args=())
def my_op(x: Tensor) -> Tensor:
    return x.sin()

@my_op.register_fake
def _(x):
    return torch.empty_like(x)

1
2
3

# 使用梯度检查点节省内存
def forward_with_checkpoint(self, x):
    return checkpoint(self.layer, x)

1
2
3
4
5
6
7

with torch.profiler.profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True
) as prof:
    model(inputs)

print(prof.key_averages().table(sort_by="cuda_time_total"))

1
2
3

# 内存使用分析
torch.cuda.memory_summary()
torch.cuda.memory_stats()