Struct GpuContext 

pub struct GpuContext { /* private fields */ }

GPU context for compute operations.

Manages GPU device, queue, and cached compute pipelines for FST operations. Create once and reuse for multiple operations to amortize initialization cost.

Initialization

GPU context creation is an async operation that:

1. Enumerates available GPU adapters
2. Requests a device with required features
3. Compiles and caches compute shader pipelines

Use pollster::block_on() for synchronous initialization:

use arcweight::gpu::GpuContext;

let ctx = pollster::block_on(GpuContext::new())?;
println!("Using GPU: {}", ctx.adapter_info().name);

Thread Safety

GpuContext is Send but not Sync. Each thread should have its own context, or use synchronization primitives to share access.

Resource Management

GPU resources are automatically cleaned up when the context is dropped. Buffers created from this context remain valid as long as the context lives.

Implementations

impl GpuContext

pub async fn new() -> Result<Self>

Create a new GPU context.

Initializes the GPU backend by requesting a high-performance adapter, creating a device, and compiling compute shader pipelines.

Async

This is an async operation. Use pollster::block_on() for synchronous initialization in non-async contexts.

Errors

Returns an error if:

• No compatible GPU adapter is found
• Device creation fails
• Shader compilation fails

Examples

use arcweight::gpu::GpuContext;

// Async context
async fn init_gpu() -> arcweight::Result<GpuContext> {
    GpuContext::new().await
}

// Synchronous context
let ctx = pollster::block_on(GpuContext::new())?;

pub fn adapter_info(&self) -> &AdapterInfo

Get information about the GPU adapter.

Returns details about the GPU including name, vendor, and backend.

pub fn device(&self) -> &Device

Get the underlying wgpu device.

Use this for advanced GPU operations not provided by this module.

pub fn queue(&self) -> &Queue

Get the wgpu command queue.

Use this for submitting custom command buffers.

pub fn create_weight_buffer(&self, weights: &[f32]) -> Result<GpuWeightBuffer>

Create a weight buffer on the GPU.

Uploads weight data to GPU memory for use in compute operations. The buffer is created with storage and copy usage flags.

Arguments

• weights - Slice of f32 weights to upload

Complexity

• Time: O(n) for data transfer
• Space: O(n) GPU memory

Examples

use arcweight::gpu::GpuContext;

let ctx = pollster::block_on(GpuContext::new())?;
let weights = vec![1.0f32, 2.0, 3.0, 4.0];
let gpu_buffer = ctx.create_weight_buffer(&weights)?;
assert_eq!(gpu_buffer.len(), 4);

pub fn reduce_min_gpu(&self, weights: &GpuWeightBuffer) -> Result<f32>

Compute minimum value using GPU parallel reduction.

Implements tropical semiring addition (minimum) using a multi-pass parallel reduction algorithm optimized for GPU execution.

Algorithm

Uses a tree-based parallel reduction:

1. Each workgroup (256 threads) reduces 256 elements to 1
2. Results are written to intermediate buffer
3. Repeat until single value remains

Complexity

• Time: O(n / p + log p) where p = number of GPU cores
• Space: O(n / 256) for intermediate buffers
• Passes: ceil(log_256(n))

Examples

use arcweight::gpu::GpuContext;

let ctx = pollster::block_on(GpuContext::new())?;
let weights: Vec<f32> = (0..10000).map(|i| i as f32).collect();
let buffer = ctx.create_weight_buffer(&weights)?;

let min = ctx.reduce_min_gpu(&buffer)?;
assert_eq!(min, 0.0);

pub fn reduce_sum_gpu(&self, weights: &GpuWeightBuffer) -> Result<f32>

Compute sum using GPU parallel reduction.

Implements probability semiring addition (sum) using a multi-pass parallel reduction algorithm.

Complexity

Same as reduce_min_gpu.

pub fn reduce_min(&self, weights: &GpuWeightBuffer) -> Result<f32>

Compute minimum value in a weight buffer (CPU fallback)

Uses parallel reduction on the GPU for O(log n) depth.

pub fn reduce_sum(&self, weights: &GpuWeightBuffer) -> Result<f32>

Compute sum of values in a weight buffer (CPU fallback)

pub fn read_buffer(&self, buffer: &GpuWeightBuffer) -> Result<Vec<f32>>

Read a buffer back to CPU memory

pub fn elementwise_min_gpu( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise minimum of two weight buffers using GPU

pub fn elementwise_min( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise minimum of two weight buffers

pub fn elementwise_add_gpu( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise addition of two weight buffers using GPU

pub fn elementwise_add( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise addition of two weight buffers (for tropical multiplication)

pub fn elementwise_mul_gpu( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise multiplication of two weight buffers using GPU

pub fn elementwise_mul( &self, a: &GpuWeightBuffer, b: &GpuWeightBuffer, ) -> Result<GpuWeightBuffer>

Element-wise multiplication (for probability semiring)

pub fn log_sum_exp_gpu(&self, weights: &GpuWeightBuffer) -> Result<f32>

Compute log-sum-exp for log semiring using GPU

This is numerically stable using the max-subtract trick.

pub fn batch_reduce_min(&self, weights_list: &[&[f32]]) -> Result<Vec<f32>>

Batch process multiple weight arrays

Efficiently processes multiple arrays in a single GPU dispatch

Trait Implementations

impl Debug for GpuContext

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.