架构设计

Tiny-DL-Inference 的系统架构和设计原则。

概述

Tiny-DL-Inference 采用分层架构，将设备管理、张量操作、神经网络算子和高级推理执行的关注点分离。

┌─────────────────────────────────────────────────────────────────┐
│                    应用层                                         │
│  ┌─────────────┐  ┌─────────────┐  ┌─────────────────────────┐  │
│  │ 模型加载器   │  │  基准测试   │  │      推理运行器          │  │
│  └─────────────┘  └─────────────┘  └─────────────────────────┘  │
├─────────────────────────────────────────────────────────────────┤
│                    算子层                                         │
│  ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐ ┌───────────────┐  │
│  │ 卷积   │ │ 池化   │ │  ReLU  │ │ Softmax│ │卷积+偏置+ReLU │  │
│  └────────┘ └────────┘ └────────┘ └────────┘ └───────────────┘  │
├─────────────────────────────────────────────────────────────────┤
│                    核心层                                         │
│  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐  │
│  │      张量       │  │    着色器管理    │  │    GPU上下文    │  │
│  └─────────────────┘  └─────────────────┘  └─────────────────┘  │
├─────────────────────────────────────────────────────────────────┤
│                    设备层                                         │
│  ┌─────────────────────────────────────────────────────────────┐│
│  │                    WebGPU 适配器/设备                        ││
│  └─────────────────────────────────────────────────────────────┘│
└─────────────────────────────────────────────────────────────────┘

层描述

设备层

与 GPU 硬件交互的基础层。

组件：

GPUAdapter - 物理 GPU 选择
GPUDevice - 带命令队列的逻辑设备
GPUQueue - 命令提交

职责：

硬件抽象
命令缓冲区提交
同步原语

核心层

提供 GPU 编程的基本抽象。

GPUContext

集中式资源管理：

设备初始化和配置
缓冲区分配和生命周期
命令编码和提交
延迟资源销毁

typescript

class GPUContext {
  private adapter: GPUAdapter;
  private device: GPUDevice;
  private pendingCleanup: Set<Promise<void>>;
  
  // 延迟初始化
  async initialize(config?: GPUContextConfig): Promise<void>;
  
  // 资源管理
  deferDestroy(buffer: GPUBuffer): void;
  async sync(): Promise<void>;
  
  // 命令提交
  createCommandEncoder(): GPUCommandEncoder;
  submit(buffers: GPUCommandBuffer[]): void;
}

Tensor

多维数组抽象：

GPU 缓冲区所有权
布局元数据（NCHW/NHWC）
通过视图零拷贝重塑

typescript

class Tensor {
  readonly shape: number[];
  readonly layout: DataLayout;
  readonly buffer: GPUBuffer;
  
  // 数据传输
  upload(data: Float32Array): Promise<void>;
  download(): Promise<Float32Array>;
  
  // 视图操作（零拷贝）
  reshape(newShape: number[]): Tensor;
  
  // 布局转换
  convertLayout(target: DataLayout): Promise<Tensor>;
}

视图系统：

视图共享相同的 GPU 缓冲区
destroy() 对视图无操作
实现高效的 Flatten 操作

算子层

神经网络操作实现为计算着色器。

Operator 基类

typescript

abstract class Operator {
  protected context: GPUContext;
  protected pipeline: GPUComputePipeline;
  protected shaderModule: GPUShaderModule;
  
  // 着色器编译
  protected abstract compileShader(): string;
  
  // 形状推断
  abstract computeOutputShape(
    inputShape: TensorShape,
    params?: OperatorParams
  ): TensorShape;
  
  // 执行
  abstract forward(
    inputs: Tensor[],
    params?: OperatorParams
  ): Promise<Tensor>;
  
  // 延迟初始化
  protected ensureInitialized(): void;
}

着色器编译流程

Operator.forward()
    ↓
ensureInitialized() ──首次调用──→ compileShader()
    ↓                                    ↓
    ←←←←←←←←← 管道创建 ←←←←←←←←←
    ↓
编码命令 → 提交 → 返回输出张量

应用层

用于模型执行的高级 API。

InferenceEngine

编排整个推理流程：

算子注册和管理
模型加载和权重初始化
逐层执行
内存清理

typescript

class InferenceEngine {
  private context: GPUContext;
  private operators: Map<string, Operator>;
  private weights: Map<string, Tensor>;
  
  async infer(input: Tensor): Promise<Tensor> {
    // 1. 验证输入
    // 2. 按拓扑顺序执行层
    // 3. 清理中间张量
    // 4. 返回输出
  }
}

执行流程：

输入 ──┐
       ├→ [卷积] → [ReLU] → [池化] → [展平] → [全连接] → [Softmax] ──→ 输出
权重───┘

内存管理

缓冲区生命周期

创建 → 使用 → 销毁
   │      │       │
   │      │       └── deferDestroy() 在 GPU 工作后排队
   │      └────────── 编码命令，提交到队列
   └───────────────── GPUContext.createBuffer()

延迟销毁

通过延迟缓冲区销毁防止释放后使用错误：

typescript

// 安全模式
const tempBuffer = context.createBuffer({...});
context.submit(commandsUsingBuffer);
context.deferDestroy(tempBuffer);
// 在 GPU 工作完成后销毁缓冲区

视图语义

typescript

// 原始张量拥有缓冲区
const original = new Tensor(context, [1, 3, 224, 224]);

// 视图共享缓冲区（零开销）
const view = original.reshape([1, 3 * 224 * 224]);

