rustytorch/.claude/plans/quizzical-swimming-ullman.md

CUDA Implementation Gaps Deep Fix Plan

Executive Summary

Two critical CUDA issues identified and ready to fix:

1. CUDA Data Transfer Bug - cuda_matmul() returns zeros because it writes to cloned handles instead of actual storage
2. cuDNN Module Rewrite - 44 private .desc field accesses need migration to cudarc 0.18.x safe API

---

Issue 1: CUDA Data Transfer Bug (CRITICAL - Quick Fix)

Root Cause

The cuda_matmul() function uses cuda_slice_clone() which only clones the CudaSlice handle pointer, NOT the underlying GPU memory. cuBLAS writes to this disconnected handle, and the result is never reflected in the Storage.

Broken pattern in cuda_matmul() (lines 587-594):

let a_slice = self.storage.cuda_slice_clone()?;  // Handle clone only
let b_slice = other.storage.cuda_slice_clone()?;
let mut c_slice = Arc::get_mut(&mut result_storage)?.cuda_slice_clone()?;
cublas.gemm(config, &b_slice, &a_slice, &mut c_slice)?;  // Writes to disconnected handle!

Working pattern in cuda_matmul_out() (lines 308-320):

let a_guard = self.storage.lock_cuda_slice()?;
let b_guard = other.storage.lock_cuda_slice()?;
let mut out_guard = out.storage.lock_cuda_slice()?;
let a_slice = a_guard.cuda_slice()?;
let b_slice = b_guard.cuda_slice()?;
let c_slice = out_guard.cuda_slice_mut()?;  // Direct mutable reference!
cublas.gemm(config, b_slice, a_slice, c_slice)?;  // Writes directly to storage

Fix

File: crates/core/rtx-tensor/src/tensor/matrix_multiplication.rs

Changes (lines 474-618):

1. Replace cuda_slice_clone() with lock_cuda_slice() pattern
2. Use cuda_slice() and cuda_slice_mut() for direct storage access
3. Add stream synchronization after cuBLAS GEMM

fn cuda_matmul(&self, other: &Self, device_id: usize) -> Result<Self> {
    // ... existing setup code ...

    // Create result tensor
    let result_shape = self.shape.matmul_shape(&other.shape)?;
    let mut result = Self::zeros_like_shape(&result_shape, &self.device)?;

    // Use lock pattern for direct storage access
    let a_guard = self.storage.lock_cuda_slice()?;
    let b_guard = other.storage.lock_cuda_slice()?;
    let mut result_guard = result.storage.lock_cuda_slice()?;

    let a_slice = a_guard.cuda_slice()?;
    let b_slice = b_guard.cuda_slice()?;
    let c_slice = result_guard.cuda_slice_mut()?;

    // Execute cuBLAS GEMM - writes directly to result storage
    unsafe {
        cublas.gemm(config, b_slice, a_slice, c_slice)?;
    }

    // Drop guards before returning
    drop(result_guard);
    drop(b_guard);
    drop(a_guard);

    Ok(result)
}

Also Fix

File: crates/core/rtx-tensor/src/storage/core.rs (line 1764)

Update misleading comment:

/// Get cloned CUDA slice handle (NOT a deep copy!)
/// WARNING: This only clones the slice handle pointer, not GPU memory.
/// Use lock_cuda_slice() + cuda_slice_mut() for mutable access.

---

Issue 2: cuDNN Module Rewrite (HIGH - Larger Refactor)

Current Problem

The cuDNN module accesses private .desc fields (44 occurrences across 7 files):

• descriptor.desc - direct private field access
• self.descriptor.desc - same issue

cudarc 0.18.x safe API provides factory methods on Cudnn handle instead.

cudarc 0.18.x Safe API

Old pattern (broken):

let descriptor = TensorDescriptor::new()?;
unsafe {
    cudnn_sys::cudnnSetTensor4dDescriptor(descriptor.desc, ...);  // PRIVATE FIELD
}

New pattern (cudarc 0.18.x):

// Use Cudnn handle factory methods
let descriptor = cudnn.create_4d_tensor::<f32>(format, [n, c, h, w])?;

// Use operation structs for convolution
let conv_op = ConvForward { conv: &conv_desc, x: &x_desc, w: &w_desc, y: &y_desc };
let algo = conv_op.pick_algorithm()?;
let workspace_size = conv_op.get_workspace_size(algo)?;
unsafe { conv_op.launch(algo, workspace, (alpha, beta), x, w, y)?; }

Files to Modify

