Omar SobhandClaude Sonnet 4.6
228137555f
G0 (Critical): Replace 45 unimplemented!() panics across three GPU backends - rtx-backend-cuda: sin/cos/tanh via PTX, relu/sigmoid/leaky_relu/elu via activation.rs, pow/clamp/gt_scalar via unary.rs, var/var_dim host-side, conv2d/max_pool2d/avg_pool2d CPU fallback in new ops/conv.rs; new PTX kernels in element_wise.cu - rtx-backend-rocm: all 15 ops via CPU round-trip (to_vec → compute → from_slice) - rtx-backend-sycl: all 15 ops via CPU round-trip (to_host → compute → from_data) G2 (High): Re-add rtx-distributed to workspace - Vendor 4 minimal RNCCL stub crates at crates/vendor/rnccl/* - Update rtx-distributed RNCCL path deps to point at stubs (../../../../RNCCL/* → ../../vendor/rnccl/*) - Remove rtx-distributed from workspace exclude list, add to members G5 (Medium): Re-enable rtx-tts (213 tests restored) - Fix 15 rtx-nn API drift issues: LayerNorm::new, Conv1d::from_config, Conv1dPadding::Zeros, Dropout::new(p, device), tensor methods (relu/tanh/sigmoid/cat/stack), squeeze(Some(n)), to_vec() turbofish removal, Tensor::randn with &[...] slices G8 (Low): Quantum stubs + multimodal forward bug - rtx-timeseries: remove dead quantum/neuromorphic TODO comment blocks (no module files exist) - rtx-multimodal/fusion/transformer.rs: wire TransformerBlock loop in forward() - rtx-multimodal/fusion/strategies.rs: wire bottleneck_layers loop in forward() - rtx-transformers/architectures/transformer_block.rs: add forward() method (pre-norm residuals; full attention+FFN pending when those sub-layers are wired) Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
479 lines
14 KiB
Rust
479 lines
14 KiB
Rust
//! # RustyTorch++ ROCm Backend
|
|
//!
|
|
//! AMD GPU backend implementation using HIP runtime.
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - **HIP Runtime**: AMD's CUDA-compatible API for GPU computing
|
|
//! - **rocBLAS**: Optimized matrix operations
|
|
//! - **MIOpen**: Deep learning primitives
|
|
//! - **Source Compatibility**: HIP kernels are ~95% CUDA-compatible
|
|
//!
|
|
//! ## Architecture
|
|
//!
|
|
//! ```text
|
|
//! RocmBackend
|
|
//! ├── RocmTensorPrimitive - HIP device memory
|
|
//! ├── RocmDevice - Device context and streams
|
|
//! └── Ops
|
|
//! ├── Basic - Add, mul, etc. (HIP kernels)
|
|
//! ├── GEMM - Matrix multiply (rocBLAS)
|
|
//! └── Attention - Flash Attention (ported from CUDA)
|
|
//! ```
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! use rtx_backend_rocm::{RocmBackend, RocmDevice};
|
|
//! use rtx_backend::Backend;
|
|
//!
|
|
//! let device = RocmDevice::new(0)?;
|
|
//! let a = RocmBackend::zeros([1024, 1024], &device);
|
|
//! let b = RocmBackend::randn([1024, 1024], &device);
|
|
//! let c = RocmBackend::matmul(&a, &b);
|
|
//! ```
|
|
|
|
#![warn(missing_docs)]
|
|
|
|
mod device;
|
|
mod error;
|
|
/// HIP FFI bindings for low-level GPU access.
|
|
pub mod hip_ffi;
|
|
/// HIP kernel implementations for GPU operations.
|
|
#[cfg(feature = "hip-runtime")]
|
|
pub mod hip_kernels;
|
|
/// Operations module - re-exported for tests and direct access.
|
|
pub mod ops;
|
|
mod tensor;
|
|
|
|
pub use device::{RocmDevice, RocmDeviceInfo};
|
|
pub use error::{RocmBackendError, RocmBackendResult};
|
|
pub use tensor::{RocmBuffer, RocmTensorPrimitive};
|
|
|
|
use rtx_backend::{Backend, BoolU8};
|
|
|
|
/// ROCm backend for RustyTorch++.
|
|
///
|
|
/// This backend uses AMD GPUs via HIP runtime and provides:
|
|
/// - rocBLAS for matrix operations
|
|
/// - HIP kernels (CUDA source-compatible)
|
|
/// - Multi-GPU support
|
|
#[derive(Clone, Debug, Default)]
|
|
pub struct RocmBackend;
|
|
|
|
impl Backend for RocmBackend {
|
|
type TensorPrimitive<const D: usize> = RocmTensorPrimitive<D>;
|
|
type Device = RocmDevice;
|
|
type FloatElem = f32;
|
|
type IntElem = i32;
|
|
type BoolElem = BoolU8;
|
|
|
|
fn name() -> &'static str {
|
|
"rocm"
|
|
}
|
|
|
|
fn seed(seed: u64) {
|
|
ops::seed_rng(seed);
|
|
}
|
|
|
|
// ==================== Tensor Creation ====================
|
|
|
|
fn zeros<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
|
ops::creation::zeros(shape, device)
|
|
}
|
|
|
|
fn ones<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
|
ops::creation::ones(shape, device)
|
|
}
|
|
|
|
fn full<const D: usize>(
|
|
shape: [usize; D],
|
|
fill_value: Self::FloatElem,
|
|
device: &Self::Device,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::creation::full(shape, fill_value, device)
|
|
}
|
|
|
|
fn rand<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
|
ops::creation::rand(shape, device)
|
|
}
|
|
|
|
fn randn<const D: usize>(shape: [usize; D], device: &Self::Device) -> Self::TensorPrimitive<D> {
|
|
ops::creation::randn(shape, device)
|
|
}
|
|
|
|
fn from_data<const D: usize>(
|
|
data: &[Self::FloatElem],
|
|
shape: [usize; D],
|
|
device: &Self::Device,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::creation::from_data(data, shape, device)
|
|
}
|
|
|
|
// ==================== Basic Operations ====================
|
|
|
|
fn add<const D: usize>(
|
|
lhs: Self::TensorPrimitive<D>,
|
|
rhs: Self::TensorPrimitive<D>,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::basic::add(&lhs, &rhs)
|
|
}
|
|
|
|
fn sub<const D: usize>(
|
|
lhs: Self::TensorPrimitive<D>,
|
|
rhs: Self::TensorPrimitive<D>,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::basic::sub(&lhs, &rhs)
|
|
}
|
|
|
|
fn mul<const D: usize>(
|
|
lhs: Self::TensorPrimitive<D>,
|
|
rhs: Self::TensorPrimitive<D>,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::basic::mul(&lhs, &rhs)
|
|
}
|
|
|
|
fn div<const D: usize>(
|
|
lhs: Self::TensorPrimitive<D>,
|
|
rhs: Self::TensorPrimitive<D>,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::basic::div(&lhs, &rhs)
|
|
}
|
|
|
|
fn matmul(
|
|
lhs: Self::TensorPrimitive<2>,
|
|
rhs: Self::TensorPrimitive<2>,
|
|
) -> Self::TensorPrimitive<2> {
|
|
ops::gemm::matmul(&lhs, &rhs)
|
|
}
|
|
|
|
fn bmm(
|
|
lhs: Self::TensorPrimitive<3>,
|
|
rhs: Self::TensorPrimitive<3>,
|
|
) -> Self::TensorPrimitive<3> {
|
|
ops::gemm::bmm(&lhs, &rhs)
|
|
}
|
|
|
|
// ==================== Unary Operations ====================
|
|
|
|
fn neg<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::neg(&tensor)
|
|
}
|
|
|
|
fn exp<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::exp(&tensor)
|
|
}
|
|
|
|
fn log<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::log(&tensor)
|
|
}
|
|
|
|
fn sqrt<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::sqrt(&tensor)
|
|
}
|
|
|
|
fn abs<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::abs(&tensor)
|
|
}
|
|
|
|
fn sin<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::sin(&tensor)
|
|
}
|
|
|
|
fn cos<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::unary::cos(&tensor)
|
|
}
|
|
|
|
fn pow<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
exp: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::unary::pow(&tensor, exp)
|
|
}
|
|
|
|
fn clamp<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
min: Self::FloatElem,
|
|
max: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::comparison::clamp(&tensor, Some(min), Some(max))
|
|
}
|
|
|
|
// ==================== Activation Functions ====================
|
|
|
|
fn relu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::activation::relu(&tensor)
|
|
}
|
|
|
|
fn sigmoid<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::activation::sigmoid(&tensor)
|
|
}
|
|
|
|
fn tanh<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::activation::tanh(&tensor)
|
|
}
|
|
|
|
// ==================== Reduction Operations ====================
|
|
|
|
fn sum<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
|
ops::reduction::sum(&tensor)
|
|
}
|
|
|
|
fn sum_dim<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
dim: usize,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::reduction::sum_dim(&tensor, dim)
|
|
}
|
|
|
|
fn mean<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
|
ops::reduction::mean(&tensor)
|
|
}
|
|
|
|
fn mean_dim<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
dim: usize,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::reduction::mean_dim(&tensor, dim)
|
|
}
|
|
|
|
fn var<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
|
ops::reduction::var(&tensor)
|
|
}
|
|
|
|
fn var_dim<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
dim: usize,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::reduction::var_dim(&tensor, dim)
|
|
}
|
|
|
|
fn max<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
|
ops::reduction::max(&tensor)
|
|
}
|
|
|
|
fn min<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<1> {
|
|
ops::reduction::min(&tensor)
|
|
}
|
|
|
|
// ==================== Shape Operations ====================
|
|
|
|
fn shape<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> [usize; D] {
|
|
tensor.shape
|
|
}
|
|
|
|
fn reshape<const D1: usize, const D2: usize>(
|
|
tensor: Self::TensorPrimitive<D1>,
|
|
shape: [usize; D2],
|
|
) -> Self::TensorPrimitive<D2> {
|
|
ops::shape::reshape(tensor, shape)
|
|
}
|
|
|
|
fn transpose<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::shape::transpose(&tensor)
|
|
}
|
|
|
|
fn swap_dims<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
dim1: usize,
|
|
dim2: usize,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::shape::swap_dims(&tensor, dim1, dim2)
|
|
}
|
|
|
|
// ==================== LLM-Specific Operations ====================
|
|
|
|
fn flash_attention(
|
|
query: Self::TensorPrimitive<4>,
|
|
key: Self::TensorPrimitive<4>,
|
|
value: Self::TensorPrimitive<4>,
|
|
mask: Option<&Self::TensorPrimitive<4>>,
|
|
scale: Self::FloatElem,
|
|
causal: bool,
|
|
) -> Self::TensorPrimitive<4> {
|
|
ops::attention::flash_attention(&query, &key, &value, mask, scale, causal)
|
|
}
|
|
|
|
fn softmax<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
dim: usize,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::activation::softmax(&tensor, dim)
|
|
}
|
|
|
|
fn layer_norm<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
weight: &Self::TensorPrimitive<1>,
|
|
bias: Option<&Self::TensorPrimitive<1>>,
|
|
eps: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::normalization::layer_norm(&tensor, weight, bias, eps)
|
|
}
|
|
|
|
fn rms_norm<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
weight: &Self::TensorPrimitive<1>,
|
|
eps: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::normalization::rms_norm(&tensor, weight, eps)
|
|
}
|
|
|
|
fn rope<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
cos: &Self::TensorPrimitive<2>,
|
|
sin: &Self::TensorPrimitive<2>,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::attention::rope(&tensor, cos, sin)
|
|
}
|
|
|
|
fn gelu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::activation::gelu(&tensor)
|
|
}
|
|
|
|
fn silu<const D: usize>(tensor: Self::TensorPrimitive<D>) -> Self::TensorPrimitive<D> {
|
|
ops::activation::silu(&tensor)
|
|
}
|
|
|
|
fn leaky_relu<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
negative_slope: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::activation::leaky_relu(&tensor, negative_slope)
|
|
}
|
|
|
|
fn elu<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
alpha: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::activation::elu(&tensor, alpha)
|
|
}
|
|
|
|
// ==================== Comparison Operations ====================
|
|
|
|
fn gt_scalar<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
value: Self::FloatElem,
|
|
) -> Self::TensorPrimitive<D> {
|
|
ops::comparison::gt_scalar(&tensor, value)
|
|
}
|
|
|
|
// ==================== Convolution Operations ====================
|
|
|
|
fn conv2d(
|
|
input: Self::TensorPrimitive<4>,
|
|
weight: &Self::TensorPrimitive<4>,
|
|
bias: Option<&Self::TensorPrimitive<1>>,
|
|
stride: [usize; 2],
|
|
padding: [usize; 2],
|
|
dilation: [usize; 2],
|
|
groups: usize,
|
|
) -> Self::TensorPrimitive<4> {
|
|
let config = ops::convolution::Conv2dConfig {
|
|
kernel_size: (weight.shape[2], weight.shape[3]),
|
|
stride: (stride[0], stride[1]),
|
|
padding: (padding[0], padding[1]),
|
|
dilation: (dilation[0], dilation[1]),
|
|
groups,
|
|
};
|
|
ops::convolution::conv2d(&input, weight, bias, &config)
|
|
}
|
|
|
|
// ==================== Pooling Operations ====================
|
|
|
|
fn max_pool2d(
|
|
input: Self::TensorPrimitive<4>,
|
|
kernel_size: [usize; 2],
|
|
stride: [usize; 2],
|
|
padding: [usize; 2],
|
|
) -> Self::TensorPrimitive<4> {
|
|
let config = ops::pooling::Pool2dConfig {
|
|
kernel_size: (kernel_size[0], kernel_size[1]),
|
|
stride: (stride[0], stride[1]),
|
|
padding: (padding[0], padding[1]),
|
|
dilation: (1, 1),
|
|
ceil_mode: false,
|
|
};
|
|
ops::pooling::max_pool2d(&input, &config)
|
|
}
|
|
|
|
fn avg_pool2d(
|
|
input: Self::TensorPrimitive<4>,
|
|
kernel_size: [usize; 2],
|
|
stride: [usize; 2],
|
|
padding: [usize; 2],
|
|
count_include_pad: bool,
|
|
) -> Self::TensorPrimitive<4> {
|
|
let config = ops::pooling::Pool2dConfig {
|
|
kernel_size: (kernel_size[0], kernel_size[1]),
|
|
stride: (stride[0], stride[1]),
|
|
padding: (padding[0], padding[1]),
|
|
dilation: (1, 1),
|
|
ceil_mode: false,
|
|
};
|
|
ops::pooling::avg_pool2d(&input, &config, count_include_pad)
|
|
}
|
|
|
|
// ==================== Device Management ====================
|
|
|
|
fn device<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Self::Device {
|
|
tensor.device.clone()
|
|
}
|
|
|
|
fn to_device<const D: usize>(
|
|
tensor: Self::TensorPrimitive<D>,
|
|
device: &Self::Device,
|
|
) -> Self::TensorPrimitive<D> {
|
|
if tensor.device == *device {
|
|
tensor
|
|
} else {
|
|
ops::device::copy_to_device(tensor, device)
|
|
}
|
|
}
|
|
|
|
fn to_data<const D: usize>(tensor: &Self::TensorPrimitive<D>) -> Vec<Self::FloatElem> {
|
|
ops::device::copy_to_host(tensor)
|
|
}
|
|
|
|
fn sync(device: &Self::Device) {
|
|
let _ = device.synchronize();
|
|
}
|
|
}
|
|
|
|
/// Check if ROCm is available on this system.
|
|
pub fn is_available() -> bool {
|
|
#[cfg(feature = "hip-runtime")]
|
|
{
|
|
hip_ffi::HipRuntime::is_available()
|
|
}
|
|
#[cfg(not(feature = "hip-runtime"))]
|
|
{
|
|
device::is_available()
|
|
}
|
|
}
|
|
|
|
/// Get the number of available AMD GPUs.
|
|
pub fn device_count() -> usize {
|
|
hip_ffi::HipRuntime::device_count()
|
|
}
|
|
|
|
/// Check if HIP runtime feature is enabled.
|
|
pub fn has_hip_runtime() -> bool {
|
|
cfg!(feature = "hip-runtime")
|
|
}
|
|
|
|
/// Check if rocBLAS feature is enabled.
|
|
pub fn has_rocblas() -> bool {
|
|
cfg!(feature = "rocblas")
|
|
}
|
|
|
|
/// Check if MIOpen feature is enabled.
|
|
pub fn has_miopen() -> bool {
|
|
cfg!(feature = "miopen")
|
|
}
|
|
|
|
/// Type alias for training with ROCm + autodiff.
|
|
pub type RocmTraining = RocmBackend;
|
|
|
|
/// Type alias for inference with ROCm (no autodiff overhead).
|
|
pub type RocmInference = RocmBackend;
|