Kyanite

• Kyanite
  • 概述
  • 快速演示
  • 系统要求
  • 内部结构
    • 图中间表示法
    • 优化器
    • CPU 执行器
    • Cuda 执行器
  • 与其他 crates 的比较
    • Rust 封装现有的运行时
    • 从头开始的 Rust 项目
  • 开发

概述

Kyanite 是一个用 Rust 编写/为 Rust 编写的神经网络推理库。它可以使用 cuda/cudnn/cublas 在 CPU 或 Nvidia GPU 上运行 ONNX 文件。

它足够通用，可以运行各种类型的网络，已经与以下进行了测试

• 简单的全连接网络
• 基于 ResNet 的 CNN
• 大型语言模型，如 LLaMA
• 图像生成模型，如 Stable Diffusion。有关演示，请参阅 stable_diffusion 示例，位于 kn-runtime crate。

该框架由以下 crates 组成

• kn-graph: 核心crate，包含中间表示和 CPU 执行器。
• kn-cuda-sys: 使用 rust-bindgen 生成的 Cuda FFI 绑定。
• kn-cuda-eval: Cuda 执行器和规划器。
• kn-runtime: 用于在运行时选择 CPU 和 GPU 执行的 crate 封装器。
• kn-python: 使用 PyO3 的实验性 Python 封装器，用于运行时 crate。

快速演示

// Graph operations (using kn-graph)
// Load on onnx file into a graph
let graph = load_graph_from_onnx_path("test.onnx", false)?;
// Optimize the graph
let graph = optimize_graph(&graph, Default::default());
// Render the graph as an svg file
graph_to_svg("test.svg", &graph, false, false)?;

// Build the inputs
let batch_size = 8;
let inputs = [DTensor::F32(Tensor::zeros(IxDyn(&[batch_size, 16])))];

// CPU: (using kn-graph)
// just evaluate the graph
let outputs: Vec<DTensor> = cpu_eval_graph(&graph, batch_size, &inputs);

// GPU: (using kn-cuda-eval)
// build an executor
let device = CudaDevice::new(0).unwrap();
let mut executor = CudaExecutor::new(device, &graph, batch_size);
// run the executor on the inputs
let outputs: &[DTensor] = executor.evaluate(&inputs);

// Runtime device selection: (using kn-runtime)
let device = Device::best();
let mut prepared = device.prepare(graph, batch_size);
let outputs: Vec<DTensor> = prepared.eval( & inputs);

系统要求

要使用 CUDA crates，需要在系统上安装适当的库；它们不会自动下载

• CUDA（包括CUDA、cuBLAS、NVRTC）: 安装程序，按照说明进行操作。确保环境变量 CUDA_PATH 指向安装的根目录（即，CUDA_PATH/bin/ 应存在）。
• cuDNN: 存档文件，提取到您选择的路径。如果您选择与 CUDA_PATH 相同的路径，则无需执行其他操作。否则，将环境变量 CUDNN_PATH 设置为 cuDNN 安装的根目录（即，CUDNN_PATH/bin 应存在）。

该项目已在 CUDA v12.2 和 cuDNN 版本 v8.9.5 上进行过测试。较新版本可能也可以工作，但这不能保证，因为 CUDA 有时会更改某些函数的名称或删除它们。

内部结构

典型的流程图如下所示。下面的第二张图显示了在简单的神经网络架构上运行此流程图的结果。

图中间表示法

中心是 Graph IR，神经网络图的中间表示。

该结构是一个 SSA 风格的有向无环图，其中节点是具有形状、数据类型和计算它的操作的值。这些值是抽象的；它们还没有步长或内存位置。

操作与其他框架类似，但尽可能保持正交。一些示例操作：卷积、矩阵乘法、重塑、广播、切片、一元、二元、归约、softmax 等。有关图操作的全列表，请参阅 文档。