// 清理
view.destroy();  // 无操作：不拥有缓冲区
original.destroy();  // 实际销毁 GPU 缓冲区

数据流

前向传播

输入张量 (CPU 内存)
    ↓ upload()
输入张量 (GPU 缓冲区)
    ↓ operator.forward()
命令编码器
    ↓ 编码计算通道
命令缓冲区
    ↓ submit()
GPU 队列
    ↓ 执行着色器
输出张量 (GPU 缓冲区)
    ↓ download()
输出数组 (CPU 内存)

推理流程

typescript

// 1. 初始化
const engine = new InferenceEngine();
await engine.initialize();

// 2. 加载模型
await engine.loadModel(modelDef);
//    - 分配权重张量
//    - 注册算子

// 3. 准备输入
const input = engine.tensorFromArray(data, shape);

// 4. 执行
const output = await engine.infer(input);
//    对于每层：
//    - 获取算子
//    - 执行前向传播
//    - 存储输出
//    - 销毁中间张量

// 5. 获取结果
const predictions = await output.download();

// 6. 清理
engine.destroy();

WebGPU 集成

管道创建

typescript

// 1. 编译 WGSL 着色器
const shaderModule = device.createShaderModule({
  code: wgslCode,
  label: 'operator-shader'
});

// 2. 创建绑定组布局
const bindGroupLayout = device.createBindGroupLayout({
  entries: [
    { binding: 0, visibility: COMPUTE, buffer: { type: 'storage' } },
    { binding: 1, visibility: COMPUTE, buffer: { type: 'read-only-storage' } }
  ]
});

// 3. 创建管道布局
const pipelineLayout = device.createPipelineLayout({
  bindGroupLayouts: [bindGroupLayout]
});

// 4. 创建计算管道
const pipeline = device.createComputePipeline({
  layout: pipelineLayout,
  compute: { module: shaderModule, entryPoint: 'main' }
});

调度策略

工作组大小针对 GPU 占用率优化：

wgsl

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let idx = global_id.x;
  if (idx >= numElements) { return; }
  
  // 处理 idx 位置的元素
  output[idx] = activation(input[idx]);
}

调度大小计算：

typescript

const workgroupSize = 256;
const numWorkgroups = Math.ceil(numElements / workgroupSize);
pass.dispatchWorkgroups(numWorkgroups);

优化架构

算子融合

将多个操作组合成单个 GPU 内核：

无融合：
  卷积  →  偏置  →  ReLU
    6 次内存操作（每次读/写）

融合：
  卷积+偏置+ReLU（单个内核）
    2 次内存操作
    
加速：3 倍内存带宽减少

内存布局策略

NCHW（通道优先）：

卷积操作的自然选择
通道级操作的更好缓存局部性
PyTorch 兼容

NHWC（通道最后）：

GPU 上更好的内存合并访问
TensorFlow 兼容
提供转换工具

错误处理

错误层次结构

Error
├── WebGPUNotSupportedError
├── DeviceInitializationError
├── InvalidShapeError
└── BufferSizeError

恢复策略

错误类型	恢复方式
WebGPUNotSupported	回退到 CPU 或显示消息
DeviceInitialization	使用较低功耗偏好重试
InvalidShape	操作前验证形状
BufferSize	检查数据长度是否匹配张量大小

扩展点

添加自定义算子

typescript

class CustomOperator extends Operator {
  protected compileShader(): string {
    return `
      @compute @workgroup_size(256)
      fn main(@builtin(global_invocation_id) id: vec3<u32>) {
        // 自定义计算
      }
    `;
  }
  
  computeOutputShape(input: TensorShape): TensorShape {
    // 返回输出形状
  }
  
  async forward(inputs: Tensor[], params: OperatorParams): Promise<Tensor> {
    // 实现
  }
}

自定义模型格式

扩展 ModelLoader 以支持不同格式：

typescript

class ONNXLoader extends ModelLoader {
  async loadFromONNX(path: string): Promise<ModelDefinition> {
    // 解析 ONNX 并转换为内部格式
  }
}

性能考虑

瓶颈

CPU-GPU 传输：最小化上传/下载频率
着色器编译：尽可能缓存管道
内存分配：为可变大小输入重用缓冲区
同步：批处理操作以减少同步点

最佳实践

尽可能使用融合算子
一起处理输入批次
对临时缓冲区使用 deferDestroy()
避免不必要的下载
预分配缓冲区

架构设计 ​

概述 ​

层描述 ​

设备层 ​

核心层 ​

GPUContext ​

Tensor ​

算子层 ​

Operator 基类 ​

着色器编译流程 ​

应用层 ​

InferenceEngine ​

内存管理 ​

缓冲区生命周期 ​

延迟销毁 ​

视图语义 ​

数据流 ​

前向传播 ​

推理流程 ​

WebGPU 集成 ​

管道创建 ​

调度策略 ​

优化架构 ​

算子融合 ​

内存布局策略 ​

错误处理 ​

错误层次结构 ​

恢复策略 ​

扩展点 ​

添加自定义算子 ​

自定义模型格式 ​

性能考虑 ​

瓶颈 ​

最佳实践 ​

架构设计

概述

层描述

设备层

核心层

GPUContext

Tensor

算子层

Operator 基类

着色器编译流程

应用层

InferenceEngine

内存管理

缓冲区生命周期

延迟销毁

视图语义

数据流

前向传播

推理流程

WebGPU 集成

管道创建

调度策略

优化架构

算子融合

内存布局策略

错误处理

错误层次结构

恢复策略

扩展点

添加自定义算子

自定义模型格式

性能考虑

瓶颈

最佳实践