Omar Sobh
1381b5b633
CI / Build (macos-latest) (push) Waiting to run
CI / Test (macos-latest) (push) Blocked by required conditions
CI / Test (ubuntu-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (macos-latest) (push) Blocked by required conditions
CI / Python Bindings (maturin) (ubuntu-latest) (push) Blocked by required conditions
CI / WASM Build + Size Check (push) Blocked by required conditions
CI / Distributed Training Tests (push) Blocked by required conditions
CI / CI Success (push) Blocked by required conditions
CI / Build CPU-Only (Explicit) (push) Failing after 6s
CI / Build (ubuntu-latest) (push) Failing after 6s
Documentation / Build API Documentation (push) Failing after 6s
Documentation / Build User Guide (push) Successful in 7s
CI / Format Check (push) Failing after 13s
CI / Clippy Check (push) Failing after 37s
Performance Benchmarks / Run Benchmarks (push) Successful in 3m23s
885 lines
29 KiB
Rust
885 lines
29 KiB
Rust
//! CUDA kernel implementations for CFD operations
|
||
//!
|
||
//! This module provides real CUDA kernel compilation and execution using cudarc 0.17.3.
|
||
//! All kernels are compiled at runtime using NVRTC for maximum flexibility.
|
||
|
||
use super::*;
|
||
use cudarc::driver::{CudaContext, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
|
||
use cudarc::nvrtc::compile_ptx;
|
||
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
|
||
/// CUDA kernel manager for CFD operations
|
||
/// Handles kernel compilation, memory management, and execution
|
||
pub struct CudaKernelManager {
|
||
context: Arc<CudaContext>,
|
||
stream: Arc<CudaStream>,
|
||
modules: HashMap<String, Arc<CudaModule>>,
|
||
}
|
||
|
||
impl CudaKernelManager {
|
||
/// Create new CUDA kernel manager and compile all kernels
|
||
pub fn new(config: &CfdConfig) -> CfdResult<Self> {
|
||
let context = CudaContext::new(config.device_id as usize).map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to initialize CUDA context: {}", e))
|
||
})?;
|
||
let stream = context.default_stream();
|
||
|
||
let mut manager = Self {
|
||
context,
|
||
stream,
|
||
modules: HashMap::new(),
|
||
};
|
||
manager.compile_all_kernels()?;
|
||
Ok(manager)
|
||
}
|
||
|
||
/// Compile all kernels at startup
|
||
fn compile_all_kernels(&mut self) -> CfdResult<()> {
|
||
// Compile advection kernels
|
||
let advection_upwind_src = include_str!("cuda/advection_upwind.cu");
|
||
self.compile_kernel_source("advection_kernels", advection_upwind_src)?;
|
||
|
||
// Compile diffusion kernels
|
||
let diffusion_src = include_str!("cuda/diffusion.cu");
|
||
self.compile_kernel_source("diffusion_kernels", diffusion_src)?;
|
||
|
||
// Compile 2D implicit diffusion kernels
|
||
let diffusion_2d_src = include_str!("cuda/diffusion_2d_implicit.cu");
|
||
self.compile_kernel_source("diffusion_2d_kernels", diffusion_2d_src)?;
|
||
|
||
// Compile Poisson kernels
|
||
let poisson_src = include_str!("cuda/poisson.cu");
|
||
self.compile_kernel_source("poisson_kernels", poisson_src)?;
|
||
|
||
// Compile matrix operation kernels
|
||
let matrix_ops_src = include_str!("cuda/matrix_ops.cu");
|
||
self.compile_kernel_source("matrix_kernels", matrix_ops_src)?;
|
||
|
||
// Compile reduction kernels
|
||
let reduction_src = include_str!("cuda/reduction.cu");
|
||
self.compile_kernel_source("reduction_kernels", reduction_src)?;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Compile a CUDA kernel from source code
|
||
fn compile_kernel_source(&mut self, module_name: &str, source: &str) -> CfdResult<()> {
|
||
// Compile to PTX using NVRTC
|
||
let ptx = compile_ptx(source).map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to compile {} kernels: {}", module_name, e))
|
||
})?;
|
||
|
||
// Load PTX module into context
|
||
let module = self.context.load_module(ptx).map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to load {} module: {}", module_name, e))
|
||
})?;
|
||
|
||
self.modules.insert(module_name.to_string(), module);
|
||
Ok(())
|
||
}
|
||
|
||
/// Allocate GPU memory for f32 array
|
||
pub fn allocate_f32(&self, size: usize) -> CfdResult<CudaSlice<f32>> {
|
||
self.stream
|
||
.alloc_zeros::<f32>(size)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))
|
||
}
|
||
|
||
/// Copy data from host to device
|
||
pub fn copy_to_device(&self, host_data: &[f32]) -> CfdResult<CudaSlice<f32>> {
|
||
self.stream
|
||
.memcpy_stod(host_data)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Host to device copy failed: {}", e)))
|
||
}
|
||
|
||
/// Copy data from device to host
|
||
pub fn copy_from_device(&self, device_slice: &CudaSlice<f32>) -> CfdResult<Vec<f32>> {
|
||
self.stream
|
||
.memcpy_dtov(device_slice)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Device to host copy failed: {}", e)))
|
||
}
|
||
|
||
/// Synchronize stream
|
||
pub fn synchronize(&self) -> CfdResult<()> {
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Stream synchronization failed: {}", e)))
|
||
}
|
||
|
||
/// Get context reference
|
||
pub fn context(&self) -> &Arc<CudaContext> {
|
||
&self.context
|
||
}
|
||
|
||
/// Get stream reference
|
||
pub fn stream(&self) -> &Arc<CudaStream> {
|
||
&self.stream
|
||
}
|
||
|
||
/// Get module by name
|
||
pub fn get_module(&self, name: &str) -> CfdResult<&Arc<CudaModule>> {
|
||
self.modules
|
||
.get(name)
|
||
.ok_or_else(|| CfdError::gpu_error(&format!("Module {} not found", name)))
|
||
}
|
||
|
||
/// Perform GPU reduction to compute sum
|
||
pub fn reduce_sum(&self, input: &CudaSlice<f32>) -> CfdResult<f32> {
|
||
let n = input.len();
|
||
let block_size = 256;
|
||
let grid_size = ((n as u32 + block_size - 1) / block_size).min(1024);
|
||
|
||
// Allocate temporary buffer for partial results
|
||
let mut partial_results = self.allocate_f32(grid_size as usize)?;
|
||
|
||
// Load reduction module
|
||
let module = self.get_module("reduction_kernels")?;
|
||
|
||
// First pass: reduce to grid_size partial results
|
||
let func = module
|
||
.load_function("reduce_sum")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get reduce_sum kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(input)
|
||
.arg(&mut partial_results)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Reduce sum kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
// Second pass: final reduction
|
||
let final_func = module
|
||
.load_function("reduce_final_sum")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get final sum kernel: {}", e)))?;
|
||
|
||
let mut final_result = self.allocate_f32(1)?;
|
||
|
||
let final_config = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 32 * std::mem::size_of::<f32>() as u32,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&final_func)
|
||
.arg(&partial_results)
|
||
.arg(&mut final_result)
|
||
.arg(&(grid_size as i32))
|
||
.launch(final_config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Final sum kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
self.synchronize()?;
|
||
|
||
// Copy result back
|
||
let mut result_host = vec![0.0f32; 1];
|
||
self.stream
|
||
.memcpy_dtoh(&final_result, &mut result_host)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to copy result from GPU: {}", e)))?;
|
||
|
||
Ok(result_host[0])
|
||
}
|
||
|
||
/// Compute maximum absolute value using GPU reduction
|
||
pub fn reduce_abs_max(&self, input: &CudaSlice<f32>) -> CfdResult<f32> {
|
||
let n = input.len();
|
||
let block_size = 256;
|
||
let grid_size = ((n as u32 + block_size - 1) / block_size).min(1024);
|
||
|
||
// Allocate temporary buffer for partial results
|
||
let mut partial_results = self
|
||
.stream
|
||
.alloc_zeros::<f32>(grid_size as usize)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
|
||
let module = self.get_module("reduction_kernels")?;
|
||
let reduce_func = module
|
||
.load_function("reduce_abs_max")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get reduction kernel: {}", e)))?;
|
||
|
||
// First reduction pass
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&reduce_func)
|
||
.arg(input)
|
||
.arg(&mut partial_results)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Reduction kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
// Second reduction pass if needed
|
||
if grid_size > 1 {
|
||
let mut final_result = self
|
||
.stream
|
||
.alloc_zeros::<f32>(1)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
|
||
let final_func = module.load_function("reduce_final_max").map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to get final reduction kernel: {}", e))
|
||
})?;
|
||
|
||
let final_config = LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&final_func)
|
||
.arg(&partial_results)
|
||
.arg(&mut final_result)
|
||
.arg(&(grid_size as i32))
|
||
.launch(final_config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Final reduction failed: {}", e)))?;
|
||
}
|
||
|
||
self.stream.synchronize()?;
|
||
let result_host = self.copy_from_device(&final_result)?;
|
||
Ok(result_host[0])
|
||
} else {
|
||
self.stream.synchronize()?;
|
||
let result_host = self.copy_from_device(&partial_results)?;
|
||
Ok(result_host[0])
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Advection kernel wrapper for various schemes
|
||
pub struct AdvectionKernel {
|
||
scheme: AdvectionScheme,
|
||
stream: Arc<CudaStream>,
|
||
module: Arc<CudaModule>,
|
||
}
|
||
|
||
impl AdvectionKernel {
|
||
pub fn new(manager: &CudaKernelManager, scheme: AdvectionScheme) -> CfdResult<Self> {
|
||
Ok(Self {
|
||
scheme,
|
||
stream: manager.stream.clone(),
|
||
module: manager.get_module("advection_kernels")?.clone(),
|
||
})
|
||
}
|
||
|
||
pub fn apply(
|
||
&self,
|
||
phi: &CudaSlice<f32>,
|
||
phi_new: &mut CudaSlice<f32>,
|
||
velocity: f32,
|
||
dt: f32,
|
||
dx: f32,
|
||
) -> CfdResult<()> {
|
||
let n = phi.len();
|
||
let grid_size = (n as u32 + 255) / 256;
|
||
let block_size = 256;
|
||
|
||
let kernel_name = match self.scheme {
|
||
AdvectionScheme::Upwind => "advection_1d",
|
||
AdvectionScheme::Central => "advection_central_1d",
|
||
AdvectionScheme::Quick => "advection_quick_1d",
|
||
AdvectionScheme::Weno => "advection_weno_1d",
|
||
};
|
||
|
||
let func = self.module.load_function(kernel_name).map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
|
||
})?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(phi)
|
||
.arg(phi_new)
|
||
.arg(&velocity)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn apply_2d(
|
||
&self,
|
||
phi: &CudaSlice<f32>,
|
||
phi_new: &mut CudaSlice<f32>,
|
||
u: &CudaSlice<f32>,
|
||
v: &CudaSlice<f32>,
|
||
dt: f32,
|
||
dx: f32,
|
||
dy: f32,
|
||
nx: usize,
|
||
ny: usize,
|
||
) -> CfdResult<()> {
|
||
let grid_dim_x = (nx as u32 + 15) / 16;
|
||
let grid_dim_y = (ny as u32 + 15) / 16;
|
||
|
||
let kernel_name = match self.scheme {
|
||
AdvectionScheme::Upwind => "advection_2d",
|
||
AdvectionScheme::Central => "advection_central_2d",
|
||
AdvectionScheme::Quick => "advection_quick_2d",
|
||
AdvectionScheme::Weno => "advection_weno_2d",
|
||
};
|
||
|
||
let func = self.module.load_function(kernel_name).map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
|
||
})?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_dim_x, grid_dim_y, 1),
|
||
block_dim: (16, 16, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(phi)
|
||
.arg(phi_new)
|
||
.arg(u)
|
||
.arg(v)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&dy)
|
||
.arg(&(nx as i32))
|
||
.arg(&(ny as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
/// Diffusion kernel wrapper for various schemes
|
||
pub struct DiffusionKernel {
|
||
scheme: DiffusionScheme,
|
||
stream: Arc<CudaStream>,
|
||
diffusion_module: Arc<CudaModule>,
|
||
diffusion_2d_module: Arc<CudaModule>,
|
||
}
|
||
|
||
impl DiffusionKernel {
|
||
pub fn new(manager: &CudaKernelManager, scheme: DiffusionScheme) -> CfdResult<Self> {
|
||
Ok(Self {
|
||
scheme,
|
||
stream: manager.stream.clone(),
|
||
diffusion_module: manager.get_module("diffusion_kernels")?.clone(),
|
||
diffusion_2d_module: manager.get_module("diffusion_2d_kernels")?.clone(),
|
||
})
|
||
}
|
||
|
||
pub fn apply(
|
||
&self,
|
||
temp: &CudaSlice<f32>,
|
||
temp_new: &mut CudaSlice<f32>,
|
||
alpha: f32,
|
||
dt: f32,
|
||
dx: f32,
|
||
) -> CfdResult<()> {
|
||
let n = temp.len();
|
||
let grid_size = (n as u32 + 255) / 256;
|
||
let block_size = 256;
|
||
|
||
let kernel_name = match self.scheme {
|
||
DiffusionScheme::Explicit => "diffusion_explicit_1d",
|
||
DiffusionScheme::Implicit => "diffusion_implicit_1d",
|
||
DiffusionScheme::CrankNicolson => "diffusion_crank_nicolson_1d",
|
||
};
|
||
|
||
let func = self
|
||
.diffusion_module
|
||
.load_function(kernel_name)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Failed to get kernel {}: {}", kernel_name, e))
|
||
})?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
match self.scheme {
|
||
DiffusionScheme::Explicit => unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(temp)
|
||
.arg(temp_new)
|
||
.arg(&alpha)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
},
|
||
DiffusionScheme::Implicit | DiffusionScheme::CrankNicolson => {
|
||
// For implicit schemes, use iterative solver with convergence check
|
||
let max_iterations = 100;
|
||
let tolerance = 1e-6;
|
||
let mut residual = 1.0;
|
||
let mut iteration = 0;
|
||
|
||
while residual > tolerance && iteration < max_iterations {
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(temp)
|
||
.arg(&mut *temp_new)
|
||
.arg(&alpha)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
self.stream.synchronize().map_err(|e| {
|
||
CfdError::gpu_error(&format!("Synchronization failed: {}", e))
|
||
})?;
|
||
|
||
// TODO: Compute actual residual on GPU
|
||
residual *= 0.9; // Simplified convergence
|
||
iteration += 1;
|
||
}
|
||
return Ok(());
|
||
}
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
Ok(())
|
||
}
|
||
|
||
pub fn apply_2d(
|
||
&self,
|
||
temp: &CudaSlice<f32>,
|
||
temp_new: &mut CudaSlice<f32>,
|
||
alpha: f32,
|
||
dt: f32,
|
||
dx: f32,
|
||
dy: f32,
|
||
nx: usize,
|
||
ny: usize,
|
||
) -> CfdResult<()> {
|
||
let grid_dim_x = (nx as u32 + 15) / 16;
|
||
let grid_dim_y = (ny as u32 + 15) / 16;
|
||
|
||
match self.scheme {
|
||
DiffusionScheme::Explicit => {
|
||
let func = self
|
||
.diffusion_module
|
||
.load_function("diffusion_explicit_2d")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_dim_x, grid_dim_y, 1),
|
||
block_dim: (16, 16, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(temp)
|
||
.arg(temp_new)
|
||
.arg(&alpha)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&dy)
|
||
.arg(&(nx as i32))
|
||
.arg(&(ny as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
Ok(())
|
||
}
|
||
DiffusionScheme::Implicit | DiffusionScheme::CrankNicolson => {
|
||
// Implement 2D implicit/Crank-Nicolson using ADI (Alternating Direction Implicit)
|
||
let kernel_name = match self.scheme {
|
||
DiffusionScheme::Implicit => "diffusion_implicit_2d_adi",
|
||
DiffusionScheme::CrankNicolson => "diffusion_crank_nicolson_2d_adi",
|
||
_ => unreachable!(),
|
||
};
|
||
|
||
let func = self
|
||
.diffusion_2d_module
|
||
.load_function(kernel_name)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_dim_x, grid_dim_y, 1),
|
||
block_dim: (16, 16, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
// ADI requires two sweeps: x-direction then y-direction
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(temp)
|
||
.arg(temp_new)
|
||
.arg(&alpha)
|
||
.arg(&dt)
|
||
.arg(&dx)
|
||
.arg(&dy)
|
||
.arg(&(nx as i32))
|
||
.arg(&(ny as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
Ok(())
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Poisson equation solver kernel
|
||
pub struct PoissonKernel {
|
||
stream: Arc<CudaStream>,
|
||
module: Arc<CudaModule>,
|
||
kernel_manager: Arc<CudaKernelManager>,
|
||
}
|
||
|
||
impl PoissonKernel {
|
||
pub fn new(manager: &Arc<CudaKernelManager>) -> CfdResult<Self> {
|
||
Ok(Self {
|
||
stream: manager.stream.clone(),
|
||
module: manager.get_module("poisson_kernels")?.clone(),
|
||
kernel_manager: manager.clone(),
|
||
})
|
||
}
|
||
|
||
/// [`Self::solve_2d`] under its older name.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn solve_jacobi_2d(
|
||
&self,
|
||
phi: &mut CudaSlice<f32>,
|
||
source: &CudaSlice<f32>,
|
||
nx: usize,
|
||
ny: usize,
|
||
dx: f32,
|
||
dy: f32,
|
||
max_iterations: usize,
|
||
tolerance: f32,
|
||
) -> CfdResult<usize> {
|
||
self.solve_2d(phi, source, nx, ny, dx, dy, max_iterations, tolerance)
|
||
}
|
||
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn solve_2d(
|
||
&self,
|
||
phi: &mut CudaSlice<f32>,
|
||
source: &CudaSlice<f32>,
|
||
nx: usize,
|
||
ny: usize,
|
||
dx: f32,
|
||
dy: f32,
|
||
max_iterations: usize,
|
||
tolerance: f32,
|
||
) -> CfdResult<usize> {
|
||
let grid_dim_x = (nx as u32 + 15) / 16;
|
||
let grid_dim_y = (ny as u32 + 15) / 16;
|
||
|
||
// Calculate factors for Poisson equation
|
||
let dx2_inv = 1.0f32 / (dx * dx);
|
||
let dy2_inv = 1.0f32 / (dy * dy);
|
||
let factor = 1.0f32 / (2.0f32 * (dx2_inv + dy2_inv));
|
||
|
||
// Allocate temporary arrays for iteration
|
||
let mut phi_temp = self
|
||
.stream
|
||
.alloc_zeros::<f32>(nx * ny)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
let mut residual = self
|
||
.stream
|
||
.alloc_zeros::<f32>(nx * ny)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
|
||
let jacobi_func = self
|
||
.module
|
||
.load_function("poisson_jacobi_2d")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get Jacobi kernel: {}", e)))?;
|
||
|
||
let residual_func = self
|
||
.module
|
||
.load_function("poisson_residual_2d")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get residual kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_dim_x, grid_dim_y, 1),
|
||
block_dim: (16, 16, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
for iteration in 0..max_iterations {
|
||
// Perform Jacobi iteration
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&jacobi_func)
|
||
.arg(&mut phi_temp)
|
||
.arg(&*phi)
|
||
.arg(source)
|
||
.arg(&factor)
|
||
.arg(&dx2_inv)
|
||
.arg(&dy2_inv)
|
||
.arg(&(nx as i32))
|
||
.arg(&(ny as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Jacobi kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
// Swap buffers
|
||
std::mem::swap(&mut *phi, &mut phi_temp);
|
||
|
||
// Check convergence every 10 iterations
|
||
if iteration % 10 == 0 {
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&residual_func)
|
||
.arg(&mut residual)
|
||
.arg(&*phi)
|
||
.arg(source)
|
||
.arg(&dx2_inv)
|
||
.arg(&dy2_inv)
|
||
.arg(&(nx as i32))
|
||
.arg(&(ny as i32))
|
||
.launch(config)
|
||
.map_err(|e| {
|
||
CfdError::gpu_error(&format!("Residual kernel launch failed: {}", e))
|
||
})?;
|
||
}
|
||
|
||
// Use GPU reduction to find max residual
|
||
let max_residual = self.kernel_manager.reduce_abs_max(&residual)?;
|
||
|
||
if max_residual < tolerance {
|
||
return Ok(iteration + 1);
|
||
}
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
}
|
||
|
||
Ok(max_iterations)
|
||
}
|
||
}
|
||
|
||
/// Matrix operations kernel for sparse linear algebra
|
||
pub struct MatrixOpsKernel {
|
||
stream: Arc<CudaStream>,
|
||
module: Arc<CudaModule>,
|
||
kernel_manager: Arc<CudaKernelManager>,
|
||
}
|
||
|
||
impl MatrixOpsKernel {
|
||
pub fn new(manager: &Arc<CudaKernelManager>) -> CfdResult<Self> {
|
||
Ok(Self {
|
||
stream: manager.stream.clone(),
|
||
module: manager.get_module("matrix_kernels")?.clone(),
|
||
kernel_manager: manager.clone(),
|
||
})
|
||
}
|
||
|
||
pub fn dot_product(&self, x: &CudaSlice<f32>, y: &CudaSlice<f32>) -> CfdResult<f32> {
|
||
let n = x.len();
|
||
let grid_size = (n as u32 + 255) / 256;
|
||
let block_size = 256;
|
||
|
||
// Allocate result buffer for partial sums
|
||
let mut partial_sums = self
|
||
.stream
|
||
.alloc_zeros::<f32>(grid_size as usize)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
|
||
let func = self
|
||
.module
|
||
.load_function("dot_product_partial")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: block_size as u32 * std::mem::size_of::<f32>() as u32,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(x)
|
||
.arg(y)
|
||
.arg(&mut partial_sums)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
|
||
// Use GPU reduction to compute final sum
|
||
let result = self.kernel_manager.reduce_sum(&partial_sums)?;
|
||
Ok(result)
|
||
}
|
||
|
||
/// `y = T x` for the tridiagonal matrix with `diagonal` (n) and
|
||
/// `off_diagonal` (n − 1) — the `tridiagonal_matvec` kernel.
|
||
pub fn tridiagonal_matvec(
|
||
&self,
|
||
diagonal: &CudaSlice<f32>,
|
||
off_diagonal: &CudaSlice<f32>,
|
||
x: &CudaSlice<f32>,
|
||
y: &mut CudaSlice<f32>,
|
||
) -> CfdResult<()> {
|
||
let n = x.len();
|
||
let func = self
|
||
.module
|
||
.load_function("tridiagonal_matvec")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
let config = LaunchConfig {
|
||
grid_dim: ((n as u32 + 255) / 256, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(diagonal)
|
||
.arg(off_diagonal)
|
||
.arg(x)
|
||
.arg(y)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Compute L2 norm of vector
|
||
pub fn vector_norm(&self, x: &CudaSlice<f32>) -> CfdResult<f32> {
|
||
let n = x.len();
|
||
let grid_size = (n as u32 + 255) / 256;
|
||
let block_size = 256;
|
||
|
||
let mut partial_sums = self
|
||
.stream
|
||
.alloc_zeros::<f32>(grid_size as usize)
|
||
.map_err(|e| CfdError::gpu_error(&format!("GPU allocation failed: {}", e)))?;
|
||
|
||
let func = self
|
||
.module
|
||
.load_function("l2_norm_partial")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: block_size as u32 * std::mem::size_of::<f32>() as u32,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(x)
|
||
.arg(&mut partial_sums)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
|
||
let partial_host = self
|
||
.stream
|
||
.memcpy_dtov(&partial_sums)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Device to host copy failed: {}", e)))?;
|
||
|
||
Ok(partial_host.iter().sum::<f32>().sqrt())
|
||
}
|
||
|
||
/// Compute y = alpha * x + y (AXPY operation)
|
||
pub fn axpy(&self, alpha: f32, x: &CudaSlice<f32>, y: &mut CudaSlice<f32>) -> CfdResult<()> {
|
||
let n = x.len();
|
||
let grid_size = (n as u32 + 255) / 256;
|
||
let block_size = 256;
|
||
|
||
let func = self
|
||
.module
|
||
.load_function("vector_axpy")
|
||
.map_err(|e| CfdError::gpu_error(&format!("Failed to get kernel: {}", e)))?;
|
||
|
||
let config = LaunchConfig {
|
||
grid_dim: (grid_size, 1, 1),
|
||
block_dim: (block_size, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
};
|
||
|
||
unsafe {
|
||
self.stream
|
||
.launch_builder(&func)
|
||
.arg(&alpha)
|
||
.arg(x)
|
||
.arg(y)
|
||
.arg(&(n as i32))
|
||
.launch(config)
|
||
.map_err(|e| CfdError::gpu_error(&format!("Kernel launch failed: {}", e)))?;
|
||
}
|
||
|
||
self.stream
|
||
.synchronize()
|
||
.map_err(|e| CfdError::gpu_error(&format!("Synchronization failed: {}", e)))?;
|
||
|
||
Ok(())
|
||
}
|
||
}
|