可以使用 图构建器 API 在代码中直接构建图，但为了方便，存在一个 ONNX 加载器。它可以读取 ONNX 文件并将支持的子集操作转换为 IR 所支持的。

由于图 IR 比 ONNX 规范更为正交，因此许多 ONNX 操作被分解成单独的步骤，以下是一些示例

• ONNX 二元操作隐式广播其操作数，但在 IR 中这是一项单独的操作。
• ONNX 卷积和矩阵乘法有一个内置的可选偏置操作数；这也变成了一个单独的广播加二元加法操作。

要确定 ONNX 操作是否受支持，请检查 load.rs 中 visit_node 函数顶层匹配语句的分支。许多常见操作已经实现，添加更多操作不应太难。

有关典型图的更大示例，请参阅 stable_diffusion_piece.svg，这是从稳定扩散模型的开头取出的一个小部分。

优化器

图可以由 优化器 优化（可选）。由于图是追加的，因此返回一个新的图。

目前实现的优化有

• 常量折叠
• 将连续的仿射（偏置、缩放、批归一化）操作融合成一个单一的偏置+缩放操作。
• 将连续的裁剪操作（relu、min、max）融合成一个单一的 min+max 操作。
• 强度降低：用乘以逆常数替换除以常数。
• 识别 layernorm 模板（reduce、subtract、power、reduce、divide）并用 layernorm 操作符替换。

CPU 执行器

最后，需要执行这个图。有一个简单的CPU执行器，它直接运行每个操作。在这里没有尝试进行任何主要的优化，除了使用BLAS例程进行矩阵乘法和im2col进行卷积。这个执行器尽可能简单是很重要的，因为它作为检查GPU执行器正确性的单元测试的基础。

Cuda 执行器

运行这些图的第二种（更实用）方式是使用Cuda执行器。这涉及到通过Cuda规划器运行图，该规划器输出预定的Cuda操作调度并分配必要的内存缓冲区。这个步骤被分成一个单独的步骤，这样昂贵的规划步骤只需要在每个网络架构上执行一次；生成的计划然后可以在执行器中多次重用。

规划器有以下主要职责

• 确定张量的内存布局：步长和内存偏移量
  • 这隐式地处理了大多数reshape、广播、步长...操作。
  • 如果可能，也会重用缓冲区，以最小化总内存使用。这里有很大的改进空间；目前，这只是一个单遍算法。

• 决定卷积和矩阵乘法运行哪些cuDNN/cuBLAS操作。如果可能，将这些操作融合在一起。以下是一些示例
  • cuDNN支持一个“卷积 + 剩余 + 偏置 + relu”操作
  • cuBLAS矩阵乘法可以包括输入矩阵的转置，以及等效地通过交换输入输出的输出矩阵的转置。
  • cuDNN和cuBLAS操作有时包括一个“标量”参数，它乘以一些操作数

• 使用基于NVRTC（运行时编译）的autokernel框架编译剩余的标量和复合操作的自定义内核。
  • 由autokernel处理的操作包括：标量操作、减少、softmax、layernorm、gather。
  • 使用手工编写的内核模板，并在编译前在运行时替换细节，如张量形状、步长、标量操作...
  • 这里发生了更多的操作融合
    • 多个标量操作被编译成一个单独的内核
    • 常量标量被内联
    • 一些复合内核支持融合输入或输出标量操作

这种最终的运算融合可以非常显著并节省大量的主存冗余传输。可以通过手动为每个使用的操作组合编写内核来实现相同的表现，但组合爆炸和相关的维护会非常庞大。

以下是一个带有一些手写澄清注释的生成标量内核的示例

#include "util.cu"

