1  # API 参考
  2  
  3  项目中公共 API 的完整参考，按模块组织。
  4  
  5  ## 🔬 API 稳定性
  6  
  7  | 模块 | 稳定性 | 说明 |
  8  |------|--------|------|
  9  | `common.benchmark` | **稳定** | 核心工具，很少变更 |
 10  | `common.utils` | **稳定** | 配置加载、模型工具 |
 11  | `01_TVM` | **测试版** | 可能随 TVM 更新而演进 |
 12  | `02_ORT` | **稳定** | 标准 ORT 扩展模式 |
 13  | `03_CUTLASS` | **测试版** | 积极开发中 |
 14  | `04_cuTile` | **实验性** | 研究原型 |
 15  | `05_Triton` | **测试版** | 积极开发中 |
 16  
 17  ---
 18  
 19  ## 通用模块
 20  
 21  ### `common.benchmark.timer`
 22  
 23  用于基准测试的性能计时工具。
 24  
 25  #### `TimingResult`
 26  
 27  包含计时统计的数据类。
 28  
 29  ```python
 30  @dataclass
 31  class TimingResult:
 32      mean_ms: float        # 平均延迟（毫秒）
 33      std_ms: float         # 标准差
 34      min_ms: float         # 观察到的最小延迟
 35      max_ms: float         # 观察到的最大延迟
 36      iterations: int       # 测量迭代次数
 37  ```
 38  
 39  #### `benchmark_function()`
 40  
 41  使用自动预热和迭代控制来基准测试函数。
 42  
 43  ```python
 44  def benchmark_function(
 45      func: Callable[[], Any],
 46      warmup_iters: int = 10,
 47      bench_iters: int = 100,
 48      sync_cuda: bool = True,
 49  ) -> TimingResult
 50  ```
 51  
 52  **参数：**
 53  - `func`: 要基准测试的可调用对象（应该无参数）
 54  - `warmup_iters`: 测量前的预热迭代次数（默认：10）
 55  - `bench_iters`: 测量迭代次数（默认：100）
 56  - `sync_cuda`: 如可用，在前后同步 CUDA（默认：True）
 57  
 58  **返回：** 包含计时统计的 `TimingResult`
 59  
 60  **示例：**
 61  ```python
 62  from common.benchmark import benchmark_function
 63  import torch
 64  
 65  def matmul_op():
 66      a = torch.randn(1024, 1024, device='cuda')
 67      b = torch.randn(1024, 1024, device='cuda')
 68      return torch.matmul(a, b)
 69  
 70  result = benchmark_function(matmul_op, warmup_iters=5, bench_iters=50)
 71  print(f"Mean: {result.mean_ms:.3f} ms")
 72  ```
 73  
 74  **复杂度：**
 75  - 时间：O(迭代次数 × 函数时间)
 76  - 空间：O(1) 额外
 77  
 78  #### `benchmark_cuda_function()`
 79  
 80  显式流同步的 CUDA 优化版本。
 81  
 82  ```python
 83  def benchmark_cuda_function(
 84      func: Callable[[], Any],
 85      warmup_iters: int = 10,
 86      bench_iters: int = 100,
 87  ) -> TimingResult
 88  ```
 89  
 90  **说明：** 始终同步 CUDA 设备。用于 `func` 涉及 CUDA 操作时。
 91  
 92  #### `CUDATimer`
 93  
 94  用于细粒度 CUDA 计时的上下文管理器。
 95  
 96  ```python
 97  class CUDATimer:
 98      def __enter__(self) -> "CUDATimer": ...
 99      def __exit__(self, *args) -> None: ...
100      def elapsed_ms(self) -> float: ...
101  ```
102  
103  **示例：**
104  ```python
105  from common.benchmark.timer import CUDATimer
106  
107  with CUDATimer() as timer:
108      result = model(input_tensor)
109      torch.cuda.synchronize()
110  
111  print(f"Elapsed: {timer.elapsed_ms():.3f} ms")
112  ```
113  
114  ---
115  
116  ### `common.benchmark.reporter`
117  
118  基准测试结果报告和序列化工具。
119  
120  #### `BenchmarkEntry`
121  
122  单个基准测试结果条目。
123  
124  ```python
125  @dataclass
126  class BenchmarkEntry:
127      module_name: str
128      test_name: str
129      hardware: str
130      latency_ms: float
131      throughput: float | None = None
132      memory_mb: float | None = None
133      timestamp: str = field(default_factory=lambda: datetime.now().isoformat())
134  ```
135  
136  #### `BenchmarkReport`
137  
138  支持序列化的多条基准条目容器。
139  
140  ```python
141  class BenchmarkReport:
142      def __init__(self) -> None: ...
143      
144      def add_entry(self, entry: BenchmarkEntry) -> None: ...
145      def to_markdown_table(self) -> str: ...
146      def to_dict(self) -> dict: ...
147      def to_json(self, indent: int = 2) -> str: ...
148      def save_json(self, path: str) -> None: ...
149      
150      @classmethod
151      def from_json(cls, path: str) -> "BenchmarkReport": ...
152  ```
153  
154  **示例：**
155  ```python
156  from common.benchmark import BenchmarkEntry, BenchmarkReport
157  
158  report = BenchmarkReport()
159  report.add_entry(BenchmarkEntry(
160      module_name="TVM",
161      test_name="resnet50_baseline",
162      hardware="A100",
163      latency_ms=8.5,
164      throughput=117.6
165  ))
166  
167  # 导出为 Markdown
168  print(report.to_markdown_table())
169  
170  # 保存为 JSON 供后续分析
171  report.save_json("benchmark_results.json")
172  ```
173  
174  #### `generate_comparison_table()`
175  
176  生成计时结果的并排对比。
177  
178  ```python
179  def generate_comparison_table(
180      results: dict[str, TimingResult],
181      baseline_key: str = "pytorch_eager",
182  ) -> str
183  ```
184  
185  **参数：**
186  - `results`: 名称到 TimingResult 的字典映射
187  - `baseline_key`: 用作加速比计算基准的键
188  
189  **返回：** Markdown 格式的对比表格
190  
191  #### `format_environment_info()`
192  
193  捕获系统环境信息。
194  
195  ```python
196  def format_environment_info() -> dict[str, str]
197  ```
198  
199  **返回：** 包含 Python 版本、CUDA 版本、GPU 信息等的字典
200  
201  ---
202  
203  ### `common.utils.config_loader`
204  
205  YAML 配置加载工具。
206  
207  #### `load_benchmark_config()`
208  
209  ```python
210  def load_benchmark_config(config_path: str | None = None) -> dict[str, Any]
211  ```
212  
213  **参数：**
214  - `config_path`: YAML 配置路径。如为 None，使用 `configs/benchmark.yaml`
215  
216  **返回：** 配置字典
217  
218  #### `load_module_defaults()`
219  
220  ```python
221  def load_module_defaults(config_path: str | None = None) -> dict[str, Any]
222  ```
223  
224  **参数：**
225  - `config_path`: YAML 配置路径。如为 None，使用 `configs/module_defaults.yaml`
226  
227  **返回：** 模块默认值字典
228  
229  #### `get_module_config()`
230  
231  ```python
232  def get_module_config(
233      module_name: str,
234      config_path: str | None = None
235  ) -> dict[str, Any]
236  ```
237  
238  **参数：**
239  - `module_name`: 模块名称（如 "01_TVM", "02_ORT"）
240  - `config_path`: 可选的配置路径覆盖
241  
242  **返回：** 模块特定配置
243  
244  ---
245  
246  ### `common.utils.model_loader`
247  
248  模型加载工具。
249  
250  #### `ModelInfo`
251  
252  ```python
253  @dataclass
254  class ModelInfo:
255      name: str
256      input_shape: tuple[int, ...]
257      num_params: int
258      framework: str
259  ```
260  
261  #### `load_model()`
262  
263  ```python
264  def load_model(
265      model_name: str = "resnet50",
266      framework: str = "pytorch",
267      pretrained: bool = True,
268      device: str = "cuda",
269  ) -> tuple[Any, ModelInfo]
270  ```
271  
272  **参数：**
273  - `model_name`: 模型架构（resnet50, resnet18 等）
274  - `framework`: 使用的框架（pytorch, onnx）
275  - `pretrained`: 加载预训练权重
276  - `device`: 加载模型的设备
277  
278  **返回：** (模型, ModelInfo) 元组
279  
280  **支持的模型：**
281  - resnet18, resnet50, resnet101
282  - mobilenet_v2
283  - efficientnet_b0
284  
285  ---
286  
287  ### `common.utils.gelu`
288  
289  GELU 激活函数参考实现。
290  
291  #### `gelu_tanh_approx()`
292  
293  使用 tanh 近似计算 GELU。
294  
295  ```python
296  def gelu_tanh_approx(x: np.ndarray) -> np.ndarray
297  ```
298  
299  **参数：**
300  - `x`: 任意形状的输入数组
301  
302  **返回：** 与输入形状相同的 GELU 激活输出数组
303  
304  #### `gelu_erf()`
305  
306  使用精确误差函数计算 GELU。
307  
308  ```python
309  def gelu_erf(x: np.ndarray) -> np.ndarray
310  ```
311  
312  #### `gelu_backward()`
313  
314  计算 GELU 对输入的梯度。
315  
316  ```python
317  def gelu_backward(x: np.ndarray, grad_output: np.ndarray) -> np.ndarray
318  ```
319  
320  ---
321  
322  ### `common.utils.path_utils`
323  
324  路径工具函数。
325  
326  #### `get_project_root()`
327  
328  获取项目根目录。
329  
330  ```python
331  def get_project_root() -> Path
332  ```
333  
334  **返回：** 项目根目录的 Path 对象
335  
336  #### `ensure_project_root_in_path()`
337  
338  确保项目根目录在 sys.path 中。
339  
340  ```python
341  def ensure_project_root_in_path() -> None
342  ```
343  
344  ---
345  
346  ## 模块 01：TVM 端到端优化
347  
348  ### `1_import_and_baseline.py`
349  
350  #### `load_pytorch_model()`
351  
352  ```python
353  def load_pytorch_model(
354      model_name: str = "resnet50",
355      pretrained: bool = True,
356      device: str = "cuda",
357  ) -> tuple[nn.Module, tuple[int, ...]]
358  ```
359  
360  **返回：** (模型, 输入形状)
361  
362  #### `convert_to_relay()`
363  
364  ```python
365  def convert_to_relay(
366      model: nn.Module,
367      input_shape: tuple[int, ...],
368      input_name: str = "input0",
369  ) -> tuple[tvm.IRModule, dict]
370  ```
371  
372  **返回：** (relay 模块, 参数)
373  
374  **示例：**
375  ```python
376  model = torchvision.models.resnet50(pretrained=True)
377  mod, params = convert_to_relay(model, (1, 3, 224, 224))
378  ```
379  
380  ### `2_auto_scheduler_tuning.py`
381  
382  #### `run_auto_scheduler()`
383  
384  ```python
385  def run_auto_scheduler(
386      mod: tvm.IRModule,
387      params: dict,
388      target: str = "cuda",
389      num_trials: int = 1000,
390      log_file: str = "tuning_logs/resnet50.json",
391      early_stopping: int = 600,
392  ) -> tvm.IRModule
393  ```
394  
395  **参数：**
396  - `mod`: Relay IR 模块
397  - `params`: 模型参数
398  - `target`: 目标设备（"cuda", "llvm"）
399  - `num_trials`: 调优试验次数
400  - `log_file`: 调优日志保存路径
401  - `early_stopping`: N 次试验无改进后停止
402  
403  **返回：** 优化后的 IR 模块
404  
405  ### `3_tensorir_manual_schedule.py`
406  
407  #### `create_matmul_module()`
408  
409  ```python
410  def create_matmul_module(
411      m: int,
412      n: int,
413      k: int
414  ) -> tvm.IRModule
415  ```
416  
417  **返回：** M×N×K 矩阵乘法的 TensorIR 模块
418  
419  #### `apply_schedule_primitives()`
420  
421  ```python
422  def apply_schedule_primitives(
423      mod: tvm.IRModule,
424      block_size: int = 32,
425      k_tile: int = 8,
426  ) -> tvm.tir.Schedule
427  ```
428  
429  **参数：**
430  - `mod`: 输入 TensorIR 模块
431  - `block_size`: 线程块瓦片大小
432  - `k_tile`: K 维度瓦片大小
433  
434  **返回：** 调度后的 TensorIR
435  
436  #### `apply_gpu_schedule()`
437  
438  高级 GPU 调度与线程绑定。
439  
440  ```python
441  def apply_gpu_schedule(
442      mod: tvm.IRModule,
443      block_m: int = 128,
444      block_n: int = 128,
445      thread_tile: int = 32,
446  ) -> tvm.tir.Schedule
447  ```
448  
449  #### `benchmark_schedule()`
450  
451  ```python
452  def benchmark_schedule(
453      sch: tvm.tir.Schedule,
454      target: str = "llvm",
455      m: int = 1024,
456      n: int = 1024,
457      k: int = 1024,
458      warmup_iters: int = 5,
459      bench_iters: int = 20,
460  ) -> TimingResult
461  ```
462  
463  #### `verify_correctness()`
464  
465  ```python
466  def verify_correctness(
467      sch: tvm.tir.Schedule,
468      target: str = "llvm",
469      m: int = 1024,
470      n: int = 1024,
471      k: int = 1024,
472      rtol: float = 1e-5,
473  ) -> bool
474  ```
475  
476  **返回：** 如数值误差在容差内返回 True
477  
478  ---
479  
480  ## 模块 02：ORT 自定义 CUDA 算子
481  
482  ### CUDA 内核
483  
484  #### `LaunchGeluKernel()`
485  
486  C++ 签名：
487  ```cpp
488  cudaError_t LaunchGeluKernel(
489      const float* input,
490      float* output,
491      int64_t count,
492      cudaStream_t stream
493  );
494  ```
495  
496  **参数：**
497  - `input`: 输入数组设备指针
498  - `output`: 输出数组设备指针
499  - `count`: 元素数量
500  - `stream`: 异步执行的 CUDA 流
501  
502  **返回：** `cudaError_t` 状态
503  
504  ### C++ 自定义算子
505  
506  #### `GeluCustomOp`
507  
508  ```cpp
509  struct GeluCustomOp : Ort::CustomOpBase<GeluCustomOp, GeluKernel> {
510      const char* GetName() const;
511      const char* GetExecutionProviderType() const;
512      size_t GetInputTypeCount() const;
513      ONNXTensorElementDataType GetInputType(size_t index) const;
514      size_t GetOutputTypeCount() const;
515      ONNXTensorElementDataType GetOutputType(size_t index) const;
516  };
517  ```
518  
519  ### Python API
520  
521  #### `register_custom_ops()`
522  
523  ```python
524  def register_custom_ops(
525      session_options: ort.SessionOptions
526  ) -> str
527  ```
528  
529  **参数：**
530  - `session_options`: ONNX Runtime 会话选项
531  
532  **返回：** 注册库的路径
533  
534  **示例：**
535  ```python
536  import onnxruntime as ort
537  from custom_ops import register_custom_ops
538  
539  sess_options = ort.SessionOptions()
540  lib_path = register_custom_ops(sess_options)
541  session = ort.InferenceSession("model.onnx", sess_options)
542  ```
543  
544  #### `create_gelu_test_model()`
545  
546  ```python
547  def create_gelu_test_model(
548      input_shape: list[int] = [1, 256],
549      output_path: str | None = None,
550  ) -> onnx.ModelProto
551  ```
552  
553  **参数：**
554  - `input_shape`: 输入张量形状
555  - `output_path`: 可选的模型保存路径
556  
557  **返回：** 包含自定义 GELU 算子的 ONNX 模型
558  
559  #### `run_inference()`
560  
561  ```python
562  def run_inference(
563      model_path: str,
564      input_data: np.ndarray,
565      use_cuda: bool = True,
566  ) -> np.ndarray
567  ```
568  
569  **返回：** 推理输出 numpy 数组
570  
571  ---
572  
573  ## 模块 03：CUTLASS Hopper GEMM
574  
575  ### `hopper_gemm.h`
576  
577  #### `GemmConfig`
578  
579  ```cpp
580  struct GemmConfig {
581      int M{4096};        // 矩阵 M 维度
582      int N{4096};        // 矩阵 N 维度
583      int K{4096};        // 矩阵 K 维度
584      int block_M{128};   // 线程块 M 瓦片
585      int block_N{128};   // 线程块 N 瓦片
586      int block_K{64};    // 线程块 K 瓦片
587      float alpha{1.0f};  // C = alpha * AB + beta * C
588      float beta{0.0f};
589  };
590  ```
591  
592  #### `gemm_wgmma()`
593  
594  主 CUTLASS 支持的 GEMM 函数。
595  
596  ```cpp
597  cudaError_t gemm_wgmma(
598      const half* A,
599      const half* B,
600      half* C,
601      const GemmConfig& config,
602      cudaStream_t stream = nullptr
603  );
604  ```
605  
606  **约定：** 所有矩阵（A, B, C）均为行主序。
607  
608  **参数：**
609  - `A`: 行主序 M×K 矩阵
610  - `B`: 行主序 K×N 矩阵
611  - `C`: 行主序 M×N 输出
612  - `config`: GEMM 配置
613  - `stream`: CUDA 流
614  
615  **返回：** `cudaError_t`
616  
617  **示例：**
618  ```cpp
619  GemmConfig config{.M = 4096, .N = 4096, .K = 4096};
620  half *A, *B, *C;  // 分配并初始化
621  
622  cudaError_t err = gemm_wgmma(A, B, C, config);
623  if (err != cudaSuccess) {
624      // 处理错误
625  }
626  ```
627  
628  #### `gemm_cublas()`
629  
630  cuBLAS 参考实现（行主序接口）。
631  
632  ```cpp
633  cudaError_t gemm_cublas(
634      const half* A,
635      const half* B,
636      half* C,
637      const GemmConfig& config,
638      cudaStream_t stream = nullptr
639  );
640  ```
641  
642  #### `BenchmarkResult`
643  
644  ```cpp
645  struct BenchmarkResult {
646      float latency_ms;       // 平均延迟
647      float tflops;          // 达到的 TFLOPS
648      float bandwidth_gbps;  // 内存带宽
649      float efficiency;      // 理论峰值的百分比
650  };
651  ```
652  
653  #### `benchmark_gemm()` / `benchmark_cublas()`
654  
655  ```cpp
656  BenchmarkResult benchmark_gemm(
657      const GemmConfig& config,
658      int warmup_iters = 10,
659      int bench_iters = 100
660  );
661  
662  BenchmarkResult benchmark_cublas(
663      const GemmConfig& config,
664      int warmup_iters = 10,
665      int bench_iters = 100
666  );
667  ```
668  
669  #### 工具函数
670  
671  ```cpp
672  bool is_hopper_gpu();                    // 检查是否为 SM90
673  int get_sm_version();                    // 获取计算能力
674  float calculate_tflops(const GemmConfig& config, float latency_ms);
675  float calculate_bandwidth(const GemmConfig& config, float latency_ms);
676  bool verify_gemm(const half* C, const half* C_ref, int M, int N, float rtol = 1e-3f);
677  ```
678  
679  ---
680  
681  ## 模块 04：cuTile 下一代 CUDA
682  
683  ### `tile_gemm.py`
684  
685  #### `TileArray`
686  
687  ```python
688  class TileArray:
689      def __init__(
690          self,
691          data: np.ndarray,
692          tile_shape: Optional[Tuple[int, ...]] = None
693      ): ...
694      
695      @property
696      def shape(self) -> Tuple[int, ...]: ...
697      
698      @property
699      def tile_shape(self) -> Tuple[int, ...]: ...
700  ```
701  
702  #### `tiled_gemm()`
703  
704  ```python
705  @tile_kernel
706  def tiled_gemm(
707      A: TileArray,
708      B: TileArray,
709      block_shape: tuple[int, int, int] = (128, 128, 32),
710  ) -> TileArray
711  ```
712  
713  **参数：**
714  - `A`: 左矩阵（M×K）
715  - `B`: 右矩阵（K×N）
716  - `block_shape`: 瓦片维度（block_M, block_N, block_K）
717  
718  **返回：** 结果矩阵 C（M×N）
719  
720  ---
721  
722  ## 模块 05：Triton GPU 内核
723  
724  ### `triton_kernels.gemm`
725  
726  #### `matmul()`
727  
728  使用 Triton 计算矩阵乘法 C = A @ B。
729  
730  ```python
731  def matmul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor
732  ```
733  
734  **参数：**
735  - `a`: 形状为 (M, K) 的输入张量
736  - `b`: 形状为 (K, N) 的输入张量
737  
738  **返回：** 形状为 (M, N) 的输出张量
739  
740  **示例：**
741  ```python
742  import torch
743  from triton_kernels import matmul
744  
745  a = torch.randn(1024, 512, device='cuda', dtype=torch.float16)
746  b = torch.randn(512, 2048, device='cuda', dtype=torch.float16)
747  c = matmul(a, b)
748  ```
749  
750  #### `matmul_autotuned()`
751  
752  使用自动调优的 Triton 内核计算矩阵乘法。
753  
754  ```python
755  def matmul_autotuned(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor
756  ```
757  
758  **说明：** 自动根据输入维度和硬件选择最优块大小。推荐生产使用。
759  
760  ### `triton_kernels.flash_attention`
761  
762  #### `flash_attention()`
763  
764  内存高效的注意力计算。
765  
766  ```python
767  def flash_attention(
768      q: torch.Tensor,
769      k: torch.Tensor,
770      v: torch.Tensor,
771      softmax_scale: float | None = None,
772  ) -> torch.Tensor
773  ```
774  
775  **参数：**
776  - `q`: 查询张量 (batch, heads, seq_len, head_dim)
777  - `k`: 键张量 (batch, heads, seq_len, head_dim)
778  - `v`: 值张量 (batch, heads, seq_len, head_dim)
779  - `softmax_scale`: 可选的 softmax 缩放因子
780  
781  **返回：** 注意力输出张量
782  
783  ---
784  
785  ## 版本历史
786  
787  ### v1.0.0 中的 API 变更
788  
789  - ✅ **稳定**：`common.benchmark` 接口已定型
790  - ⚠️ **测试版**：`03_CUTLASS_Hopper_GEMM` 接口可能变更
791  - 🧪 **实验性**：`04_cuTile` 为研究原型
792  
793  ### 弃用
794  
795  当前无。
796  
797  ---
798  
799  *实现细节请参见源代码和 [架构文档](./architecture)。*