RUDA.
Документация / РусскийИсходный текст ↗

ruTENSOR Руководство пользователя

English | 简体中文 | 日本語 | Deutsch | Русский

Вычислительные библиотеки · Тензорная платформа · 中文

ruTENSOR обеспечивает сжатие тензоров по именованным осям, сокращения, физические перестановки и поэлементные операции. Входы используют RudaTensor<R>; приложение выбирает устройство Runtime. Эта библиотека отличается от платформы более высокого уровня ruda-tensor.

1. Настройте зависимости

Пакет Cargo — ruTENSOR; имя импорта Rust — rutensor. По умолчанию включен std и вычисление тензора устройства без выбора драйвера. В следующей конфигурации каталог приложения размещается рядом с исходным каталогом RUDA:

[dependencies]
rutensor = { package = "ruTENSOR", path = "../RUDA/ruTENSOR" }
ruda-core = { path = "../RUDA/ruda-core", default-features = false, features = ["std", "tensor-host-data"] }
ruda-kernel = { path = "../RUDA/ruda-kernel", default-features = false, features = ["frontend-std", "device-tensor"] }
ruda-driver-cuda = { path = "../RUDA/ruda-driver-cuda", default-features = false, features = ["std"] }

2. Свёртка тензоров, редукция и перестановка

Сохраните следующее как src/main.rs приложения:

use ruda_core::tensor::data::TensorData;
use ruda_driver_cuda::{CudaDevice, CudaRuntime};
use ruda_kernel::tensor::{readback::into_data_sync, transfer::from_data};
use rutensor::{einsum, permute, reduce, ReductionOp};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = CudaDevice::default();
    let a = from_data::<CudaRuntime>(
        TensorData::new(vec![1f32, 2., 3., 4., 5., 6.], [2, 3]), &device,
    );
    let b = from_data::<CudaRuntime>(
        TensorData::new(vec![1f32, 0., 0., 1., 1., 1.], [3, 2]), &device,
    );
    let product = einsum("ik,kj->ij", &[&a, &b])?;
    let row_sums = reduce(&a, &[0, 1], &[0], ReductionOp::Sum)?;
    let transposed = permute(&a, &[1, 0])?;

    println!("product: {:?}", into_data_sync(product).to_vec::<f32>()?);
    println!("row sums: {:?}", into_data_sync(row_sums).to_vec::<f32>()?);
    println!("transpose: {:?}", into_data_sync(transposed).to_vec::<f32>()?);
    Ok(())
}

einsum("ik,kj->ij", ...) суммирует по k и возвращает форму [2, 2]. reduce сохраняет режим 0 и уменьшает режим 1, возвращая форму [2]. permute возвращает недавно выделенный тензор [3, 2], а не представление, совместно использующее входное хранилище.

3. Выражения Einsum

einsum(expression, inputs) принимает один или несколько входов. Буквы с учетом регистра обозначают оси; правая часть стрелки выбирает и упорядочивает выходные оси.

Выражение Операция
ik,kj->ij Умножение матрицы
...ik,...kj->...ij Матричное умножение с широковещательными пакетными осями
abc,cde->abde Многомерная свёртка тензоров
ij,jk,kl->il Свёртка трёх входов
i,j->ij Внешнее произведение
ii->i Диагональ
ii-> Трассировка, возвращающая скаляр нулевого ранга
ijk->ki Редуцировать j и переставить оставшиеся оси
...i->i Редуцировать все оси, обозначенные многоточием

Для фиксированных выражений, фигур, шагов и dtypes создайте EinsumPlan::new(expression, descriptors) и повторно используйте execute(inputs). Чтобы указать точность вывода и арифметику, используйте EinsumPlan::with_options или einsum_with_options.

4. Дескрипторы и планы выполнения

A Mode — это метка i32. Одна и та же метка идентифицирует одну и ту же логическую ось в тензорах, независимо от ее физического положения. TensorDescriptor хранит экстенты, шаги элементов и хранилище dtype. OperandDescriptor добавляет метки и входное унарное преобразование.

Эти функции создают и выполняют план для D = alpha * A @ B + beta * C:

use ruda_kernel::{dsl::Runtime, tensor::RudaTensor};
use rutensor::{
    ComputeType, DType, OperandDescriptor, OperationDescriptor,
    Plan, Result, TensorDescriptor,
};

fn make_plan<R: Runtime>(
    a: &RudaTensor<R>, b: &RudaTensor<R>, c: &RudaTensor<R>,
    m: usize, n: usize,
) -> Result<Plan> {
    let operation = OperationDescriptor::contraction(
        OperandDescriptor::from_tensor(a, &[0, 2])?,
        OperandDescriptor::from_tensor(b, &[2, 1])?,
        Some(OperandDescriptor::from_tensor(c, &[0, 1])?),
        TensorDescriptor::contiguous(&[m, n], DType::F32)?,
        &[0, 1],
        ComputeType::F32,
    )?;
    Plan::new(operation)
}

fn execute<R: Runtime>(
    plan: &Plan, a: &RudaTensor<R>, b: &RudaTensor<R>, c: &RudaTensor<R>,
    alpha: f64, beta: f64,
) -> Result<RudaTensor<R>> {
    plan.execute(&[a, b, c], &[alpha, beta])
}

Планы не сохраняют входные буферы. Последующие исполнения могут использовать разные тензоры с соответствующими формами, шагами и типами d. Все входы должны находиться на одном устройстве.

Конструктор Порядок ввода выполнения Скалярный порядок
contraction A, B; C необязателен alpha; beta при наличии C
sum_product Все входы произведения; C необязателен alpha; beta при наличии C
reduction A; C необязателен alpha; beta при наличии C
permutation A alpha
elementwise_binary A, B alpha, beta
elementwise_trinary A, B, C alpha, beta, gamma

Plan::execute выделяет выход. Plan::execute_into принимает и возвращает выходной тензор, форма, шаги и dtype которого соответствуют выходному дескриптору. Его буфер должен принадлежать единолично, без совместного использования с входами или другими представлениями.

Используйте TensorDescriptor::new(extents, strides, dtype) для дополненных или измененных макетов вывода; выходные оси не должны перекрываться. Явные дескрипторы операций могут добавлять дополнительные оси широковещания через выходную форму.

5. Сокращение и поэлементные операции

reduce(input, input_modes, output_modes, operation) уменьшает количество меток, отсутствующих в выводе; уменьшенные оси удалены:

ReductionOp Операция Пустая редукция
Sum Сумма 0
Product Произведение 1
Min Минимум Положительная бесконечность
Max Максимум Отрицательная бесконечность

elementwise_binary и elementwise_trinary выравнивают и транслируют именованные оси с помощью BinaryOp::{Add, Mul, Min, Max}. Троичные операции оценивают (alpha * op(A) op_ab beta * op(B)) op_abc gamma * op(C).

Выберите Identity, Negate, Abs, Sqrt, Exp, Log, Sin, Cos, Tanh, Relu, Reciprocal или от Conjugate до OperandDescriptor::with_unary. Log — натуральный логарифм; Conjugate равен Identity для реальных значений. Мин и Макс распространяются на NaNs.

6. Точность, хранение и ошибки

ruTENSOR предоставляет Ruda Rust API, а не NVIDIA cuTENSOR C ABI. Устройство должно поддерживать выбранные типы хранилища и вычислений.