282 lines
7.0 KiB
Rust
282 lines
7.0 KiB
Rust
// Copyright (c) 2024 RustyTorch Team
|
|
// Licensed under MIT OR Apache-2.0
|
|
//! GPU Backend Abstraction for CFD Operations
|
|
//!
|
|
//! This module provides a backend-agnostic interface for GPU-accelerated CFD kernels,
|
|
//! supporting both CUDA (NVIDIA) and Metal (Apple Silicon) backends.
|
|
//!
|
|
//! # Architecture
|
|
//!
|
|
//! ```text
|
|
//! GpuBackend trait
|
|
//! ├── CudaBackend (NVIDIA GPUs via cudarc)
|
|
//! └── MetalBackend (Apple Silicon via rtx-metal)
|
|
//! ```
|
|
//!
|
|
//! # Example
|
|
//!
|
|
//! ```rust,ignore
|
|
//! let backend = GpuBackend::auto_detect()?;
|
|
//! let buffer = backend.allocate::<f32>(1024)?;
|
|
//! backend.launch_kernel("poisson_jacobi", config, &[buffer])?;
|
|
//! ```
|
|
|
|
use crate::error::CfdResult;
|
|
use std::sync::Arc;
|
|
|
|
/// Configuration for kernel launches
|
|
#[derive(Debug, Clone)]
|
|
pub struct LaunchConfig {
|
|
/// Grid dimensions (width, height, depth)
|
|
pub grid: (usize, usize, usize),
|
|
/// Thread group / block dimensions
|
|
pub block: (usize, usize, usize),
|
|
/// Shared memory size in bytes
|
|
pub shared_memory: usize,
|
|
}
|
|
|
|
impl LaunchConfig {
|
|
/// Create a 1D launch configuration
|
|
pub fn new_1d(n: usize, block_size: usize) -> Self {
|
|
let grid = (n + block_size - 1) / block_size;
|
|
Self {
|
|
grid: (grid, 1, 1),
|
|
block: (block_size, 1, 1),
|
|
shared_memory: 0,
|
|
}
|
|
}
|
|
|
|
/// Create a 2D launch configuration
|
|
pub fn new_2d(nx: usize, ny: usize, block_x: usize, block_y: usize) -> Self {
|
|
let grid_x = (nx + block_x - 1) / block_x;
|
|
let grid_y = (ny + block_y - 1) / block_y;
|
|
Self {
|
|
grid: (grid_x, grid_y, 1),
|
|
block: (block_x, block_y, 1),
|
|
shared_memory: 0,
|
|
}
|
|
}
|
|
|
|
/// Create a 3D launch configuration
|
|
pub fn new_3d(
|
|
nx: usize,
|
|
ny: usize,
|
|
nz: usize,
|
|
block_x: usize,
|
|
block_y: usize,
|
|
block_z: usize,
|
|
) -> Self {
|
|
let grid_x = (nx + block_x - 1) / block_x;
|
|
let grid_y = (ny + block_y - 1) / block_y;
|
|
let grid_z = (nz + block_z - 1) / block_z;
|
|
Self {
|
|
grid: (grid_x, grid_y, grid_z),
|
|
block: (block_x, block_y, block_z),
|
|
shared_memory: 0,
|
|
}
|
|
}
|
|
|
|
/// Set shared memory size
|
|
pub fn with_shared_memory(mut self, size: usize) -> Self {
|
|
self.shared_memory = size;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// GPU buffer handle - opaque type for backend-specific storage
|
|
pub struct GpuBuffer {
|
|
/// Buffer ID for the backend
|
|
pub id: usize,
|
|
/// Size in elements
|
|
pub len: usize,
|
|
/// Element size in bytes
|
|
pub element_size: usize,
|
|
}
|
|
|
|
impl GpuBuffer {
|
|
/// Get total size in bytes
|
|
pub fn size_bytes(&self) -> usize {
|
|
self.len * self.element_size
|
|
}
|
|
}
|
|
|
|
/// Backend type enumeration
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum BackendType {
|
|
/// NVIDIA CUDA backend
|
|
Cuda,
|
|
/// Apple Metal backend
|
|
Metal,
|
|
/// CPU fallback
|
|
Cpu,
|
|
}
|
|
|
|
impl std::fmt::Display for BackendType {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
BackendType::Cuda => write!(f, "CUDA"),
|
|
BackendType::Metal => write!(f, "Metal"),
|
|
BackendType::Cpu => write!(f, "CPU"),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// GPU Backend trait for CFD operations
|
|
///
|
|
/// This trait abstracts the underlying GPU framework (CUDA/Metal),
|
|
/// providing a unified interface for CFD kernel execution.
|
|
pub trait GpuBackend: Send + Sync {
|
|
/// Get the backend type
|
|
fn backend_type(&self) -> BackendType;
|
|
|
|
/// Get device name
|
|
fn device_name(&self) -> &str;
|
|
|
|
/// Check if backend is available
|
|
fn is_available(&self) -> bool;
|
|
|
|
/// Allocate a buffer on the device
|
|
fn allocate(&self, len: usize, element_size: usize) -> CfdResult<GpuBuffer>;
|
|
|
|
/// Free a buffer
|
|
fn free(&self, buffer: &GpuBuffer) -> CfdResult<()>;
|
|
|
|
/// Copy data from host to device
|
|
fn copy_to_device(&self, buffer: &GpuBuffer, data: &[u8]) -> CfdResult<()>;
|
|
|
|
/// Copy data from device to host
|
|
fn copy_to_host(&self, buffer: &GpuBuffer, data: &mut [u8]) -> CfdResult<()>;
|
|
|
|
/// Synchronize all pending operations
|
|
fn synchronize(&self) -> CfdResult<()>;
|
|
|
|
/// Get maximum threads per block/threadgroup
|
|
fn max_threads_per_block(&self) -> usize;
|
|
|
|
/// Get device memory size in bytes
|
|
fn device_memory(&self) -> usize;
|
|
}
|
|
|
|
/// CPU fallback backend for testing and non-GPU systems
|
|
pub struct CpuBackend {
|
|
name: String,
|
|
}
|
|
|
|
impl CpuBackend {
|
|
/// Create a new CPU backend
|
|
pub fn new() -> Self {
|
|
Self {
|
|
name: "CPU".to_string(),
|
|
}
|
|
}
|
|
}
|
|
|
|
impl Default for CpuBackend {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl GpuBackend for CpuBackend {
|
|
fn backend_type(&self) -> BackendType {
|
|
BackendType::Cpu
|
|
}
|
|
|
|
fn device_name(&self) -> &str {
|
|
&self.name
|
|
}
|
|
|
|
fn is_available(&self) -> bool {
|
|
true
|
|
}
|
|
|
|
fn allocate(&self, len: usize, element_size: usize) -> CfdResult<GpuBuffer> {
|
|
Ok(GpuBuffer {
|
|
id: 0,
|
|
len,
|
|
element_size,
|
|
})
|
|
}
|
|
|
|
fn free(&self, _buffer: &GpuBuffer) -> CfdResult<()> {
|
|
Ok(())
|
|
}
|
|
|
|
fn copy_to_device(&self, _buffer: &GpuBuffer, _data: &[u8]) -> CfdResult<()> {
|
|
Ok(())
|
|
}
|
|
|
|
fn copy_to_host(&self, _buffer: &GpuBuffer, _data: &mut [u8]) -> CfdResult<()> {
|
|
Ok(())
|
|
}
|
|
|
|
fn synchronize(&self) -> CfdResult<()> {
|
|
Ok(())
|
|
}
|
|
|
|
fn max_threads_per_block(&self) -> usize {
|
|
1
|
|
}
|
|
|
|
fn device_memory(&self) -> usize {
|
|
// Return approximate available system memory
|
|
8 * 1024 * 1024 * 1024 // 8 GB default
|
|
}
|
|
}
|
|
|
|
/// Auto-detect and create the best available GPU backend
|
|
pub fn auto_detect_backend() -> Arc<dyn GpuBackend> {
|
|
// Try CUDA first (on any platform)
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
// TODO: Implement CUDA backend detection
|
|
tracing::info!("CUDA feature enabled but backend not yet implemented");
|
|
}
|
|
|
|
// Try Metal on macOS
|
|
#[cfg(all(target_os = "macos", feature = "metal"))]
|
|
{
|
|
use rtx_metal::MetalDevice;
|
|
if MetalDevice::is_available() {
|
|
tracing::info!("Metal backend available");
|
|
// TODO: Create and return Metal backend
|
|
}
|
|
}
|
|
|
|
// Fallback to CPU
|
|
tracing::info!("Using CPU fallback backend");
|
|
Arc::new(CpuBackend::new())
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_launch_config_1d() {
|
|
let config = LaunchConfig::new_1d(1024, 256);
|
|
assert_eq!(config.grid, (4, 1, 1));
|
|
assert_eq!(config.block, (256, 1, 1));
|
|
}
|
|
|
|
#[test]
|
|
fn test_launch_config_2d() {
|
|
let config = LaunchConfig::new_2d(64, 64, 16, 16);
|
|
assert_eq!(config.grid, (4, 4, 1));
|
|
assert_eq!(config.block, (16, 16, 1));
|
|
}
|
|
|
|
#[test]
|
|
fn test_cpu_backend() {
|
|
let backend = CpuBackend::new();
|
|
assert_eq!(backend.backend_type(), BackendType::Cpu);
|
|
assert!(backend.is_available());
|
|
}
|
|
|
|
#[test]
|
|
fn test_auto_detect() {
|
|
let backend = auto_detect_backend();
|
|
assert!(backend.is_available());
|
|
}
|
|
}
|