一个完全在 GPU 上运行的 CPU
Hacker News 摘要原标题:A CPU that runs entirely on GPU
这个项目是一个在 GPU 上完全运行的 CPU,其寄存器、内存、标志位和程序计数器全部由张量(Tensor)表示。它在算术逻辑单元(ALU)的设计上非常独特:所有的算术和逻辑操作都是通过训练好的神经网络模型来完成的,没有任何硬编码的算术逻辑。
核心设计与工作原理
整个 CPU 架构常驻在 GPU 中,利用 PyTorch 张量来管理状态。指令解码、ALU 调度和状态更新都在 GPU 设备端发生,数据不需要返回主机 CPU。
该系统的 ALU 操作通过一系列经过训练的 .pt 模型进行路由:
• 加法与减法:采用神经网络实现的 Kogge-Stone 进位先行算法(CLA),经过 8 次神经传递完成,保证了 100% 的准确率。
• 乘法:使用学习到的字节对查找表(Look-Up Table),即使是 64 位乘法也能高效处理。
• 位运算:与、或、异或等逻辑运算使用向量化的真值表神经网络。
• 移位操作:左右移位采用基于注意力机制的位路由技术。
• 除法:通过神经网络减法实现的恢复余数除法算法。
• 高级数学函数:正弦、余弦、平方根、指数、对数以及反正切函数同样通过专门训练的深度神经网络或残差网络模型实现。
训练模型库
项目内置了 23 个模型,总容量约为 135 MB,其中 13 个模型被激活并用于指令执行。这些模型已经通过了 347 项自动化测试,确保在整数算术上达到 100% 的准确度。此外,项目甚至包含了一个基于 Qwen2.5-Coder-1.5B LoRA 的大语言模型适配器用于指令解码。
性能与研究发现
在苹果芯片(Apple Silicon)的 MPS 后端基准测试中,该神经 CPU 表现出一些与传统 CPU 截然不同的特性:
1. 乘法速度超过加法:在神经 CPU 中,由于字节对查找表具有零顺序依赖性,其乘法速度比加法快 12 倍。而在传统硬件中,乘法通常比加法慢。
2. 经典的并行算法有效:将硬件设计中的 Kogge-Stone 算法应用于神经网络后,加减法速度提升了 3.3 倍。
3. 向量化优势:通过向量化处理,移位操作的速度提升了 6.5 倍,有效弥补了注意力机制的计算开销。
在执行模式上,程序可以以每周期 136 到 262 微秒的速度运行,大约相当于每秒执行 4975 条指令(IPS)。
硬件模拟规格
该系统模拟的是 64 位 ARM64 架构:
• 寄存器:31 个 64 位寄存器,在 GPU 上以 torch.int64 存储。
• 内存:扁平化的字节寻址空间,以 torch.uint8 存储。
• 标志位:N、Z、C、V 状态标志均为 GPU 常驻张量。
• 指令集:支持完整的 ARM64 指令解码,并能执行真实的二进制编码。
执行模式与技术实现
项目提供了两种主要的执行模式:
• 神经模式:默认模式。每一次 ALU 算术运算都是对神经网络模型的一次前向推理。
• 快速模式:保持 GPU 张量架构,但 ALU 运算改用 GPU 原生算术指令(如 torch.add),在苹果芯片上批量处理时速度可达 1.35M IPS。
为了追求极致性能,项目还包含了基于 Metal 的计算内核:
• MLX Metal:利用苹果 MLX 框架编写的自定义 Metal 计算内核。
• Rust Metal:使用 Rust 语言直接通过原生 Metal API 实现的 GPU 指令循环,支持 GPU 端的系统调用处理、基本块缓存、神经调度和乱序执行。
项目结构与演示
项目包含了完整的开发组件:
• ncpu:核心逻辑,分为神经模式、模型模式和张量内核。
• kernels:底层 GPU 内核代码,包含 Metal 和 Rust 实现。
• models:训练好的算术、逻辑、位移和数学模型。
• demos:包含一个著名的 DOOM 射线投射器演示,该程序所有的算术运算(如 DDA 算法中的每一处计算)全部通过神经网络模型完成。在神经模式下约 2.5 FPS,而在快速模式下可达 5000 FPS。
• paper:包含详细的架构分析研究论文。
用户可以通过简单的 pip install 完成环境安装,并使用 main.py 运行汇编程序或二进制固件。该项目通过 MIT 协议开源。