rustytorch/crates/specialized/rtx-cfd/src/turbulence/k_epsilon_gpu.rs

//! GPU-accelerated k-epsilon turbulence model
//!
//! This module provides CUDA-accelerated implementation of the k-epsilon turbulence model
//! for RANS simulations, including transport equation solving, production term calculation,
//! and eddy viscosity computation.

use super::KEpsilonConstants;
use crate::kernels::CudaKernelManager;
use crate::{CfdConfig, CfdError, CfdResult};
use cudarc::driver::{CudaModule, CudaSlice, LaunchConfig, PushKernelArg};
use std::sync::Arc;

/// GPU memory buffers for k-epsilon model
struct KEpsilonGpuBuffers {
    /// Grid dimensions
    nx: usize,
    ny: usize,
    nz: usize,

    /// Turbulent kinetic energy
    k: CudaSlice<f32>,
    k_old: CudaSlice<f32>,
    /// Dissipation rate
    epsilon: CudaSlice<f32>,
    epsilon_old: CudaSlice<f32>,

    /// Velocity gradients
    dudx: CudaSlice<f32>,
    dudy: CudaSlice<f32>,
    dudz: CudaSlice<f32>,
    dvdx: CudaSlice<f32>,
    dvdy: CudaSlice<f32>,
    dvdz: CudaSlice<f32>,
    dwdx: CudaSlice<f32>,
    dwdy: CudaSlice<f32>,
    dwdz: CudaSlice<f32>,

    /// Production and dissipation terms
    production: CudaSlice<f32>,
    dissipation: CudaSlice<f32>,
    eddy_viscosity: CudaSlice<f32>,

    /// Source terms
    source_k: CudaSlice<f32>,
    source_epsilon: CudaSlice<f32>,

    /// Temporary arrays
    temp1: CudaSlice<f32>,
    temp2: CudaSlice<f32>,
}

/// GPU-accelerated k-epsilon turbulence model
pub struct KEpsilonGpuModel {
    /// GPU kernel manager
    kernel_manager: CudaKernelManager,
    /// Model constants
    constants: KEpsilonConstants,
    /// K-epsilon CUDA module
    k_epsilon_module: Option<Arc<CudaModule>>,
    /// GPU memory buffers
    gpu_buffers: Option<KEpsilonGpuBuffers>,
}

impl KEpsilonGpuModel {
    /// Create new k-epsilon GPU model
    pub fn new(config: &CfdConfig) -> CfdResult<Self> {
        let kernel_manager = CudaKernelManager::new(config)?;

        Ok(Self {
            kernel_manager,
            constants: KEpsilonConstants::standard(),
            k_epsilon_module: None,
            gpu_buffers: None,
        })
    }

    /// Initialize model with grid dimensions
    pub fn initialize(&mut self, nx: usize, ny: usize, nz: usize) -> CfdResult<()> {
        // Load k-epsilon kernels
        self.load_k_epsilon_kernels()?;

        // Allocate GPU buffers
        self.allocate_gpu_buffers(nx, ny, nz)?;

        Ok(())
    }

    /// Load and compile k-epsilon CUDA kernels
    fn load_k_epsilon_kernels(&mut self) -> CfdResult<()> {
        // In real implementation, would compile or load PTX for k-epsilon kernels
        // For now, load from kernel manager
        self.k_epsilon_module = Some(self.kernel_manager.get_module("k_epsilon_kernels")?.clone());
        Ok(())
    }

    /// Allocate GPU memory buffers
    fn allocate_gpu_buffers(&mut self, nx: usize, ny: usize, nz: usize) -> CfdResult<()> {
        let n = nx * ny * nz;

        let buffers = KEpsilonGpuBuffers {
            nx,
            ny,
            nz,
            k: self.kernel_manager.allocate_f32(n)?,
            k_old: self.kernel_manager.allocate_f32(n)?,
            epsilon: self.kernel_manager.allocate_f32(n)?,
            epsilon_old: self.kernel_manager.allocate_f32(n)?,
            dudx: self.kernel_manager.allocate_f32(n)?,
            dudy: self.kernel_manager.allocate_f32(n)?,
            dudz: self.kernel_manager.allocate_f32(n)?,
            dvdx: self.kernel_manager.allocate_f32(n)?,
            dvdy: self.kernel_manager.allocate_f32(n)?,
            dvdz: self.kernel_manager.allocate_f32(n)?,
            dwdx: self.kernel_manager.allocate_f32(n)?,
            dwdy: self.kernel_manager.allocate_f32(n)?,
            dwdz: self.kernel_manager.allocate_f32(n)?,
            production: self.kernel_manager.allocate_f32(n)?,
            dissipation: self.kernel_manager.allocate_f32(n)?,
            eddy_viscosity: self.kernel_manager.allocate_f32(n)?,
            source_k: self.kernel_manager.allocate_f32(n)?,
            source_epsilon: self.kernel_manager.allocate_f32(n)?,
            temp1: self.kernel_manager.allocate_f32(n)?,
            temp2: self.kernel_manager.allocate_f32(n)?,
        };

        self.gpu_buffers = Some(buffers);
        Ok(())
    }