File	Changes	Impact
rtx-tensor/src/cudnn/mod.rs	Store Arc<Cudnn> handle, expose factory methods	LOW
rtx-tensor/src/cudnn/descriptors.rs	Remove Safe wrappers, use cudarc types directly	HIGH
rtx-tensor/src/cudnn/convolution.rs	Use ConvForward struct	HIGH
rtx-tensor/src/cudnn/fused_ops.rs	Use ConvBiasActivationForward struct	MEDIUM
rtx-tensor/src/cudnn/error.rs	Update for cudarc error types	LOW
rtx-tensor/src/lib.rs	Re-enable cudnn module export	LOW

Implementation Phases

Phase 2A: Update CudnnContext (mod.rs)

pub struct CudnnContext {
    cudnn: Arc<Cudnn>,  // Store cudarc handle
    device_id: usize,
}

impl CudnnContext {
    pub fn new(device_id: usize) -> Result<Self> {
        let ctx = get_or_create_context(device_id)?;
        let stream = ctx.default_stream();
        let cudnn = Cudnn::new(stream)?;
        Ok(Self { cudnn: Arc::new(cudnn), device_id })
    }

    pub fn cudnn(&self) -> &Arc<Cudnn> { &self.cudnn }
}

Phase 2B: Simplify Descriptors (descriptors.rs)

// Remove SafeTensorDescriptor, SafeFilterDescriptor, etc.
// Use cudarc types directly

pub fn create_tensor_4d(
    cudnn: &Arc<Cudnn>,
    format: cudnnTensorFormat_t,
    dims: [i32; 4],
) -> Result<TensorDescriptor<f32>> {
    cudnn.create_4d_tensor(format, dims)
        .map_err(|e| CudnnError::from(e))
}

Phase 2C: Update Convolution (convolution.rs)

pub fn conv2d_forward(
    ctx: &CudnnContext,
    input: &TensorDescriptor<f32>,
    filter: &FilterDescriptor<f32>,
    conv: &ConvDescriptor<f32>,
    output: &TensorDescriptor<f32>,
    input_data: &CudaSlice<f32>,
    filter_data: &CudaSlice<f32>,
    output_data: &mut CudaSlice<f32>,
) -> Result<()> {
    let op = ConvForward { conv, x: input, w: filter, y: output };
    let algo = op.pick_algorithm()?;
    let workspace_size = op.get_workspace_size(algo)?;

    // Allocate workspace if needed
    let workspace = if workspace_size > 0 {
        Some(ctx.allocate_workspace(workspace_size)?)
    } else {
        None
    };

    unsafe {
        op.launch(algo, workspace.as_deref(), (1.0f32, 0.0f32),
                  input_data, filter_data, output_data)?;
    }
    Ok(())
}

---

Implementation Order

Priority 1: CUDA Data Transfer Fix (30 min)

1. Fix cuda_matmul() to use lock pattern
2. Fix cuda_matmul() FP16 path similarly
3. Update misleading cuda_slice_clone() comment
4. Run cargo test -p rtx-tensor --features cuda cublas - should pass

Priority 2: cuDNN Module Rewrite (4-6 hours)

1. Update CudnnContext to store Arc<Cudnn>
2. Simplify descriptors.rs to use factory methods
3. Update convolution.rs to use ConvForward struct
4. Update fused_ops.rs to use ConvBiasActivationForward
5. Re-enable cudnn module in lib.rs
6. Run cargo test -p rtx-tensor --features cuda cudnn

---

Test Commands

# After CUDA data transfer fix
cargo test -p rtx-tensor --features cuda cublas

# After cuDNN rewrite
cargo test -p rtx-tensor --features cuda cudnn
cargo test -p rtx-kernel --features cuda conv

# Full verification
cargo check -p rtx-tensor --features cuda
cargo check -p rtx-kernel --features cuda

---

Critical Files

Priority	File	Line	Issue
P1	rtx-tensor/src/tensor/matrix_multiplication.rs	587-594	cuda_slice_clone() bug
P1	rtx-tensor/src/storage/core.rs	1764	Misleading comment
P2	rtx-tensor/src/cudnn/descriptors.rs	107, 197, 241, 312, 362, 376, 427	Private .desc access
P2	rtx-tensor/src/cudnn/convolution.rs	243-252, 376-381, 423-426, 483-486	Private .desc access
P2	rtx-tensor/src/cudnn/fused_ops.rs	191-206, 266-269	Private .desc access
P2	rtx-tensor/src/cudnn/mod.rs	-	CudnnContext update
P2	rtx-tensor/src/lib.rs	42-43	Re-enable cudnn export