Lib.rs

alkomp 是一个使用 Rust 编写的 GPGPU 库，用于执行计算操作。它设计用于在 WebGPU 上工作，使计算代码能够运行在 DirectX、Vulkan、Metal 上，最终还能运行在 OpenCL 和浏览器中。

Python 绑定围绕 numpy 数组进行，下面提供了一个示例。

示例

创建你的项目：cargo new --bin gpuproject

添加到 cargo.toml

[dependencies]
alkomp = {git = "https://github.com/RustyBamboo/alkomp", branch = "main"}

修改 src/main.rs。

例如，以下代码在 GPU 上运行 Collatz 序列。

use alkomp;
fn main() {
    let code = "
        #version 450
        layout(local_size_x = 1) in;

        layout(set = 0, binding = 0) buffer PrimeIndices {
            uint[] indices;
        };

        uint collatz_iterations(uint n) {
            uint i = 0;
            while(n != 1) {
                if (mod(n, 2) == 0) {
                    n = n / 2;
                }
                else {
                    n = (3 * n) + 1;
                }
                i++;
            }
            return i;
        }

        void main() {
            uint index = gl_GlobalInvocationID.x;
            indices[index] = collatz_iterations(indices[index]);
        }";

        let mut spirv = alkomp::glslhelper::GLSLCompile::new(&code);
        let shader = spirv.compile("main").unwrap();

        let arr: Vec<u32> = vec![1, 2, 3, 4];

        let mut device = alkomp::Device::new(0);
        let data_gpu = device.to_device(arr.as_slice());

        let args = alkomp::ParamsBuilder::new()
            .param(Some(&data_gpu))
            .build(Some(0));

        let compute = device.compile("main", &shader, &args.0).unwrap();

        device.call(compute, (arr.len() as u32, 1, 1), &args.1);

        let collatz = futures::executor::block_on(device.get(&data_gpu)).unwrap();
        let collatz = &collatz[0..collatz.len() - 1];

        assert_eq!(&[0, 1, 7, 2], &collatz[..]);
}

Python 包装器（与 numpy 一起使用）

除了编写 Rust 代码外，还可以编写与 alkomp 接口的 Python 代码。目前，Python 接口设计为专门与 numpy ndarrays 一起工作。这意味着您可以使用 data_gpu = device.to_device(my_np_array) 快速将 numpy 数组发送到 GPU，并使用 device.call(...) 运行计算。 to_device 返回一个记录 GPU 缓冲区内存位置的对象，以及形状和类型。为了检索缓冲区的内容：device.get(data_gpu)。 get 函数返回一个与 my_np_array 相同形状的 numpy。

要构建 Python 库，请阅读 这篇文档。

以下是一个示例，我们将使用 Python 进行与上面相同的计算

#!python3

import alkompy
import numpy as np

arr = np.array(range(1,5), dtype=np.uint32)

# Retrieve a GPU device
dev = alkompy.Device(0)

# Send data to a GPU
data_gpu = dev.to_device(arr)

# GLSL code to compile
code = """
    #version 450
    layout(local_size_x = 1) in;
    
    layout(set = 0, binding = 0) buffer PrimeIndices {
        uint[] indices;
    };

    uint collatz_iterations(uint n) {
        uint i = 0;
        while(n != 1) {
            if (mod(n, 2) == 0) {
                n = n / 2;
            }
            else {
                n = (3 * n) + 1;
            }
            i++;
        }
        return i;
    }
    
    void main() {
        uint index = gl_GlobalInvocationID.x;
        indices[index] = collatz_iterations(indices[index]);
    }"""

shader = alkompy.compile_glsl(code)

# Run the shader and specifying the order of the bindings
dev.run("main", shader, (len(arr), 1, 1), [data_gpu])

result = dev.get(data_gpu)
assert((result == np.array([0, 1, 7, 2])).all())

待办事项

• 构建 GPU 常用操作的 crate
• Python 绑定
• 将 rust-gpu 集成到原生计算机着色器编写中

目前，在GPU上运行的计算内核代码并非原生用Rust编写。Shaderc用于将GLSL编译为SPIR-V。