Crate tsai_compute 

tsai_compute

A heterogeneous compute abstraction layer for tsai-rs.

This crate provides a unified interface for compute operations across different hardware backends:

• CPU: SIMD-optimized (AVX2, AVX-512, NEON) with optional NUMA awareness
• Metal: Apple GPU via Metal framework (macOS/iOS)
• CUDA: NVIDIA GPU via CUDA toolkit
• Vulkan: Cross-platform GPU via Vulkan
• OpenCL: Cross-platform GPU/CPU via OpenCL
• ROCm: AMD GPU via ROCm/HIP

Features

• Runtime device discovery and selection
• Unified buffer and memory management
• Command encoding and synchronization
• Workload-aware scheduling
• Burn framework integration (optional)

Quick Start

use tsai_compute::{HardwareDiscovery, DevicePool};

// Discover all available devices
let pool = HardwareDiscovery::discover_all()?;
pool.print_summary();

// Get the best device
let device = pool.best_device().expect("No devices available");
println!("Using: {}", device.name());

Feature Flags

• cpu (default): CPU backend with SIMD
• numa: NUMA-aware memory allocation
• cuda: NVIDIA CUDA support
• metal: Apple Metal support (macOS only)
• vulkan: Vulkan compute support
• opencl: OpenCL support
• rocm: AMD ROCm support
• burn-bridge: Burn framework integration

Re-exports

pub use backend::CommandEncoder;

pub use backend::ComputeBackend;

pub use backend::Fence;

pub use device::ComputeDevice;

pub use device::DeviceCapabilities;

pub use device::DeviceFeature;

pub use device::DeviceId;

pub use device::DevicePool;

pub use device::DeviceType;

pub use device::SelectionStrategy;

pub use device::SimdLevel;

pub use discovery::get_device_pool;

pub use discovery::get_discovery_time_us;

pub use discovery::refresh_device_pool;

pub use discovery::HardwareDiscovery;

pub use error::ComputeError;

pub use error::ComputeResult;

pub use memory::Buffer;

pub use memory::BufferUsage;

pub use memory::MemoryPool;

pub use scheduler::Priority;

pub use scheduler::RoundRobinScheduler;

pub use scheduler::Scheduler;

pub use scheduler::SimpleScheduler;

pub use scheduler::Workload;

pub use scheduler::WorkloadScheduler;

pub use backend::cpu::CpuBackend;

pub use backend::cpu::CpuBuffer;

pub use backend::cpu::CpuDevice;

Modules

backend

Compute backend abstractions.

bridge

Integration bridges with other frameworks.

device

Device abstractions for heterogeneous compute.

discovery

Hardware discovery orchestration.

error

Unified error types for tsai_compute.

memory

Memory management abstractions.

prelude

Prelude module for convenient imports.

scheduler

Workload scheduling for heterogeneous compute.