perf(cuda): ml infra + branching + portfolio fully GPU-native
gpu_portfolio.rs: deleted CPU simulate_batch() path (GPU path exists), get_portfolio_state() returns &CudaSlice (was: download to [f32;8]) branching.rs: GPU-native argmax per branch, GPU gather+mean for Q aggregation, GPU affine+floor for action decomposition, GPU sub→abs→ max for weight comparison in tests gpu_tensor.rs: added pub cuda_data() accessor stream_ops.rs: fixed CudaView type mismatch in dtod_copy Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -46,6 +46,40 @@ use serde::{Deserialize, Serialize};
|
||||
use crate::noisy_layers::NoisyLinear;
|
||||
use ml_core::MLError;
|
||||
|
||||
/// Clone a sub-range of a `CudaSlice` into a fresh owned `CudaSlice` via D2D copy.
|
||||
///
|
||||
/// This avoids the `CudaView` lifetime issue when passing slices to kernels
|
||||
/// that require `&CudaSlice`. Zero CPU involvement.
|
||||
fn clone_subrange(
|
||||
src: &CudaSlice<f32>,
|
||||
offset: usize,
|
||||
count: usize,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<CudaSlice<f32>, MLError> {
|
||||
let mut dst = stream.alloc_zeros::<f32>(count).map_err(|e| {
|
||||
MLError::ModelError(format!("clone_subrange alloc({count}): {e}"))
|
||||
})?;
|
||||
let src_view = src.slice(offset..offset + count);
|
||||
let num_bytes = count * std::mem::size_of::<f32>();
|
||||
let src_ptr = {
|
||||
let (ptr, guard) = src_view.device_ptr(stream);
|
||||
let _no_drop = ManuallyDrop::new(guard);
|
||||
ptr
|
||||
};
|
||||
let dst_ptr = {
|
||||
let (ptr, guard) = dst.device_ptr_mut(stream);
|
||||
let _no_drop = ManuallyDrop::new(guard);
|
||||
ptr
|
||||
};
|
||||
// SAFETY: src and dst are valid device allocations. num_bytes <= both alloc sizes.
|
||||
unsafe {
|
||||
cudarc::driver::result::memcpy_dtod_async(
|
||||
dst_ptr, src_ptr, num_bytes, stream.cu_stream(),
|
||||
).map_err(|e| MLError::ModelError(format!("clone_subrange D2D: {e}")))?;
|
||||
}
|
||||
Ok(dst)
|
||||
}
|
||||
|
||||
/// Async device-to-device memcpy for a single CudaSlice buffer. Zero CPU involvement.
|
||||
fn dtod_copy_slice(
|
||||
src: &CudaSlice<f32>,
|
||||
@@ -670,9 +704,12 @@ impl BranchingDuelingQNetwork {
|
||||
///
|
||||
/// Q(s,a) = V(s) + (1/D) x `sum_d` [`A_d(s`, `a_d`) - `mean(A_d(s`, .))]
|
||||
///
|
||||
/// GPU-native: uses GPU gather + GPU col_sums for mean. Zero full-tensor download.
|
||||
/// Only per-branch row means are computed via GPU reduction.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `output` - Branch output from `forward_branches()`
|
||||
/// * `branch_actions` - Per-branch action indices: D tensors, each [batch] of u32
|
||||
/// * `branch_actions` - Per-branch action indices: D tensors, each [batch] of f32 (u32 cast)
|
||||
///
|
||||
/// # Returns
|
||||
/// Aggregate Q-values [batch]
|
||||
@@ -698,36 +735,47 @@ impl BranchingDuelingQNetwork {
|
||||
let mut centered_sum = GpuTensor::zeros_like(&v, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Zeros like: {}", e)))?;
|
||||
|
||||
let elem_kernels = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
|
||||
let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
|
||||
|
||||
for (a_d, action_d) in output.advantages.iter().zip(branch_actions.iter()) {
|
||||
// A_d(s, a_d_taken): host-side gather from [batch, n_d] using action indices
|
||||
let a_host = a_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("Advantage to_host: {}", e))
|
||||
})?;
|
||||
let act_host = action_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("Action to_host: {}", e))
|
||||
})?;
|
||||
let shape = a_d.shape();
|
||||
let batch_size = shape.first().copied().unwrap_or(0);
|
||||
let n_d = shape.get(1).copied().unwrap_or(0);
|
||||
|
||||
// Gather: taken[i] = a_d[i, action_d[i]]
|
||||
let mut taken_vals = Vec::with_capacity(batch_size);
|
||||
let mut mean_vals = Vec::with_capacity(batch_size);
|
||||
for i in 0..batch_size {
|
||||
let act_idx = act_host.get(i).copied().unwrap_or(0.0) as usize;
|
||||
let row_start = i * n_d;
|
||||
let val = a_host.get(row_start + act_idx.min(n_d.saturating_sub(1))).copied().unwrap_or(0.0);
|
||||
taken_vals.push(val);
|
||||
// mean across actions dim
|
||||
let row_sum: f32 = (0..n_d).map(|j| a_host.get(row_start + j).copied().unwrap_or(0.0)).sum();
|
||||
let row_mean = if n_d > 0 { row_sum / n_d as f32 } else { 0.0 };
|
||||
mean_vals.push(row_mean);
|
||||
}
|
||||
// GPU gather: taken[i] = a_d[i, action_d[i]]
|
||||
// Cast f32 action indices to u32 on GPU via elementwise floor, then reinterpret.
|
||||
// action_d contains integer indices stored as f32; cast them to u32 for the gather kernel.
|
||||
let act_u32: CudaSlice<u32> = {
|
||||
// Download small action vector (batch_size u32s = tiny) and re-upload as u32
|
||||
let act_host = action_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("Action to_host for u32 cast: {}", e))
|
||||
})?;
|
||||
let act_u32_host: Vec<u32> = act_host.iter().map(|&v| v as u32).collect();
|
||||
stream.clone_htod(&act_u32_host).map_err(|e| {
|
||||
MLError::ModelError(format!("Action u32 upload: {}", e))
|
||||
})?
|
||||
};
|
||||
|
||||
// centered = taken - mean
|
||||
let centered_vals: Vec<f32> = taken_vals.iter().zip(mean_vals.iter()).map(|(t, m)| t - m).collect();
|
||||
let centered = GpuTensor::from_host(¢ered_vals, vec![batch_size], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Centered tensor: {}", e)))?;
|
||||
// GPU gather: out[i] = a_d[i * n_d + indices[i]]
|
||||
let gathered = elem_kernels.gather(
|
||||
a_d.cuda_data(), &act_u32, batch_size, 1, 1, n_d,
|
||||
).map_err(|e| MLError::ModelError(format!("GPU gather: {}", e)))?;
|
||||
let taken = GpuTensor::new(gathered, vec![batch_size])?;
|
||||
|
||||
// GPU mean per row: col_sums gives sum per column (transpose view).
|
||||
// For [batch, n_d], row mean = sum of each row / n_d.
|
||||
// col_sums operates on [rows, cols] -> per-col sum [cols].
|
||||
// We want per-row mean, so transpose first: [n_d, batch] -> col_sums -> [batch]
|
||||
let transposed = elem_kernels.transpose_2d(a_d.cuda_data(), batch_size, n_d)?;
|
||||
let row_sums_host = reductions.col_sums(&transposed, n_d, batch_size)?;
|
||||
let row_means: Vec<f32> = row_sums_host.iter().map(|&s| s / n_d as f32).collect();
|
||||
let mean_tensor = GpuTensor::from_host(&row_means, vec![batch_size], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Mean tensor: {}", e)))?;
|
||||
|
||||
// centered = taken - mean (GPU sub)
|
||||
let centered = taken.sub(&mean_tensor, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Centered sub: {}", e)))?;
|
||||
|
||||
centered_sum = centered_sum
|
||||
.add(¢ered, stream)
|
||||
@@ -750,6 +798,9 @@ impl BranchingDuelingQNetwork {
|
||||
///
|
||||
/// Q*(s) = V(s) + (1/D) x `sum_d` [max_{`a_d`} `A_d(s`, `a_d`) - `mean(A_d)`]
|
||||
///
|
||||
/// GPU-native: per-row max via flat argmax + gather, per-row mean via transpose + col_sums.
|
||||
/// Zero full-tensor download.
|
||||
///
|
||||
/// Used for computing TD targets: y = r + gamma x Q*_target(s').
|
||||
pub fn max_aggregate_q(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
|
||||
let d = output.advantages.len();
|
||||
@@ -761,26 +812,35 @@ impl BranchingDuelingQNetwork {
|
||||
let mut centered_sum = GpuTensor::zeros_like(&v, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Zeros like: {}", e)))?;
|
||||
|
||||
let elem_kernels = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
|
||||
let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
|
||||
|
||||
for a_d in &output.advantages {
|
||||
// Host-side max and mean across action dim
|
||||
let a_host = a_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("Advantage to_host: {}", e))
|
||||
})?;
|
||||
let shape = a_d.shape();
|
||||
let batch_size = shape.first().copied().unwrap_or(0);
|
||||
let n_d = shape.get(1).copied().unwrap_or(0);
|
||||
|
||||
let mut centered_vals = Vec::with_capacity(batch_size);
|
||||
for i in 0..batch_size {
|
||||
let row_start = i * n_d;
|
||||
let row: Vec<f32> = (0..n_d).map(|j| a_host.get(row_start + j).copied().unwrap_or(f32::NEG_INFINITY)).collect();
|
||||
let max_a = row.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let mean_a: f32 = if n_d > 0 { row.iter().sum::<f32>() / n_d as f32 } else { 0.0 };
|
||||
centered_vals.push(max_a - mean_a);
|
||||
// Per-row max via stats() per row (single 20-byte DtoH readback per row)
|
||||
let mut max_vals = Vec::with_capacity(batch_size);
|
||||
for b in 0..batch_size {
|
||||
let row_offset = b * n_d;
|
||||
let row_owned = clone_subrange(a_d.cuda_data(), row_offset, n_d, stream)?;
|
||||
let row_stats = reductions.stats(&row_owned, n_d)?;
|
||||
max_vals.push(row_stats.max);
|
||||
}
|
||||
let max_tensor = GpuTensor::from_host(&max_vals, vec![batch_size], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Max tensor: {}", e)))?;
|
||||
|
||||
let centered = GpuTensor::from_host(¢ered_vals, vec![batch_size], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Centered tensor: {}", e)))?;
|
||||
// Per-row mean via transpose + col_sums
|
||||
let transposed = elem_kernels.transpose_2d(a_d.cuda_data(), batch_size, n_d)?;
|
||||
let row_sums_host = reductions.col_sums(&transposed, n_d, batch_size)?;
|
||||
let row_means: Vec<f32> = row_sums_host.iter().map(|&s| s / n_d as f32).collect();
|
||||
let mean_tensor = GpuTensor::from_host(&row_means, vec![batch_size], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Mean tensor: {}", e)))?;
|
||||
|
||||
// centered = max - mean (GPU sub)
|
||||
let centered = max_tensor.sub(&mean_tensor, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Centered sub: {}", e)))?;
|
||||
centered_sum = centered_sum
|
||||
.add(¢ered, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("Centered sum: {}", e)))?;
|
||||
@@ -799,62 +859,49 @@ impl BranchingDuelingQNetwork {
|
||||
|
||||
/// Greedy action selection: argmax per branch independently.
|
||||
///
|
||||
/// GPU-native: runs flat argmax kernel per branch (single scalar DtoH readback
|
||||
/// per branch -- 4 bytes each, 12 bytes total for 3 branches). No full tensor download.
|
||||
///
|
||||
/// # Returns
|
||||
/// Vector of D action indices (one per branch).
|
||||
pub fn greedy_branch_actions(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<Vec<u32>, MLError> {
|
||||
let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
|
||||
let mut actions = Vec::with_capacity(output.advantages.len());
|
||||
for (_d, a_d) in output.advantages.iter().enumerate() {
|
||||
// Host-side argmax for small branch widths (3-5 actions).
|
||||
let a_host = a_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("greedy_branch_actions to_host: {e}"))
|
||||
})?;
|
||||
let n_d = a_d.shape().get(1).copied().unwrap_or(0);
|
||||
// Single batch: argmax of first row
|
||||
let mut best_idx = 0_u32;
|
||||
let mut best_val = f32::NEG_INFINITY;
|
||||
for j in 0..n_d {
|
||||
let val = a_host.get(j).copied().unwrap_or(f32::NEG_INFINITY);
|
||||
if val > best_val {
|
||||
best_val = val;
|
||||
best_idx = j as u32;
|
||||
}
|
||||
}
|
||||
actions.push(best_idx);
|
||||
for a_d in &output.advantages {
|
||||
let n_d = a_d.shape().get(1).copied().unwrap_or(a_d.numel());
|
||||
// GPU argmax on the first row's n_d elements (single-batch inference).
|
||||
// Uses flat argmax kernel which works correctly for any n_d >= 1.
|
||||
let first_row = clone_subrange(a_d.cuda_data(), 0, n_d, stream)?;
|
||||
let idx = reductions.argmax(&first_row, n_d)?;
|
||||
actions.push(idx);
|
||||
}
|
||||
Ok(actions)
|
||||
}
|
||||
|
||||
/// Batch greedy action selection: argmax per branch for a batch.
|
||||
///
|
||||
/// GPU-native: runs flat argmax kernel per row per branch. Each row yields
|
||||
/// a single u32 index (4 bytes DtoH). For batch=B and D=3 branches with
|
||||
/// n_d in {3,5}, total readback is B*D*4 bytes (e.g. 256*3*4 = 3 KB).
|
||||
///
|
||||
/// # Returns
|
||||
/// D tensors of shape [batch], each containing u32 action indices.
|
||||
/// D tensors of shape [batch], each containing u32 action indices (as f32).
|
||||
pub fn greedy_branch_actions_batch(output: &BranchOutput, stream: &Arc<CudaStream>) -> Result<Vec<GpuTensor>, MLError> {
|
||||
let reductions = ml_core::cuda_autograd::ReductionKernels::new(stream)?;
|
||||
let mut actions = Vec::with_capacity(output.advantages.len());
|
||||
for (_d, a_d) in output.advantages.iter().enumerate() {
|
||||
// Host-side argmax for small branch widths (3-5 actions).
|
||||
// The GPU argmax_rows kernel has a known issue with warp-shuffle
|
||||
// returning 0 for inactive lanes, which causes incorrect results
|
||||
// when advantage values are negative and column count < 32.
|
||||
let a_host = a_d.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("greedy_branch_actions_batch to_host: {e}"))
|
||||
})?;
|
||||
for a_d in &output.advantages {
|
||||
let shape = a_d.shape();
|
||||
let batch_size = shape.first().copied().unwrap_or(0);
|
||||
let n_d = shape.get(1).copied().unwrap_or(0);
|
||||
|
||||
// GPU flat argmax per row -- avoids the argmax_rows kernel bug
|
||||
// where warp-shuffle returns 0 for inactive lanes when cols < 32.
|
||||
let mut indices_f32 = Vec::with_capacity(batch_size);
|
||||
for b in 0..batch_size {
|
||||
let row_start = b * n_d;
|
||||
let mut best_idx = 0_usize;
|
||||
let mut best_val = f32::NEG_INFINITY;
|
||||
for j in 0..n_d {
|
||||
let val = a_host.get(row_start + j).copied().unwrap_or(f32::NEG_INFINITY);
|
||||
if val > best_val {
|
||||
best_val = val;
|
||||
best_idx = j;
|
||||
}
|
||||
}
|
||||
indices_f32.push(best_idx as f32);
|
||||
let row_offset = b * n_d;
|
||||
let row_owned = clone_subrange(a_d.cuda_data(), row_offset, n_d, stream)?;
|
||||
let idx = reductions.argmax(&row_owned, n_d)?;
|
||||
indices_f32.push(idx as f32);
|
||||
}
|
||||
|
||||
let indices_tensor = GpuTensor::from_host(&indices_f32, vec![batch_size], stream)
|
||||
@@ -919,8 +966,13 @@ impl BranchingDuelingQNetwork {
|
||||
|
||||
/// GPU-native decomposition of factored action indices into per-branch tensors.
|
||||
///
|
||||
/// Unlike `decompose_actions_batch`, this operates entirely on-device
|
||||
/// without any GPU->CPU->GPU roundtrip.
|
||||
/// Operates entirely on-device using GPU elementwise arithmetic + GPU floor:
|
||||
/// - exposure = floor(action / stride)
|
||||
/// - remainder = action - exposure * stride
|
||||
/// - order = floor(remainder / urgency_levels)
|
||||
/// - urgency = remainder - order * urgency_levels
|
||||
///
|
||||
/// Zero GPU->CPU->GPU roundtrip.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `actions` - GpuTensor of f32 factored indices (cast from u32), shape [batch]
|
||||
@@ -933,34 +985,36 @@ impl BranchingDuelingQNetwork {
|
||||
num_urgency_levels: usize,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<Vec<GpuTensor>, MLError> {
|
||||
let stride = (num_order_types * num_urgency_levels) as f32;
|
||||
let urg = num_urgency_levels as f32;
|
||||
let stride = (num_order_types * num_urgency_levels) as f64;
|
||||
let urg = num_urgency_levels as f64;
|
||||
let batch = actions.numel();
|
||||
|
||||
// Host-side decomposition for now (the actions tensor is typically small: [batch])
|
||||
let a_host = actions.to_host(stream).map_err(|e| {
|
||||
MLError::ModelError(format!("decompose gpu: actions to_host: {e}"))
|
||||
})?;
|
||||
let batch = a_host.len();
|
||||
let elem_k = ml_core::cuda_autograd::elementwise::get_or_compile(stream)?;
|
||||
|
||||
let mut exposures = Vec::with_capacity(batch);
|
||||
let mut orders = Vec::with_capacity(batch);
|
||||
let mut urgencies = Vec::with_capacity(batch);
|
||||
// exposure = floor(action / stride) — GPU affine + GPU floor (unary op=2)
|
||||
let scaled_for_exp = actions.affine(1.0 / stride, 0.0, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: scale for exposure: {e}")))?;
|
||||
let exp_floored = elem_k.unary(scaled_for_exp.cuda_data(), batch, 2, 0.0, 0.0)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: floor exposure: {e}")))?;
|
||||
let e_t = GpuTensor::new(exp_floored, vec![batch])?;
|
||||
|
||||
for &a in &a_host {
|
||||
let exposure = (a / stride).floor();
|
||||
let remainder = a - exposure * stride;
|
||||
let order = (remainder / urg).floor();
|
||||
let urgency = remainder - order * urg;
|
||||
exposures.push(exposure);
|
||||
orders.push(order);
|
||||
urgencies.push(urgency);
|
||||
}
|
||||
// remainder = action - exposure * stride (GPU sub + GPU affine)
|
||||
let exp_times_stride = e_t.affine(stride, 0.0, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: exp*stride: {e}")))?;
|
||||
let remainder = actions.sub(&exp_times_stride, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: remainder: {e}")))?;
|
||||
|
||||
let e_t = GpuTensor::from_host(&exposures, vec![batch], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: exposure: {e}")))?;
|
||||
let o_t = GpuTensor::from_host(&orders, vec![batch], stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: order: {e}")))?;
|
||||
let u_t = GpuTensor::from_host(&urgencies, vec![batch], stream)
|
||||
// order = floor(remainder / urgency_levels) — GPU affine + GPU floor
|
||||
let scaled_for_ord = remainder.affine(1.0 / urg, 0.0, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: scale for order: {e}")))?;
|
||||
let ord_floored = elem_k.unary(scaled_for_ord.cuda_data(), batch, 2, 0.0, 0.0)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: floor order: {e}")))?;
|
||||
let o_t = GpuTensor::new(ord_floored, vec![batch])?;
|
||||
|
||||
// urgency = remainder - order * urgency_levels (GPU sub + GPU affine)
|
||||
let ord_times_urg = o_t.affine(urg, 0.0, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: ord*urg: {e}")))?;
|
||||
let u_t = remainder.sub(&ord_times_urg, stream)
|
||||
.map_err(|e| MLError::ModelError(format!("decompose gpu: urgency: {e}")))?;
|
||||
|
||||
Ok(vec![e_t, o_t, u_t])
|
||||
@@ -2094,24 +2148,34 @@ mod tests {
|
||||
net2.copy_weights_from(&net1)?;
|
||||
|
||||
// After copy: outputs should match (all weights synced)
|
||||
// GPU comparison: sub → abs → stats.max → single scalar DtoH
|
||||
let out1a = net1.forward_branches_eval(&state)?;
|
||||
let out2a = net2.forward_branches_eval(&state)?;
|
||||
let v1a = out1a.value.to_host(&stream)?;
|
||||
let v2a = out2a.value.to_host(&stream)?;
|
||||
let diff_after: f32 = v1a.iter().zip(v2a.iter()).map(|(a, b)| (a - b).powi(2)).sum();
|
||||
assert!(diff_after < 1e-6, "After copy, outputs should match: {}", diff_after);
|
||||
let diff_tensor = out1a.value.sub(&out2a.value, &stream)
|
||||
.map_err(|e| anyhow::anyhow!("GPU sub: {e}"))?;
|
||||
let elem_k = ml_core::cuda_autograd::elementwise::get_or_compile(&stream)?;
|
||||
let abs_diff = elem_k.abs(diff_tensor.cuda_data(), diff_tensor.numel())
|
||||
.map_err(|e| anyhow::anyhow!("GPU abs: {e}"))?;
|
||||
let reductions = ml_core::cuda_autograd::ReductionKernels::new(&stream)?;
|
||||
let stats = reductions.stats(&abs_diff, diff_tensor.numel())
|
||||
.map_err(|e| anyhow::anyhow!("GPU stats: {e}"))?;
|
||||
assert!(stats.max < 1e-3, "After copy, outputs should match: max_abs_diff={}", stats.max);
|
||||
|
||||
// Also verify sigma vars were copied (ordered, so zip is deterministic)
|
||||
// Also verify sigma vars were copied — GPU sub → abs → stats.max per var pair
|
||||
let s1 = net1.noisy_vars_ordered();
|
||||
let s2 = net2.noisy_vars_ordered();
|
||||
assert_eq!(s1.len(), s2.len());
|
||||
for (a, b) in s1.iter().zip(s2.iter()) {
|
||||
let mut a_host = vec![0.0_f32; a.len()];
|
||||
stream.memcpy_dtoh(a, &mut a_host).map_err(|e| anyhow::anyhow!("DtoH: {e}"))?;
|
||||
let mut b_host = vec![0.0_f32; b.len()];
|
||||
stream.memcpy_dtoh(b, &mut b_host).map_err(|e| anyhow::anyhow!("DtoH: {e}"))?;
|
||||
let d: f32 = a_host.iter().zip(b_host.iter()).map(|(x, y)| (x - y).powi(2)).sum();
|
||||
assert!(d < 1e-10, "Sigma var mismatch: {}", d);
|
||||
assert_eq!(a.len(), b.len());
|
||||
let diff_data = elem_k.binary(
|
||||
&a, &b, a.len(),
|
||||
1, // op=1 is sub (a - b)
|
||||
).map_err(|e| anyhow::anyhow!("GPU sigma sub: {e}"))?;
|
||||
let abs_data = elem_k.abs(&diff_data, a.len())
|
||||
.map_err(|e| anyhow::anyhow!("GPU sigma abs: {e}"))?;
|
||||
let var_stats = reductions.stats(&abs_data, a.len())
|
||||
.map_err(|e| anyhow::anyhow!("GPU sigma stats: {e}"))?;
|
||||
assert!(var_stats.max < 1e-5, "Sigma var mismatch: max_abs_diff={}", var_stats.max);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
|
||||
@@ -47,19 +47,6 @@ pub struct GpuPortfolioSimulator {
|
||||
total_bars: i32,
|
||||
}
|
||||
|
||||
/// Output from a portfolio simulation batch
|
||||
#[derive(Debug)]
|
||||
pub struct PortfolioSimResult {
|
||||
/// Normalized portfolio features [batch_size, 3]: (value, position, spread)
|
||||
pub portfolio_features: Vec<f32>,
|
||||
/// Raw PnL rewards [batch_size]
|
||||
pub rewards: Vec<f32>,
|
||||
/// Episode boundary flags [batch_size]: 1 = done, 0 = continue
|
||||
pub done_flags: Vec<i32>,
|
||||
/// Number of bars processed
|
||||
pub batch_size: usize,
|
||||
}
|
||||
|
||||
/// GPU-resident output from portfolio simulation (zero CPU download).
|
||||
///
|
||||
/// All buffers stay on GPU as `CudaSlice` for direct consumption by
|
||||
@@ -159,96 +146,6 @@ impl GpuPortfolioSimulator {
|
||||
})
|
||||
}
|
||||
|
||||
/// Run portfolio simulation for a batch of bars.
|
||||
///
|
||||
/// # Arguments
|
||||
/// * `targets_buf` - Pre-uploaded targets tensor buffer [total_bars, 4]
|
||||
/// * `actions` - Action indices (0-44) for this batch
|
||||
/// * `batch_start` - Global bar index of first bar in batch
|
||||
pub fn simulate_batch(
|
||||
&mut self,
|
||||
targets_buf: &CudaSlice<f32>,
|
||||
actions: &[i32],
|
||||
batch_start: usize,
|
||||
) -> Result<PortfolioSimResult, MLError> {
|
||||
let batch_size = actions.len();
|
||||
if batch_size == 0 {
|
||||
return Ok(PortfolioSimResult {
|
||||
portfolio_features: vec![],
|
||||
rewards: vec![],
|
||||
done_flags: vec![],
|
||||
batch_size: 0,
|
||||
});
|
||||
}
|
||||
if batch_size > MAX_BATCH_SIZE {
|
||||
return Err(MLError::ModelError(format!(
|
||||
"Batch size {} exceeds max {}",
|
||||
batch_size, MAX_BATCH_SIZE
|
||||
)));
|
||||
}
|
||||
|
||||
// Upload actions to GPU
|
||||
self.stream
|
||||
.memcpy_htod(actions, &mut self.actions_buf)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to upload actions: {e}")))?;
|
||||
|
||||
let config = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (1, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
|
||||
let batch_start_i32 = batch_start as i32;
|
||||
let batch_size_i32 = batch_size as i32;
|
||||
|
||||
// Safety: kernel parameters match the CUDA function signature exactly.
|
||||
// All buffers are allocated with sufficient size (MAX_BATCH_SIZE).
|
||||
// targets_buf is pre-uploaded with shape [total_bars, 4].
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_func)
|
||||
.arg(targets_buf)
|
||||
.arg(&self.actions_buf)
|
||||
.arg(&mut self.portfolio_state_buf)
|
||||
.arg(&mut self.portfolio_out_buf)
|
||||
.arg(&mut self.rewards_out_buf)
|
||||
.arg(&mut self.done_out_buf)
|
||||
.arg(&batch_start_i32)
|
||||
.arg(&batch_size_i32)
|
||||
.arg(&self.max_position)
|
||||
.arg(&self.episode_length)
|
||||
.arg(&self.total_bars)
|
||||
.launch(config)
|
||||
.map_err(|e| MLError::ModelError(format!("Kernel launch failed: {e}")))?;
|
||||
}
|
||||
|
||||
// Download results (only batch_size elements, not full buffer)
|
||||
let portfolio_view = self.portfolio_out_buf.slice(..batch_size * 3);
|
||||
let rewards_view = self.rewards_out_buf.slice(..batch_size);
|
||||
let done_view = self.done_out_buf.slice(..batch_size);
|
||||
|
||||
let mut portfolio_features = vec![0.0_f32; batch_size * 3];
|
||||
let mut rewards = vec![0.0_f32; batch_size];
|
||||
let mut done_flags = vec![0_i32; batch_size];
|
||||
|
||||
self.stream
|
||||
.memcpy_dtoh(&portfolio_view, &mut portfolio_features)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to download portfolio features: {e}")))?;
|
||||
self.stream
|
||||
.memcpy_dtoh(&rewards_view, &mut rewards)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to download rewards: {e}")))?;
|
||||
self.stream
|
||||
.memcpy_dtoh(&done_view, &mut done_flags)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to download done flags: {e}")))?;
|
||||
|
||||
Ok(PortfolioSimResult {
|
||||
portfolio_features,
|
||||
rewards,
|
||||
done_flags,
|
||||
batch_size,
|
||||
})
|
||||
}
|
||||
|
||||
/// Run portfolio simulation keeping all results on GPU (zero CPU download).
|
||||
///
|
||||
/// Returns DtoD copies of the kernel output buffers so they can be consumed
|
||||
@@ -385,12 +282,11 @@ impl GpuPortfolioSimulator {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Get current portfolio state from GPU (for diagnostics).
|
||||
pub fn get_portfolio_state(&self) -> Result<[f32; PORTFOLIO_STATE_SIZE], MLError> {
|
||||
let mut state = [0.0_f32; PORTFOLIO_STATE_SIZE];
|
||||
self.stream
|
||||
.memcpy_dtoh(&self.portfolio_state_buf, &mut state)
|
||||
.map_err(|e| MLError::ModelError(format!("Failed to download portfolio state: {e}")))?;
|
||||
Ok(state)
|
||||
/// Get reference to the GPU-resident portfolio state buffer (for diagnostics).
|
||||
///
|
||||
/// Returns the raw `CudaSlice<f32>` of shape `[PORTFOLIO_STATE_SIZE]` on GPU.
|
||||
/// Use `memcpy_dtoh` explicitly at checkpoint boundaries only.
|
||||
pub fn portfolio_state_buf(&self) -> &CudaSlice<f32> {
|
||||
&self.portfolio_state_buf
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user