用Go开发无cgo的CUDA绑定用于ML任务——第三周——开源 [P]

在我们工作中，由于公司最近切换到了Rust，我们使用CUDA与Rust。Rust有相当不错的Driver API绑定，但让我不禁疑惑，为什么我们在Go里就没有一个不需要cgo的像样方案呢。过去一个月我主要构建ML工具，而Go是我几乎所有事情的首选语言。问题在于，大多数Go的CUDA项目仍然需要cgo以及完整的工具包在构建时。这破坏了交叉编译，并且使Docker镜像变得巨大，在处理机器学习项目时很烦人。所以上个月我开始捣鼓一个概念验证，使用purego在运行时加载[libcuda.so](http://libcuda.so)。完全不使用cgo。最大的痛点是线程亲和性。CUDA按线程保存上下文，所以goroutine切换会导致问题。我构建了一个简单的执行器，通过runtime.LockOSThread锁定一个OS线程，并通过一个通道传递所有调用。目前使用起来大概是这样： func run() error { cuda.Init() dev, _ := cuda.GetDevice(0) ctx, _ := dev.Primary() defer ctx.Close() a, _ := cuda.Alloc[float32](ctx, 1024) b, _ := cuda.Alloc[float32](ctx, 1024) c, _ := cuda.Alloc[float32](ctx, 1024) stream, _ := ctx.NewStream() start, _ := ctx.NewEvent() stop, _ := ctx.NewEvent() start.Record(stream) fn.LaunchOn(bg, stream, cfg, cuda.Arg(a), cuda.Arg(b), cuda.Arg(c), cuda.ArgValue(int32(1024)), ) stop.Record(stream) stop.Synchronize() duration, _ := start.Elapsed(stop) fmt.Printf("GPU time: %v\n", duration) return nil } 在我的4070 Ti上，一个1000万元素的向量加法显示CPU计时大约160微秒，但实际GPU事件计时为434微秒。这个差异让我惊讶。这个项目仍然非常早期，进展缓慢，因为我只在周末编码，而且我对CUDA完全是个新手。正在慢慢添加Graph和多GPU支持。这非常早期，所以请把它当作一个学习CUDA的仓库，但我学习CUDA觉得很有趣。我想你们中有些人可能也会觉得它有意思。仓库在[github.com/eitamring/gocudrv](http://github.com/eitamring/gocudrv)，如果你想看看。如果任何拥有5xxx系列显卡的人想试试，那就太好了（眨眼）