当 Go 遇上 GPU：用 CUDA 释放千倍算力的实战指南

go-cuda-cgo-demo/
├── main.go       # Go 主程序 (CGO 入口，负责内存分配和调度)
├── matrix.cu     # CUDA 内核代码 (在 GPU 上运行的 C++ 代码)
└── matrix.h      # C 头文件 (声明导出函数，供 CGO 识别)

// matrix.cu
#include <cuda_runtime.h>
#include <stdio.h>

// CUDA Kernel: 每个线程计算 C[row][col] 的值
__global__ void matrixMulKernel(float *a, float *b, float *c, int width) {
    // 根据 Block ID 和 Thread ID 计算当前线程的全局坐标
    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < width && col < width) {
        float sum = 0;
        // 计算点积
        for (int k = 0; k < width; k++) {
            sum += a[row * width + k] * b[k * width + col];
        }
        c[row * width + col] = sum;
    }
}

extern "C" {
    // 供 Go 调用的 C 包装函数
    // 负责显存分配、数据拷贝和内核启动
    void runMatrixMul(float *h_a, float *h_b, float *h_c, int width) {
        int size = width * width * sizeof(float);
        float *d_a, *d_b, *d_c;

        // 1. 分配 GPU 显存 (Device Memory)
        cudaMalloc((void **)&d_a, size);
        cudaMalloc((void **)&d_b, size);
        cudaMalloc((void **)&d_c, size);

        // 2. 将数据从 Host (CPU内存) 复制到 Device (GPU显存)
        // 这一步通常是性能瓶颈，应尽量减少
        cudaMemcpy(d_a, h_a, size, cudaMemcpyHostToDevice);
        cudaMemcpy(d_b, h_b, size, cudaMemcpyHostToDevice);

        // 3. 定义 Grid 和 Block 维度
        // 每个 Block 包含 16x16 = 256 个线程
        dim3 threadsPerBlock(16, 16);
        // Grid 包含足够多的 Block 以覆盖整个矩阵
        dim3 numBlocks((width + threadsPerBlock.x - 1) / threadsPerBlock.x,
                       (width + threadsPerBlock.y - 1) / threadsPerBlock.y);

        // 4. 启动内核！成千上万个线程开始并行计算
        matrixMulKernel<<<numBlocks, threadsPerBlock>>>(d_a, d_b, d_c, width);

        // 5. 将计算结果从 Device 传回 Host
        cudaMemcpy(h_c, d_c, size, cudaMemcpyDeviceToHost);

        // 6. 释放 GPU 内存
        cudaFree(d_a);
        cudaFree(d_b);
        cudaFree(d_c);
    }
}

// matrix.h
#ifndef MATRIX_H
#define MATRIX_H

void runMatrixMul(float *a, float *b, float *c, int width);

#endif

// go-cuda-cgo-demo/main.go
package main

/*
#cgo LDFLAGS: -L. -lmatrix -L/usr/local/cuda/lib64 -lcudart
#include "matrix.h"
*/
import "C"
import (
    "fmt"
    "math/rand"
    "time"
    "unsafe"
)

const width = 1024 // 矩阵大小 1024x1024，共 100万次计算

func main() {
    size := width * width
    h_a := make([]float32, size)
    h_b := make([]float32, size)
    h_c := make([]float32, size)

    // 初始化矩阵数据
    rand.Seed(time.Now().UnixNano())
    for i := 0; i < size; i++ {
        h_a[i] = rand.Float32()
        h_b[i] = rand.Float32()
    }

    fmt.Printf("Starting Matrix Multiplication (%dx%d) on GPU...\n", width, width)
    start := time.Now()

    // 调用 CUDA 函数
    // 使用 unsafe.Pointer 获取切片的底层数组指针，传递给 C
    C.runMatrixMul(
        (*C.float)(unsafe.Pointer(&h_a[0])),
        (*C.float)(unsafe.Pointer(&h_b[0])),
        (*C.float)(unsafe.Pointer(&h_c[0])),
        C.int(width),
    )

    // 注意：在更复杂的场景中，需要使用 runtime.KeepAlive(h_a)
    // 来确保 Go GC 不会在 CGO 调用期间回收切片内存。

    elapsed := time.Since(start)
    fmt.Printf("Done. Time elapsed: %v\n", elapsed)

    // 简单验证：检查左上角元素
    fmt.Printf("Result[0][0] = %f\n", h_c[0])
}

nvcc -c matrix.cu -o matrix.o
ar rcs libmatrix.a matrix.o

# 链接本地的 libmatrix.a 和系统的 CUDA 运行时库
go build -o gpu-cgo-demo main.go

./gpu-cgo-demo

Starting Matrix Multiplication (1024x1024) on GPU...
Done. Time elapsed: 611.815451ms
Result[0][0] = 262.440918

# 编译为共享库 libmatrix.so
nvcc -shared -Xcompiler -fPIC matrix.cu -o libmatrix.so

// go-cuda-purego-demo/main.go
package main

import (
    "fmt"
    "math/rand"
    "runtime"
    "time"

    "github.com/ebitengine/purego"
)

const width = 1024

func main() {
    // 1. 加载动态库
    // 注意：在运行时，libmatrix.so 和 libcuder.so 必须在 LD_LIBRARY_PATH 中
    libMatrix, err := purego.Dlopen("libmatrix.so", purego.RTLD_NOW|purego.RTLD_GLOBAL)
    if err != nil {
        panic(err)
    }

    // 还需要加载 CUDA 运行时库，因为 libmatrix 依赖它
    _, err = purego.Dlopen("/usr/local/cuda/lib64/libcudart.so", purego.RTLD_NOW|purego.RTLD_GLOBAL)
    if err != nil {
        panic(err)
    }

    // 2. 注册 C 函数符号
    var runMatrixMul func(a, b, c *float32, w int)
    purego.RegisterLibFunc(&runMatrixMul, libMatrix, "runMatrixMul")

    // 3. 准备数据 (与 CGO 版本相同)
    size := width * width
    h_a := make([]float32, size)
    h_b := make([]float32, size)
    h_c := make([]float32, size)

    rand.Seed(time.Now().UnixNano())
    for i := 0; i < size; i++ {
        h_a[i] = rand.Float32()
        h_b[i] = rand.Float32()
    }

    fmt.Println("Starting Matrix Multiplication via PureGo...")
    start := time.Now()

    // 4. 直接调用！无需 CGO 类型转换
    runMatrixMul(&h_a[0], &h_b[0], &h_c[0], width)

    // 5. 极其重要：保持内存存活
    // PureGo 调用是纯汇编实现，Go GC 无法感知堆栈上的指针引用
    // 必须显式保活，否则在计算期间 h_a 等可能被 GC 回收！
    runtime.KeepAlive(h_a)
    runtime.KeepAlive(h_b)
    runtime.KeepAlive(h_c)

    fmt.Printf("Done. Time: %v\n", time.Since(start))
    fmt.Printf("Result[0][0] = %f\n", h_c[0])
}

# 无需 CGO，直接在go-cuda-purego-demo下运行
# 确保当前目录在 LD_LIBRARY_PATH 中
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:.
CGO_ENABLED=0 go run main.go
Starting Matrix Multiplication via PureGo...
Done. Time: 584.397195ms
Result[0][0] = 260.088806