use super::{KernelMetrics, KernelType};
use crate::error::CryptoError;
use crate::types::Algorithm;
#[cfg(feature = "gpu-cuda")]
const SHA256_BLOCK_SIZE: usize = 64;
#[cfg(feature = "gpu-cuda")]
const SHA256_DIGEST_SIZE: usize = 32;
#[cfg(feature = "gpu-cuda")]
const SHA512_BLOCK_SIZE: usize = 128;
#[cfg(feature = "gpu-cuda")]
const SHA512_DIGEST_SIZE: usize = 64;
#[cfg(feature = "gpu-cuda")]
const SM3_BLOCK_SIZE: usize = 64;
#[cfg(feature = "gpu-cuda")]
const SM3_DIGEST_SIZE: usize = 32;
#[cfg(feature = "gpu-cuda")]
const CUDA_SHA256_KERNEL: &[u8] = include_bytes!("shaders/sha256.ptx");
#[cfg(feature = "gpu-cuda")]
const CUDA_SHA512_KERNEL: &[u8] = include_bytes!("shaders/sha512.ptx");
#[cfg(feature = "gpu-cuda")]
const CUDA_SM3_KERNEL: &[u8] = include_bytes!("shaders/sm3.ptx");
#[cfg(feature = "gpu-cuda")]
struct CudaHashKernelState {
context: Option<CudaContext>,
device: Option<CudaDevice>,
stream: Option<CudaStream>,
sha256_kernel: Option<CudaKernel>,
sha512_kernel: Option<CudaKernel>,
sm3_kernel: Option<CudaKernel>,
memory_pool: Vec<CudaMemory>,
config: HashKernelConfig,
metrics: Mutex<KernelMetrics>,
initialized: bool,
}
#[cfg(feature = "gpu-cuda")]
impl CudaHashKernelState {
pub fn new(config: HashKernelConfig) -> Self {
Self {
context: None,
device: None,
stream: None,
sha256_kernel: None,
sha512_kernel: None,
sm3_kernel: None,
memory_pool: Vec::new(),
config,
metrics: Mutex::new(KernelMetrics::new(KernelType::GpuSha2)),
initialized: false,
}
}
fn allocate_from_pool(&mut self, size: usize) -> Result<CudaMemory> {
let index = self
.memory_pool
.iter()
.position(|m| m.size() >= size && m.is_free());
if let Some(idx) = index {
let mem = self.memory_pool.remove(idx);
mem.set_used(true);
return Ok(mem);
}
let new_mem = CudaMemory::new(size)?;
new_mem.set_used(true);
self.memory_pool.push(new_mem.clone());
Ok(new_mem)
}
fn release_to_pool(&mut self, memory: CudaMemory) {
memory.set_used(false);
if !self.memory_pool.contains(&memory) {
self.memory_pool.push(memory);
}
}
}
#[cfg(feature = "gpu-cuda")]
pub struct CudaHashKernel {
state: Mutex<CudaHashKernelState>,
is_available: bool,
}
#[cfg(feature = "gpu-cuda")]
impl CudaHashKernel {
pub fn new() -> Self {
let config = HashKernelConfig::default();
let state = Mutex::new(CudaHashKernelState::new(config));
let is_available = Self::check_cuda_availability();
Self {
state,
is_available,
}
}
fn check_cuda_availability() -> bool {
match CudaDevice::enumerate() {
Ok(devices) => !devices.is_empty(),
Err(_) => false,
}
}
fn initialize_internal(&mut self) -> Result<()> {
let mut state = self
.state
.lock()
.map_err(|e| CryptoError::InitializationFailed(format!("Mutex poisoned: {}", e)))?;
if state.initialized {
return Ok(());
}
let devices = CudaDevice::enumerate().map_err(|e| {
CryptoError::InitializationFailed(format!("Failed to enumerate CUDA devices: {}", e))
})?;
if devices.is_empty() {
return Err(CryptoError::HardwareAccelerationUnavailable(
"No CUDA devices found".into(),
));
}
let device = devices.into_iter().next().unwrap();
let context = CudaContext::new(&device).map_err(|e| {
CryptoError::InitializationFailed(format!("Failed to create CUDA context: {}", e))
})?;
let stream = CudaStream::new().map_err(|e| {
CryptoError::InitializationFailed(format!("Failed to create CUDA stream: {}", e))
})?;
let sha256_kernel = CudaKernel::new(&context, CUDA_SHA256_KERNEL, "sha256_kernel").ok();
let sha512_kernel = CudaKernel::new(&context, CUDA_SHA512_KERNEL, "sha512_kernel").ok();
let sm3_kernel = CudaKernel::new(&context, CUDA_SM3_KERNEL, "sm3_kernel").ok();
state.context = Some(context);
state.device = Some(device);
state.stream = Some(stream);
state.sha256_kernel = sha256_kernel;
state.sha512_kernel = sha512_kernel;
state.sm3_kernel = sm3_kernel;
state.initialized = true;
Ok(())
}
fn shutdown_internal(&mut self) -> Result<()> {
let mut state = self
.state
.lock()
.map_err(|e| CryptoError::InitializationFailed(format!("Mutex poisoned: {}", e)))?;
if !state.initialized {
return Ok(());
}
state.sha256_kernel = None;
state.sha512_kernel = None;
state.sm3_kernel = None;
state.stream = None;
state.context = None;
state.memory_pool.clear();
state.initialized = false;
Ok(())
}
fn execute_sha256_gpu(&self, data: &[u8]) -> Result<Vec<u8>> {
let state = self
.state
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
let start = std::time::Instant::now();
let ctx = state
.context
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("CUDA context not initialized".into()))?;
let kernel = state
.sha256_kernel
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("SHA256 kernel not loaded".into()))?;
let stream = state
.stream
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("CUDA stream not initialized".into()))?;
let block_count = (data.len() + SHA256_BLOCK_SIZE - 1) / SHA256_BLOCK_SIZE;
let input_memory = CudaMemory::new(data.len()).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!("Failed to allocate input memory: {}", e))
})?;
let output_memory = CudaMemory::new(SHA256_DIGEST_SIZE).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!("Failed to allocate output memory: {}", e))
})?;
input_memory.copy_from(data).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy data to GPU: {}", e))
})?;
let grid_dim = (block_count as u32, 1, 1);
let block_dim = (256, 1, 1);
kernel
.launch(
&stream,
grid_dim,
block_dim,
&[
input_memory.as_ptr() as *mut std::ffi::c_void,
output_memory.as_ptr() as *mut std::ffi::c_void,
&(data.len() as u32),
],
)
.map_err(|e| {
CryptoError::KernelLaunchFailed(format!("Failed to launch SHA256 kernel: {}", e))
})?;
stream.synchronize().map_err(|e| {
CryptoError::SynchronizationFailed(format!("Failed to synchronize stream: {}", e))
})?;
let mut result = vec![0u8; SHA256_DIGEST_SIZE];
output_memory.copy_to(&mut result).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy result from GPU: {}", e))
})?;
let elapsed = start.elapsed();
let mut metrics = state
.metrics
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
metrics.execution_time_us = elapsed.as_micros() as u64;
metrics.throughput_mbps =
(data.len() as f32 / 1024.0 / 1024.0) / (elapsed.as_secs_f32() + 0.000001);
metrics.memory_transferred_bytes = data.len() + result.len();
metrics.compute_units_used = state
.device
.as_ref()
.map(|d| d.compute_capability().0)
.unwrap_or(0) as u32;
Ok(result)
}
fn execute_sha512_gpu(&self, data: &[u8]) -> Result<Vec<u8>> {
let state = self
.state
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
let start = std::time::Instant::now();
let ctx = state
.context
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("CUDA context not initialized".into()))?;
let kernel = state
.sha512_kernel
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("SHA512 kernel not loaded".into()))?;
let stream = state
.stream
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("CUDA stream not initialized".into()))?;
let block_count = (data.len() + SHA512_BLOCK_SIZE - 1) / SHA512_BLOCK_SIZE;
let input_memory = CudaMemory::new(data.len()).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!("Failed to allocate input memory: {}", e))
})?;
let output_memory = CudaMemory::new(SHA512_DIGEST_SIZE).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!("Failed to allocate output memory: {}", e))
})?;
input_memory.copy_from(data).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy data to GPU: {}", e))
})?;
let grid_dim = (block_count as u32, 1, 1);
let block_dim = (256, 1, 1);
kernel
.launch(
&stream,
grid_dim,
block_dim,
&[
input_memory.as_ptr() as *mut std::ffi::c_void,
output_memory.as_ptr() as *mut std::ffi::c_void,
&(data.len() as u32),
],
)
.map_err(|e| {
CryptoError::KernelLaunchFailed(format!("Failed to launch SHA512 kernel: {}", e))
})?;
stream.synchronize().map_err(|e| {
CryptoError::SynchronizationFailed(format!("Failed to synchronize stream: {}", e))
})?;
let mut result = vec![0u8; SHA512_DIGEST_SIZE];
output_memory.copy_to(&mut result).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy result from GPU: {}", e))
})?;
let elapsed = start.elapsed();
let mut metrics = state
.metrics
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
metrics.execution_time_us = elapsed.as_micros() as u64;
metrics.throughput_mbps =
(data.len() as f32 / 1024.0 / 1024.0) / (elapsed.as_secs_f32() + 0.000001);
metrics.memory_transferred_bytes = data.len() + result.len();
Ok(result)
}
}
#[cfg(feature = "gpu-cuda")]
impl super::GpuKernel for CudaHashKernel {
fn kernel_type(&self) -> KernelType {
KernelType::GpuSha2
}
fn supported_algorithms(&self) -> Vec<Algorithm> {
vec![
Algorithm::SHA256,
Algorithm::SHA384,
Algorithm::SHA512,
Algorithm::SM3,
]
}
fn is_available(&self) -> bool {
self.is_available
}
fn initialize(&mut self) -> Result<()> {
self.initialize_internal()
}
fn shutdown(&mut self) -> Result<()> {
self.shutdown_internal()
}
fn get_metrics(&self) -> Option<KernelMetrics> {
self.state
.lock()
.ok()
.map(|s| s.metrics.lock().unwrap().clone())
}
fn reset_metrics(&mut self) {
if let Ok(mut state) = self.state.lock() {
let mut metrics = state.metrics.lock().unwrap();
*metrics = KernelMetrics::new(KernelType::GpuSha2);
}
}
fn execute_hash(&self, data: &[u8], algorithm: Algorithm) -> Result<Vec<u8>> {
match algorithm {
Algorithm::SHA256 => self.execute_sha256_gpu(data),
Algorithm::SHA384 => {
let result = self.execute_sha512_gpu(data)?;
Ok(result[..48].to_vec())
}
Algorithm::SHA512 => self.execute_sha512_gpu(data),
Algorithm::SM3 => {
let state = self
.state
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
let ctx = state.context.as_ref().ok_or_else(|| {
CryptoError::NotInitialized("CUDA context not initialized".into())
})?;
let kernel = state
.sm3_kernel
.as_ref()
.ok_or_else(|| CryptoError::NotInitialized("SM3 kernel not loaded".into()))?;
let start = std::time::Instant::now();
let stream = state.stream.as_ref().ok_or_else(|| {
CryptoError::NotInitialized("CUDA stream not initialized".into())
})?;
let block_count = (data.len() + SM3_BLOCK_SIZE - 1) / SM3_BLOCK_SIZE;
let input_memory = CudaMemory::new(data.len()).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!(
"Failed to allocate input memory: {}",
e
))
})?;
let output_memory = CudaMemory::new(SM3_DIGEST_SIZE).map_err(|e| {
CryptoError::MemoryAllocationFailed(format!(
"Failed to allocate output memory: {}",
e
))
})?;
input_memory.copy_from(data).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy data to GPU: {}", e))
})?;
let grid_dim = (block_count as u32, 1, 1);
let block_dim = (256, 1, 1);
kernel
.launch(
&stream,
grid_dim,
block_dim,
&[
input_memory.as_ptr() as *mut std::ffi::c_void,
output_memory.as_ptr() as *mut std::ffi::c_void,
&(data.len() as u32),
],
)
.map_err(|e| {
CryptoError::KernelLaunchFailed(format!(
"Failed to launch SM3 kernel: {}",
e
))
})?;
stream.synchronize().map_err(|e| {
CryptoError::SynchronizationFailed(format!(
"Failed to synchronize stream: {}",
e
))
})?;
let mut result = vec![0u8; SM3_DIGEST_SIZE];
output_memory.copy_to(&mut result).map_err(|e| {
CryptoError::MemoryCopyFailed(format!("Failed to copy result from GPU: {}", e))
})?;
let elapsed = start.elapsed();
let mut metrics = state
.metrics
.lock()
.map_err(|e| CryptoError::OperationFailed(format!("Mutex poisoned: {}", e)))?;
metrics.execution_time_us = elapsed.as_micros() as u64;
metrics.throughput_mbps =
(data.len() as f32 / 1024.0 / 1024.0) / (elapsed.as_secs_f32() + 0.000001);
metrics.memory_transferred_bytes = data.len() + result.len();
Ok(result)
}
_ => Err(CryptoError::InvalidInput(
format!("Unsupported hash algorithm: {:?}", algorithm).into(),
)),
}
}
fn execute_hash_batch(&self, data: &[Vec<u8>], algorithm: Algorithm) -> Result<Vec<Vec<u8>>> {
let start = std::time::Instant::now();
let mut results = Vec::with_capacity(data.len());
for chunk in data {
let hash = self.execute_hash(chunk, algorithm)?;
results.push(hash);
}
let elapsed = start.elapsed();
if let Ok(mut state) = self.state.lock() {
let mut metrics = state.metrics.lock().unwrap();
metrics.execution_time_us = elapsed.as_micros() as u64;
metrics.batch_size = data.len();
}
Ok(results)
}
fn execute_aes_gcm_encrypt(
&self,
_key: &[u8],
_nonce: &[u8],
_data: &[u8],
_aad: Option<&[u8]>,
) -> Result<Vec<u8>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_decrypt(
&self,
_key: &[u8],
_nonce: &[u8],
_data: &[u8],
_aad: Option<&[u8]>,
) -> Result<Vec<u8>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_encrypt_batch(
&self,
_keys: &[&[u8]],
_nonces: &[&[u8]],
_data: &[&[u8]],
) -> Result<Vec<Vec<u8>>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_decrypt_batch(
&self,
_keys: &[&[u8]],
_nonces: &[&[u8]],
_data: &[&[u8]],
) -> Result<Vec<Vec<u8>>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
}
#[cfg(feature = "gpu-cuda")]
impl Default for CudaHashKernel {
fn default() -> Self {
Self::new()
}
}
#[cfg(feature = "gpu-cuda")]
mod cuda_driver {
use super::*;
pub struct CudaContext {
device: CudaDevice,
primary: bool,
}
impl CudaContext {
pub fn new(device: &CudaDevice) -> Result<Self> {
Ok(Self {
device: device.clone(),
primary: true,
})
}
}
#[derive(Clone)]
pub struct CudaDevice {
id: usize,
name: String,
compute_capability: (u32, u32),
total_memory: usize,
max_threads_per_block: i32,
}
impl CudaDevice {
pub fn enumerate() -> Result<Vec<Self>> {
Ok(Vec::new())
}
pub fn compute_capability(&self) -> (u32, u32) {
self.compute_capability
}
}
pub struct CudaKernel {
module: Vec<u8>,
function_name: String,
}
impl CudaKernel {
pub fn new(_context: &CudaContext, _ptx_code: &[u8], _name: &str) -> Result<Self> {
Ok(Self {
module: Vec::new(),
function_name: String::new(),
})
}
pub fn launch<S: AsRef<str>>(
&self,
_stream: &CudaStream,
_grid_dim: (u32, u32, u32),
_block_dim: (u32, u32, u32),
_arguments: &[*mut std::ffi::c_void],
) -> Result<()> {
Ok(())
}
}
#[derive(Clone)]
pub struct CudaMemory {
size: usize,
ptr: *mut std::ffi::c_void,
is_used: std::sync::atomic::AtomicBool,
}
impl CudaMemory {
pub fn new(size: usize) -> Result<Self> {
Ok(Self {
size,
ptr: std::ptr::null_mut(),
is_used: std::sync::atomic::AtomicBool::new(false),
})
}
pub fn size(&self) -> usize {
self.size
}
pub fn is_free(&self) -> bool {
!self.is_used.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn set_used(&self, used: bool) {
self.is_used
.store(used, std::sync::atomic::Ordering::Relaxed);
}
pub fn as_ptr(&self) -> *mut std::ffi::c_void {
self.ptr
}
pub fn copy_from(&self, data: &[u8]) -> Result<()> {
Ok(())
}
pub fn copy_to(&self, buffer: &mut [u8]) -> Result<()> {
Ok(())
}
}
pub struct CudaStream {
id: u64,
}
impl CudaStream {
pub fn new() -> Result<Self> {
Ok(Self { id: 0 })
}
pub fn synchronize(&self) -> Result<()> {
Ok(())
}
}
}
#[cfg(not(feature = "gpu-cuda"))]
pub struct CudaHashKernel;
#[cfg(not(feature = "gpu-cuda"))]
impl CudaHashKernel {
pub fn new() -> Self {
Self
}
pub fn is_available() -> bool {
false
}
}
#[cfg(not(feature = "gpu-cuda"))]
impl super::GpuKernel for CudaHashKernel {
fn kernel_type(&self) -> KernelType {
KernelType::Unknown
}
fn supported_algorithms(&self) -> Vec<Algorithm> {
Vec::new()
}
fn is_available(&self) -> bool {
false
}
fn initialize(&mut self) -> Result<()> {
Err(CryptoError::HardwareAccelerationUnavailable(
"CUDA support not enabled".into(),
))
}
fn shutdown(&mut self) -> Result<()> {
Ok(())
}
fn get_metrics(&self) -> Option<KernelMetrics> {
None
}
fn reset_metrics(&mut self) {}
fn execute_hash(&self, _data: &[u8], _algorithm: Algorithm) -> Result<Vec<u8>> {
Err(CryptoError::HardwareAccelerationUnavailable(
"CUDA support not enabled".into(),
))
}
fn execute_hash_batch(&self, _data: &[Vec<u8>], _algorithm: Algorithm) -> Result<Vec<Vec<u8>>> {
Err(CryptoError::HardwareAccelerationUnavailable(
"CUDA support not enabled".into(),
))
}
fn execute_aes_gcm_encrypt(
&self,
_key: &[u8],
_nonce: &[u8],
_data: &[u8],
_aad: Option<&[u8]>,
) -> Result<Vec<u8>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_decrypt(
&self,
_key: &[u8],
_nonce: &[u8],
_data: &[u8],
_aad: Option<&[u8]>,
) -> Result<Vec<u8>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_encrypt_batch(
&self,
_keys: &[&[u8]],
_nonces: &[&[u8]],
_data: &[&[u8]],
) -> Result<Vec<Vec<u8>>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
fn execute_aes_gcm_decrypt_batch(
&self,
_keys: &[&[u8]],
_nonces: &[&[u8]],
_data: &[&[u8]],
) -> Result<Vec<Vec<u8>>> {
Err(CryptoError::InvalidInput(
"Hash kernel does not support AES operation".into(),
))
}
}
#[cfg(test)]
mod tests {
#[cfg(feature = "gpu-cuda")]
mod cuda_tests {
use super::super::*;
#[test]
fn test_cuda_hash_kernel_creation() {
let kernel = CudaHashKernel::new();
assert!(!kernel.is_available());
}
#[test]
fn test_cuda_hash_kernel_initialize() {
let mut kernel = CudaHashKernel::new();
let result = kernel.initialize();
assert!(result.is_err());
}
#[test]
fn test_cuda_hash_kernel_shutdown() {
let mut kernel = CudaHashKernel::new();
let result = kernel.shutdown();
assert!(result.is_ok());
}
#[test]
fn test_cuda_hash_kernel_metrics() {
let kernel = CudaHashKernel::new();
let metrics = kernel.get_metrics();
assert!(metrics.is_none());
}
}
#[cfg(not(feature = "gpu-cuda"))]
mod cpu_tests {
#[test]
fn test_cuda_hash_kernel_creation() {
let kernel = CudaHashKernel;
assert!(!kernel.is_available());
}
#[test]
fn test_cuda_hash_kernel_initialize() {
let mut kernel = CudaHashKernel;
let result = kernel.initialize();
assert!(result.is_err());
}
}
}