书生国智科探挑战赛 · 上海人工智能实验室 × 壁仞科技
本项目完成赛道五的必选题（SpectralConv2d 算子）+ 进阶题A（PINN 热传导），核心 kernel 使用 torch.autograd.Function 实现，支持 BIREN GPU（壁仞）适配。
- SpectralConv2d / 3d: RFFT → 截断 →
ComplexMatmul2dFunction（torch.autograd.Function）→ IRFFT，相对误差 1.15e-07，ChannelSplitIrfft2 自定义 autograd Function（no_grad 推理 28ms 超越标杆） - PINN 热传导: SIREN-style sin 激活 + Adam + L-BFGS，相对 L2 误差 0.118%
- BIREN GPU 适配: compat.py 自动检测环境 + FFT fallback，显存 < 80 MB
- Agent 协作: 12 段交互日志，覆盖 6 种开发场景
track5/
├── requirements.txt # 依赖
├── spectral_conv2d/ # [必选题] FNO SpectralConv2d
│ ├── spectral_conv2d.py # 核心实现: autograd.Function + 2D/3D
│ ├── test_spectral.py # 8项测试: 形状/精度/梯度/3D/边界/kernel/DFT/性能
│ ├── test_cpp_extension.py # C++ extension 验证
│ ├── tune.py # 自动调优 (四维网格搜索)
│ └── cpp/ # torch.extension (C++/ATen)
│ ├── complex_matmul.cpp
│ ├── setup.py
│ └── complex_matmul_wrapper.py
├── pinn/ # [进阶题A] PINN 热传导
│ ├── model.py # MLP + LaplacianKernel + LaplacianStencilFunction
│ ├── train.py # Adam + L-BFGS + 解析解 + 误差计算
│ ├── visualize.py # 4 张可视化图生成
│ ├── run_pinn.py # 一键训练入口
│ ├── pinn_model.pt # 预训练权重 (L2=0.118%)
│ └── cpp/ # torch.extension (C++/ATen)
│ ├── laplacian.cpp # C++ 5-point stencil Laplacian kernel
│ └── setup.py # 编译脚本
├── biren/ # BIREN GPU 适配层
│ ├── compat.py # 环境检测 + FFT fallback + SUPA 原语
│ ├── adaptation_guide.md # 完整适配指南
│ └── supa/ # SUPA .su 原生 kernel
│ ├── complex_matmul.su # kernel 源码 (6 kernels, 2D+3D fwd/bwd)
│ ├── load_kernel.py # brcc 编译 + ctypes 加载
│ └── __init__.py
├── biren_test_results.md # BIREN 实机测试结果
├── agent_logs/ # 12 段 Agent 交互日志
│ ├── 01_operator_design.md # 算子 kernel 设计
│ ├── 02_debugging.md # 算子调试 (广播维度对齐)
│ ├── 03_model_architecture.md # 模型架构选型
│ ├── 04_performance_analysis.md # 性能瓶颈分析
│ ├── 05_visualization.md # 可视化代码生成
│ ├── 06_biren_gpu_adaptation.md # BIREN GPU 适配排查
│ ├── 07_auto_tune.md # 自动调优四维搜索
│ ├── 08_performance_audit.md # 性能数据诚实性审计
│ ├── 09_security_audit.md # 安全审计与凭据修复
│ ├── 10_performance_final.md # 性能最后冲刺
│ ├── 11_channel_split_irfft2.md # irfft2 深度优化
│ └── 12_truncated_rfft2_and_test_coverage.md # 工程收尾
└── fno/ # [新增] FNO 端到端模型
│ ├── __init__.py
│ ├── fno_model.py # FNO2d + FNO3d
│ └── run_fno.py # 训练 + 推理 benchmark
└── reports/ # 技术报告
├── spectral_conv2d_report.md
└── pinn_report.md
# 1. 安装依赖
cd track5
pip install -r requirements.txt
# 2. SpectralConv2d 测试 (8 项, ~15s)
cd spectral_conv2d
python test_spectral.py
# 3. PINN 前向推理
cd ../pinn
python -c "
import torch; from model import PINN
model = PINN(5, 128, 'sin')
model.load_state_dict(torch.load('pinn_model.pt'))
model.eval()
x = torch.linspace(0,1,100); y = torch.linspace(0,1,100)
X, Y = torch.meshgrid(x, y, indexing='ij')
T = model(X.reshape(-1,1), Y.reshape(-1,1))
print(f'Inference OK: {T.shape}')
"
# 4. BIREN GPU 兼容性检查
cd ../biren
python compat.pyInput (B, C_in, H, W)
│
▼ torch.fft.rfft2
(B, C_in, H, W//2+1) complex
│
▼ Truncate low-freq modes [modes1, modes2]
(B, C_in, modes1, modes2) complex
│
▼ ComplexMatmul2dFunction.apply() ← 核心 kernel (torch.autograd.Function)
(B, C_out, modes1, modes2) complex ← 广播乘 + reduce_sum → SUPA 映射
│
▼ Zero-pad + torch.fft.irfft2
Output (B, C_out, H, W)
复数矩阵乘 kernel:
class ComplexMatmul2dFunction(torch.autograd.Function):
"""einsum 语义: bimn,iomn->bomn (contract over C_in)"""
@staticmethod
def forward(ctx, xr, xi, wr, wi):
wr_t = wr.permute(1,0,2,3).unsqueeze(0) # align C_in dim
ac = (xr.unsqueeze(1) * wr_t).sum(dim=2) # broadcast mul + reduce
return ac - bd, ad + bc # (a+bi)(c+di)
@staticmethod
def backward(ctx, gor, goi):
# dL/dx = dL/dout * conj(W) — 复用 forward 广播模式
...| Component | Choice | Why |
|---|---|---|
| Architecture | 5 layers × 128 neurons | Exceeds requirement (≥4/≥64) |
| Activation | sin (SIREN-style) | 3× faster convergence than tanh |
| Laplacian | LaplacianKernel.compute() — double autograd.grad |
Maps to SUPA batched_hessian |
| Optimizer | Adam (1e-3) + L-BFGS fine-tune | Standard PINN 2-phase strategy |
| BC weight | 10× PDE loss | Ensures boundary enforcement |
| Resampling | Per-epoch interior points | Prevents overfitting |
| Metric | Target | Actual | Status |
|---|---|---|---|
| SpectralConv2d relative error | ≤ 1e-4 | 1.15e-07 (einsum ref) | ✅ |
| SpectralConv3d forward | — | Shape correct | ✅ |
| Gradient backprop | — | gradcheck passed (atol=1e-4) | ✅ |
| PINN relative L2 error | ≤ 1% | 0.118% | ✅ |
| Agent interaction logs | ≥ 5 logs, ≥ 3 scenarios | 12 logs, 6 scenarios | ✅ |
| BIREN GPU memory (Spectral) | 2–8 GB | ~80 MB (0.24% of 32GB) | ✅ |
| BIREN GPU memory (PINN) | < 2 GB | ~50 MB (0.15% of 32GB) | ✅ |
| Resolution | Forward | Backward | Memory |
|---|---|---|---|
| 64 × 64 | 1.04 ms | 1.39 ms | 3 MB |
| 128 × 128 | 2.26 ms | 1.95 ms | 12 MB |
| 256 × 256 | 8.37 ms | 6.39 ms | 48 MB |
from biren.compat import get_device, safe_rfft2, safe_irfft2
# Auto-detect: BIREN > CUDA > CPU
device = get_device()
# FFT with automatic SUPA fallback
x_ft = safe_rfft2(x) # try rfft2, fallback to rfft+fft
x_rec = safe_irfft2(x_ft, s) # try irfft2, fallback to ifft+irfft核心 kernel 在 BIREN 环境通过 torch_br.supa API 真实调用 SUPA 原语（不可用时自动 fallback 到 PyTorch 算子）：
.unsqueeze() * .permute()→ SUPA element-wise mul.sum(dim=2)→ SUPA reduce_sumac - bd/ad + bc→ Pure PyTorch ops (or SUPA elementwise/reduce)
详见 track5/biren/adaptation_guide.md
本项目使用 Claude Code (Opus 4.7) 辅助开发，12 段交互日志覆盖全部 6 种规定场景：
| # | Scenario | Highlights |
|---|---|---|
| 01 | 算子 kernel 设计 | einsum 原型 → torch.autograd.Function 重写 |
| 02 | 算子调试 | 3D einsum 下标冲突 + 广播维度对齐 |
| 03 | 模型架构选型 | tanh vs sin 激活对比，SIREN 选型 |
| 04 | 性能瓶颈分析 | FFT 扩展性 + CPU→GPU 迁移路径 |
| 05 | 可视化代码生成 | contourf 温度场 / 剖面 / 训练曲线 |
| 06 | BIREN GPU 适配 | API 映射表 + 显存预算 + FFT fallback |
| 07 | 自动调优 | tune.py 四维网格搜索, BIREN 实测 216 组 |
| 08 | 性能审计 | kernel-only vs 端到端修正, brsmi + official config |
| 09 | 安全审计 | 凭据泄露发现与修复, git filter-branch |
| 10 | 性能最后冲刺 | kernel='matmul' 新增, benchmark 中位数方法 |
| 11 | irfft2 深度优化 | torch.cat ~40ms → ChannelSplitIrfft2 → 28ms 超越标杆 |
| 12 | 工程收尾 | _truncated_rfft2 + 边界测试补齐 + tune.py pruned |
- GitHub: sunhao33/track5-model-operator
- Platform: Intern InkStone — track5-model-operator-dev
This project is submitted for the AI4S Competition (书生国智科探挑战赛). All rights reserved.