ruPRIM

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

Parallel primitives, reductions, scans, and indexing for Ruda.

• Cargo package: ruPRIM
• Rust crate: ruprim

Features

Quick Start

Build from the RUDA workspace:

git clone https://github.com/shuqi2077/RUDA.git
cd RUDA
cargo build --release --locked -p ruPRIM --features tensor-reduce,tensor-scan,indexing

Documentation

• User guide
• Environment setup
• Cargo features · Module exports

ruPRIM User Guide

Compute libraries · Tensor framework · 中文

ruPRIM provides device tensor reductions, cumulative scans, elementwise operations, and indexing. This page uses RudaTensor<R>, where R is a device Runtime.

1. Configure dependencies

The Cargo package is ruPRIM; its Rust import name is ruprim. Enable features for the operations you need:

This configuration places the application directory alongside the RUDA source directory. See Getting started for NVIDIA setup.

[dependencies]
ruprim = { package = "ruPRIM", path = "../RUDA/ruPRIM", default-features = false, features = ["std", "tensor-reduce", "tensor-scan", "indexing"] }
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. Sum, row reductions, and scans

This complete src/main.rs uses F32 matrix [[1, 2, 3], [4, 5, 6]] to compute its total, row means, row argmax indices, and row prefix sums. Run cargo run from the application directory:

use ruda_core::tensor::{DType, TensorMetadata, data::TensorData};
use ruda_driver_cuda::{CudaDevice, CudaRuntime};
use ruda_kernel::tensor::{readback::into_data_sync, transfer::from_data};
use ruprim::reduce::{
    components::instructions::ReduceOperationConfig,
    tensor::{KernelReduceStrategy, SumStrategy, reduce_dim, sum},
};
use ruprim::scan::cumsum;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let device = CudaDevice::default();
    let input = from_data::<CudaRuntime>(
        TensorData::new(vec![1f32, 2., 3., 4., 5., 6.], [2, 3]),
        &device,
    );
    let total = sum(
        input.clone(),
        SumStrategy::Chained(KernelReduceStrategy::Unspecified),
    )?;
    assert_eq!(total.shape().as_slice(), &[1]);
    assert_eq!(into_data_sync(total).to_vec::<f32>()?, [21.0]);

    let means = reduce_dim(
        input.clone(), None, 1,
        KernelReduceStrategy::Unspecified, ReduceOperationConfig::Mean,
    )?;
    assert_eq!(means.shape().as_slice(), &[2, 1]);
    assert_eq!(into_data_sync(means).to_vec::<f32>()?, [2.0, 5.0]);

    let indices = reduce_dim(
        input.clone(), Some(DType::I32), 1,
        KernelReduceStrategy::Unspecified, ReduceOperationConfig::ArgMax,
    )?;
    assert_eq!(into_data_sync(indices).to_vec::<i32>()?, [2, 2]);

    let prefix = cumsum(input, 1);
    assert_eq!(into_data_sync(prefix).to_vec::<f32>()?, [1., 3., 6., 4., 9., 15.]);
    Ok(())
}

sum reduces all elements and returns shape [1]. reduce_dim handles only the selected axis, normally preserving rank and setting that axis length to 1. Row means therefore have shape [2, 1], not [2].

3. Reduction operations and strategies

Full parameter order for tensor interfaces:

Select config with ReduceOperationConfig:

Index reductions require an explicit integer output dtype such as Some(DType::I32). For value reductions, pass None; output retains input dtype. Do not use output_dtype as a general cast parameter. F16/BF16 Sum, Prod, and Mean use FP32 accumulation in this path before writing the input dtype. The axis must be valid, and TopK k should be in 1..=axis length. Use reduce_dim for TopK, not whole-tensor reduce, which resets the final shape to [1].

Choose SumStrategy as follows:

• OneShot(ruda_count): explicitly sets a positive workgroup count and requires atomic add for the input dtype.
• Chained(KernelReduceStrategy::Unspecified): uses staged reductions without requiring the whole-tensor sum's global atomic add; used in the example.
• Autotune: available with tensor-reduce-autotune. Without that feature the default is OneShot(4); with it the default is Autotune.

KernelReduceStrategy offers Unspecified, Specific(ReduceStrategy), and feature-gated Autotune. Use Specific to fix a low-level strategy. Without autotuning, the default is Unspecified.

Reduction interfaces return Result<RudaTensor<R>, ReduceError>. An out-of-range axis returns InvalidAxis; missing atomic add for OneShot returns MissingAtomicAdd. sum_fallback replaces only the unsupported OneShot atomic-add case; it does not switch to CPU execution.

4. Cumulative scans

All four functions take (tensor, dim):

These are inclusive prefix operations. Output shape and dtype match the input; batch rows are processed independently. dim is zero-based and must be within the input rank. Scans return tensors directly, not Result.

The current algorithm reads the corresponding prefix for each output position, giving O(n²) total reads for an axis of length n. Account for this cost on long scan axes rather than estimating work from input element count alone.

5. Selection, gather, slicing, and integer powers

Add this function to the same file and call indexing_example(&device)?; before main returns. It reuses the earlier imports:

fn indexing_example(device: &CudaDevice) -> Result<(), Box<dyn std::error::Error>> {
    use ruda_core::tensor::element::Scalar;
    use ruprim::{
        elementwise::binary::integer_power,
        indexing::{gather, select, slice},
    };

    let input = from_data::<CudaRuntime>(
        TensorData::new(vec![1f32, 2., 3., 4., 5., 6.], [2, 3]),
        device,
    );
    let columns = from_data::<CudaRuntime>(
        TensorData::new(vec![2i32, 0], [2]), device,
    );
    let selected = select(input.clone(), 1, columns);
    assert_eq!(into_data_sync(selected).to_vec::<f32>()?, [3., 1., 6., 4.]);

    let positions = from_data::<CudaRuntime>(
        TensorData::new(vec![2i32, 0, 1, 1], [2, 2]), device,
    );
    let gathered = gather(1, input.clone(), positions);
    assert_eq!(into_data_sync(gathered).to_vec::<f32>()?, [3., 1., 5., 5.]);

    let sliced = slice(input.clone(), &[0..2, 1..3]);
    assert_eq!(into_data_sync(sliced).to_vec::<f32>()?, [2., 3., 5., 6.]);

    let squared = integer_power::scalar(input, Scalar::Int(2));
    assert_eq!(into_data_sync(squared).to_vec::<f32>()?, [1., 4., 9., 16., 25., 36.]);
    Ok(())
}

Indices and data must share a device; indices must be valid zero-based integers. For gather, use index tensors with the same rank and matching non-selected dimensions. Each example row can select different columns. Slice ranges must lie within their axes, and assignment value shape must match the selected region.

Integer powers are in ruprim::elementwise::binary::integer_power:

• scalar(input, Scalar::Int(exponent)): uses one integer exponent for all elements.
• tensor(base, exponents): uses floating-point bases and integer exponent tensors with broadcast-compatible shapes; output retains base dtype.

Use Scalar::Int(-2) for a negative exponent rather than converting the exponent to floating point. Indexing and elementwise interfaces return device tensors directly. To retain an input for other operations, pass cloned handles as in the example and keep the tensor returned by each operation.

API reference: Reductions, Scans, Indexing, Integer powers.