Skip to content

Basalt 3.13.0 — GEMM 重写 + 设备层

Latest
Latest

Choose a tag to compare

View all tags
@github-actions github-actions released this 01 Nov 04:01
· 1 commit to main since this release
v3.13.0
0de0623

Basalt 3.13.0

「CPU 到瓶颈了」这个判断，实测不支持。真正的问题在 GEMM 实现，不在硬件。

  • 坐标 — io.github.9178nekomc:basalt:3.13.0
  • 在线文档 — https://9178nekomc.github.io/basalt/
  • 规模 — 244 个公开类型 / 1632 个公开成员，全部带逐参数说明

先说 CUDA

这台机器没有 NVIDIA 显卡（只有 AMD 核显 + 两个虚拟显示器），也没有 CUDA 工具链。
CUDA 代码在这里无法编译、无法运行、无法验证——交付一堆从没执行过的内核，比不写更糟。

所以这一版做的是设备层：接 GPU 的前置条件，而且完全可以在 CPU 上验证。
CUDA 后端留到有 N 卡的机器上再实现，届时只需填内核，不用改调用方。

GEMM 重写

原实现走 i → k → j，每做一次乘加就要读一整行 B，跑到的是内存带宽而不是算力。
512³ 只有 31.8 GFLOPS，约为此 CPU 峰值的 15%。

规模 改前 改后 倍数
512³ 31.8 GFLOPS 52.0 GFLOPS 1.63×
1024³ 32.7 GFLOPS 49.5 GFLOPS 1.51×

四处改动：

  • 6×16 微内核：12 个累加器常驻寄存器；每步 k 加载 6 个 A 标量与 2 个 B 向量，
    发出 12 次 FMA——每个 FMA 摊到 0.67 次加载（8×8 内核是 1.1 次）
  • B 面板打包：按 KC × NC 收成连续内存，微内核顺序读
  • 三层分块 KC=512 / MC=30 / NC=256：工作集压进 L2
  • 运行时指令集分派：is_x86_feature_detected! 检测 AVX2 + FMA，缺失则退回标量内核

分块参数是扫出来的：KC/NC 从 128 到 512、MC 从 24 到 120 各跑一遍取最优。

扫描过程中发现一个真 bug：MC 必须是 MR 的整数倍，否则最后一个寄存器块会写穿到下一块内存。
现象是段错误，而且只在线程池路径上出现（单线程测试全过）。现在有 debug_assert 兜着。

设备层

Tensor x = engine.tensor(data, 64, 64);
x.device();                          // cpu
Tensor copy = x.to(Device.CPU());    // 复制一份新存储
x.to(Device.cuda(0));                // UnsupportedOperationException: only cpu is available
  • 张量携带设备标签，跨设备运算拒绝执行，不做隐式搬运
  • Device.cuda(0) 可以创建、比较、打印，但用它计算会显式失败，不假装成功
  • 同设备传输同样走复制路径——这条通道现在就被 5 个测试覆盖，
    而不是等第一次上 GPU 才第一次执行

两个自己挖的坑

1. 运行时 cfg! 当编译期门控

#[target_feature(enable = "avx2,fma")] 只在 x86 有效，但我用 cfg!(target_arch = "x86_64")
做分派——那是运行时的布尔值，代码在 ARM 上仍然要编译。macOS aarch64 构建直接失败。
本地是 x86_64，怎么编都过，是 CI 抓到的。

2. JNI 属性错位

插入设备层入口时，#[no_mangle] 落到了 transferTensor 上而不是 tensorShape，
导致 tensorShape 的导出符号消失。这种错误不会在编译期报，要等第一次调用才炸。
已在 DLL 里确认两个符号都存在。

验证

  • 38 个 Rust 测试（新增 3 个 GEMM 用例：对齐尺寸、含余数尺寸、多面板）
  • 20 个 Java 测试（新增 5 个设备层用例）
  • ./verify.sh 全绿（9 组套件）
  • 文档：19 页 / 244 类型 / 1632 个成员全部有说明 / 0 占位符 / 0 断链

还没做的

  • CUDA 后端：需要 N 卡机器
  • 多线程 GEMM 负载均衡：现在按行切块，大矩阵下最后一个任务偏重
  • 掩码：where / maskedFill

产物

文件 说明
basalt-3.13.0.jar 标准 jar，含四平台原生库
basalt-3.13.0-all.jar Fat jar，无外部依赖
basalt-3.13.0-sources.jar 源码
basalt-docs-3.13.0.zip 离线文档，19 页
Assets 6
Loading