    /// Set initial conditions for k and epsilon
    pub fn set_initial_conditions(&mut self, k_init: f32, epsilon_init: f32) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        // Initialize k and epsilon with uniform values
        let n = buffers.nx * buffers.ny * buffers.nz;
        let k_host = vec![k_init; n];
        let epsilon_host = vec![epsilon_init; n];

        self.kernel_manager
            .stream()
            .memcpy_htod(&k_host, &mut buffers.k)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy k to GPU: {}", e)))?;
        self.kernel_manager
            .stream()
            .memcpy_htod(&epsilon_host, &mut buffers.epsilon)
            .map_err(|e| CfdError::gpu_error(&format!("Failed to copy epsilon to GPU: {}", e)))?;

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// GPU-accelerated computation of velocity gradients
    pub fn compute_velocity_gradients(
        &mut self,
        u: &CudaSlice<f32>,
        v: &CudaSlice<f32>,
        w: &CudaSlice<f32>,
        dx: f32,
        dy: f32,
        dz: f32,
    ) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .k_epsilon_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("K-epsilon module not loaded"))?;
        let func = module
            .load_function("compute_velocity_gradients")
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to get velocity gradients kernel: {}", e))
            })?;

        let grid_dim_x = (buffers.nx as u32 + 7) / 8;
        let grid_dim_y = (buffers.ny as u32 + 7) / 8;
        let grid_dim_z = (buffers.nz as u32 + 7) / 8;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, grid_dim_z),
            block_dim: (8, 8, 8),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(u)
                .arg(v)
                .arg(w)
                .arg(&mut buffers.dudx)
                .arg(&mut buffers.dudy)
                .arg(&mut buffers.dudz)
                .arg(&mut buffers.dvdx)
                .arg(&mut buffers.dvdy)
                .arg(&mut buffers.dvdz)
                .arg(&mut buffers.dwdx)
                .arg(&mut buffers.dwdy)
                .arg(&mut buffers.dwdz)
                .arg(&dx)
                .arg(&dy)
                .arg(&dz)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .arg(&(buffers.nz as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Velocity gradients kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Compute eddy viscosity on GPU
    pub fn compute_eddy_viscosity(&mut self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .k_epsilon_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("K-epsilon module not loaded"))?;
        let func = module
            .load_function("compute_eddy_viscosity")
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to get eddy viscosity kernel: {}", e))
            })?;

        let grid_dim_x = (buffers.nx as u32 + 7) / 8;
        let grid_dim_y = (buffers.ny as u32 + 7) / 8;
        let grid_dim_z = (buffers.nz as u32 + 7) / 8;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, grid_dim_z),
            block_dim: (8, 8, 8),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&buffers.k)
                .arg(&buffers.epsilon)
                .arg(&mut buffers.eddy_viscosity)
                .arg(&self.constants.c_mu)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .arg(&(buffers.nz as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Eddy viscosity kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Compute production term P_k
    pub fn compute_production(&mut self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .k_epsilon_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("K-epsilon module not loaded"))?;
        let func = module
            .load_function("compute_production")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get production kernel: {}", e)))?;

        let grid_dim_x = (buffers.nx as u32 + 7) / 8;
        let grid_dim_y = (buffers.ny as u32 + 7) / 8;
        let grid_dim_z = (buffers.nz as u32 + 7) / 8;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, grid_dim_z),
            block_dim: (8, 8, 8),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&buffers.eddy_viscosity)
                .arg(&buffers.dudx)
                .arg(&buffers.dudy)
                .arg(&buffers.dudz)
                .arg(&buffers.dvdx)
                .arg(&buffers.dvdy)
                .arg(&buffers.dvdz)
                .arg(&buffers.dwdx)
                .arg(&buffers.dwdy)
                .arg(&buffers.dwdz)
                .arg(&mut buffers.production)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .arg(&(buffers.nz as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Production kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Compute source terms for epsilon equation
    pub fn compute_epsilon_source(&mut self) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let module = self
            .k_epsilon_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("K-epsilon module not loaded"))?;
        let func = module
            .load_function("compute_epsilon_source")
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to get epsilon source kernel: {}", e))
            })?;

        let grid_dim_x = (buffers.nx as u32 + 7) / 8;
        let grid_dim_y = (buffers.ny as u32 + 7) / 8;
        let grid_dim_z = (buffers.nz as u32 + 7) / 8;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, grid_dim_z),
            block_dim: (8, 8, 8),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&buffers.k)
                .arg(&buffers.epsilon)
                .arg(&buffers.production)
                .arg(&mut buffers.source_epsilon)
                .arg(&self.constants.c1_epsilon)
                .arg(&self.constants.c2_epsilon)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .arg(&(buffers.nz as i32))
                .launch(config)
                .map_err(|e| {
                    CfdError::gpu_error(&format!("Epsilon source kernel launch failed: {}", e))
                })?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Update k and epsilon using implicit time integration
    pub fn update_k_epsilon(&mut self, dt: f32, nu: f32) -> CfdResult<()> {
        let buffers = self
            .gpu_buffers
            .as_mut()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        // Store old values
        buffers.k_old = buffers.k.clone();
        buffers.epsilon_old = buffers.epsilon.clone();

        let module = self
            .k_epsilon_module
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("K-epsilon module not loaded"))?;
        let func = module
            .load_function("update_k_epsilon")
            .map_err(|e| CfdError::gpu_error(&format!("Failed to get update kernel: {}", e)))?;

        let grid_dim_x = (buffers.nx as u32 + 7) / 8;
        let grid_dim_y = (buffers.ny as u32 + 7) / 8;
        let grid_dim_z = (buffers.nz as u32 + 7) / 8;

        let config = LaunchConfig {
            grid_dim: (grid_dim_x, grid_dim_y, grid_dim_z),
            block_dim: (8, 8, 8),
            shared_mem_bytes: 0,
        };

        unsafe {
            self.kernel_manager
                .stream()
                .launch_builder(&func)
                .arg(&mut buffers.k)
                .arg(&mut buffers.epsilon)
                .arg(&buffers.k_old)
                .arg(&buffers.epsilon_old)
                .arg(&buffers.production)
                .arg(&buffers.source_epsilon)
                .arg(&buffers.eddy_viscosity)
                .arg(&dt)
                .arg(&nu)
                .arg(&self.constants.sigma_k)
                .arg(&self.constants.sigma_epsilon)
                .arg(&(buffers.nx as i32))
                .arg(&(buffers.ny as i32))
                .arg(&(buffers.nz as i32))
                .launch(config)
                .map_err(|e| CfdError::gpu_error(&format!("Update kernel launch failed: {}", e)))?;
        }

        self.kernel_manager.synchronize()?;
        Ok(())
    }

    /// Get eddy viscosity field from GPU
    pub fn get_eddy_viscosity(&self) -> CfdResult<Vec<f32>> {
        let buffers = self
            .gpu_buffers
            .as_ref()
            .ok_or_else(|| CfdError::gpu_error("GPU buffers not initialized"))?;

        let n = buffers.nx * buffers.ny * buffers.nz;
        let mut eddy_viscosity_host = vec![0.0f32; n];

        self.kernel_manager
            .stream()
            .memcpy_dtoh(&buffers.eddy_viscosity, &mut eddy_viscosity_host)
            .map_err(|e| {
                CfdError::gpu_error(&format!("Failed to copy eddy viscosity from GPU: {}", e))
            })?;

        self.kernel_manager.synchronize()?;
        Ok(eddy_viscosity_host)
    }

    /// Complete time step for k-epsilon model
    pub fn step(
        &mut self,
        u: &CudaSlice<f32>,
        v: &CudaSlice<f32>,
        w: &CudaSlice<f32>,
        dx: f32,
        dy: f32,
        dz: f32,
        dt: f32,
        nu: f32,
    ) -> CfdResult<()> {
        // Compute velocity gradients
        self.compute_velocity_gradients(u, v, w, dx, dy, dz)?;

        // Compute eddy viscosity
        self.compute_eddy_viscosity()?;

        // Compute production term
        self.compute_production()?;

        // Compute source terms
        self.compute_epsilon_source()?;

        // Update k and epsilon
        self.update_k_epsilon(dt, nu)?;

        Ok(())
    }
}