// constants that got inserted into the template
// this scalar operation happens on a tensor of rank 4, with 7 operands
const int RANK = 4;
const int OPERANDS = 7;
const int STRIDES_DENSE[RANK] = {648, 81, 9, 1};
const int STRIDES[OPERANDS][RANK] = {
    // these are full input tensors with normal, dense strides
    {648, 81, 9, 1},
    {648, 81, 9, 1},
    // these values have zero strides for all axes except the channel one,
    //    so these are probably biases and scaling factors
    //    that are broadcast across the other axes
    {0, 1, 0, 0},
    {0, 1, 0, 0},
    {0, 1, 0, 0},
    {0, 1, 0, 0},
    // the output tensor is just another operand
    {648, 81, 9, 1}
};

// the template function, the body of which is generated at runtime
__device__ void operation(void *pointers[OPERANDS], int offsets[OPERANDS]) {
    // all input operand memory locations are cast to the right type
    float *x0 = &((float *) pointers[0])[offsets[0]];
    float *x1 = &((float *) pointers[1])[offsets[1]];
    float *x2 = &((float *) pointers[2])[offsets[2]];
    float *x3 = &((float *) pointers[3])[offsets[3]];
    float *x4 = &((float *) pointers[4])[offsets[4]];
    float *x5 = &((float *) pointers[5])[offsets[5]];
    float *x6 = &((float *) pointers[6])[offsets[6]];
    
    // input operands are loaded
    float y0 = *x0;
    float y1 = *x1;
    
    // this is probably a residual connection
    float y2 = y0 + y1;
    
    // these 4 steps look like they're implementing a batchnorm layer  
    float y3 = *x2;
    float y4 = y2 - y3;
    float y5 = *x3;
    float y6 = y4 / y5;
    float y7 = *x4;
    float y8 = y6 * y7;
    float y9 = *x5;
    float y10 = y8 + y9;
    
    // this implements a relu6 activation function
    float y11 = 6;
    float y12 = min(y10, y11);
    float y13 = (0.0);
    float y14 = max(y12, y13);
    
    // finally the output is stored
    *x6 = y14;
}

// the kernel main function is the same for all scalar kernels
__global__ void scalar_kernel(
        int batch_size,
        Array<void *, OPERANDS> pointers
) {
    KernelInfo info = kernel_info();
    int size = batch_size * STRIDES_DENSE[0];

    // the main loop, following https://developer.nvidia.com/blog/cuda-pro-tip-write-flexible-kernels-grid-stride-loops/
    for (int flat = info.global_thread_id; flat < size; flat += info.thread_count) {
        Array<int, OPERANDS> offsets = flat_index_to_offsets<RANK, OPERANDS>(flat, STRIDES_DENSE, STRIDES);
        operation(pointers.data, &offsets[0]);
    }
}

与其他 crates 的比较

有关潜在替代方案的完整列表，请参阅Are We Learning Yet?

Rust 封装现有的运行时

• PyTorch包装器：tch
• TensorFlow包装器：tensorflow
• ONNXRuntime包装器：ort

优点

• 广泛支持许多神经网络操作
• 支持许多不同的后端（CPU、GPU（Nvidia + AMD）、TPU...）

缺点

• 对于加载ONNX文件的支持并不总是很好（但ort在这方面做得很好，正如其名称所示）
• 大型且有些黑盒的外部依赖
• 在许多情况下，运算融合较少，尽管预计未来会有所改善

性能应该与Kyanite在运算融合不太重要的情况下大致相同；所有库基本上都使用相同的底层cuDNN和cuBLAS内核。

从头开始的 Rust 项目

• tract：覆盖了ONNX规范更大的范围，但只支持CPU推理

开发

在开发这个crate时，为了更新ONNX协议，使用了prost-build crate。这需要安装protoc，并将PROTOC环境变量设置为指向可执行文件。有关详细信息，请参阅他们的安装说明（或构建脚本显示的任何错误消息）。

实际上更新协议定义时，将kn-graph/proto/onnx.proto3替换为较新版本，并运行cargo run --bin proto-to-rust。然后提交onnx.proto3文件和生成的onnx.rs文件。