ringkernel_wavesim/simulation/

kernels.rs

//! CUDA kernel definitions using the Rust DSL.
//!
//! This module contains all CUDA kernels for the wave simulation, defined in
//! a Rust DSL that gets transpiled to CUDA C at compile time.
//!
//! The generated CUDA code is designed to match the handwritten versions in
//! `shaders/fdtd_tile.cu` and `shaders/fdtd_packed.cu` exactly.

#[cfg(feature = "cuda-codegen")]
use ringkernel_cuda_codegen::{
    transpile_global_kernel, transpile_stencil_kernel, Grid, StencilConfig,
};

// ============================================================================
// Tile-Based Kernels (fdtd_tile.cu equivalent)
// ============================================================================

/// Generate the complete CUDA source for tile-based kernels.
///
/// This generates CUDA code equivalent to `shaders/fdtd_tile.cu`:
/// - `fdtd_tile_step`: Main FDTD wave equation kernel
/// - `extract_halo`: Extract halo from interior edge
/// - `inject_halo`: Inject halo to boundary region
/// - `read_interior`: Read interior cells to output buffer
/// - `apply_boundary_reflection`: Apply boundary conditions
#[cfg(feature = "cuda-codegen")]
pub fn generate_tile_kernels() -> String {
    let mut output = String::new();

    output.push_str(TILE_KERNELS_HEADER);
    output.push_str("\nextern \"C\" {\n\n");

    // Generate fdtd_tile_step kernel
    output.push_str(&generate_fdtd_tile_step());
    output.push('\n');

    // Generate extract_halo kernel
    output.push_str(&generate_extract_halo());
    output.push('\n');

    // Generate inject_halo kernel
    output.push_str(&generate_inject_halo());
    output.push('\n');

    // Generate read_interior kernel
    output.push_str(&generate_read_interior());
    output.push('\n');

    // Generate apply_boundary_reflection kernel
    output.push_str(&generate_apply_boundary_reflection());

    output.push_str("\n}  // extern \"C\"\n");

    output
}

/// Header comment for generated tile kernels.
pub const TILE_KERNELS_HEADER: &str = r#"// CUDA Kernels for Tile-Based FDTD Wave Simulation
// Generated by ringkernel-cuda-codegen from Rust DSL
//
// Buffer Layout (18x18 = 324 floats):
//   +---+----------------+---+
//   | NW|   North Halo   |NE |  <- Row 0
//   +---+----------------+---+
//   |   |   16x16 Tile   |   |  <- Rows 1-16
//   | W |    Interior    | E |
//   +---+----------------+---+
//   | SW|   South Halo   |SE |  <- Row 17
//   +---+----------------+---+
//
// Index: idx = y * 18 + x
// Interior cell (lx, ly): idx = (ly + 1) * 18 + (lx + 1)
"#;

/// Generate the main FDTD tile step kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_fdtd_tile_step() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn fdtd_tile_step(
            pressure: &[f32],
            pressure_prev: &mut [f32],
            c2: f32,
            damping: f32,
            pos: GridPos,
        ) {
            let p = pressure[pos.idx()];
            let p_prev = pressure_prev[pos.idx()];

            let p_n = pos.north(pressure);
            let p_s = pos.south(pressure);
            let p_w = pos.west(pressure);
            let p_e = pos.east(pressure);

            let laplacian = p_n + p_s + p_e + p_w - 4.0 * p;
            let p_new = 2.0 * p - p_prev + c2 * laplacian;

            pressure_prev[pos.idx()] = p_new * damping;
        }
    };

    let config = StencilConfig::new("fdtd_tile_step")
        .with_grid(Grid::Grid2D)
        .with_tile_size(16, 16)
        .with_halo(1);

    match transpile_stencil_kernel(&kernel_fn, &config) {
        Ok(cuda) => cuda,
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate the extract_halo kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_extract_halo() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn extract_halo(
            pressure: &[f32],
            halo_out: &mut [f32],
            edge: i32,
        ) {
            let i = thread_idx_x();
            if i >= 16 {
                return;
            }

            let idx = match edge {
                0 => 1 * 18 + (i + 1),       // North - extract row 1
                1 => 16 * 18 + (i + 1),      // South - extract row 16
                2 => (i + 1) * 18 + 1,       // West - extract col 1
                _ => (i + 1) * 18 + 16,      // East - extract col 16
            };

            halo_out[i as usize] = pressure[idx as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!("// Extract Halo - extracts halo data from interior edge\n// edge: 0=North, 1=South, 2=West, 3=East\n{}", cuda),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate the inject_halo kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_inject_halo() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn inject_halo(
            pressure: &mut [f32],
            halo_in: &[f32],
            edge: i32,
        ) {
            let i = thread_idx_x();
            if i >= 16 {
                return;
            }

            let idx = match edge {
                0 => 0 * 18 + (i + 1),       // North - inject to row 0
                1 => 17 * 18 + (i + 1),      // South - inject to row 17
                2 => (i + 1) * 18 + 0,       // West - inject to col 0
                _ => (i + 1) * 18 + 17,      // East - inject to col 17
            };

            pressure[idx as usize] = halo_in[i as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!("// Inject Halo - injects halo data from linear buffer to halo region\n// edge: 0=North, 1=South, 2=West, 3=East\n{}", cuda),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate the read_interior kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_read_interior() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn read_interior(
            pressure: &[f32],
            output: &mut [f32],
        ) {
            let lx = thread_idx_x();
            let ly = thread_idx_y();

            if lx >= 16 || ly >= 16 {
                return;
            }

            let src_idx = (ly + 1) * 18 + (lx + 1);
            let dst_idx = ly * 16 + lx;

            output[dst_idx as usize] = pressure[src_idx as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Read Interior - reads interior pressure to linear buffer for visualization\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate the apply_boundary_reflection kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_apply_boundary_reflection() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn apply_boundary_reflection(
            pressure: &mut [f32],
            edge: i32,
            reflection_coeff: f32,
        ) {
            let i = thread_idx_x();
            if i >= 16 {
                return;
            }

            // Calculate source and destination indices based on edge
            let src_idx = match edge {
                0 => 1 * 18 + (i + 1),       // North - reflect row 1
                1 => 16 * 18 + (i + 1),      // South - reflect row 16
                2 => (i + 1) * 18 + 1,       // West - reflect col 1
                _ => (i + 1) * 18 + 16,      // East - reflect col 16
            };

            let dst_idx = match edge {
                0 => 0 * 18 + (i + 1),       // North - to row 0
                1 => 17 * 18 + (i + 1),      // South - to row 17
                2 => (i + 1) * 18 + 0,       // West - to col 0
                _ => (i + 1) * 18 + 17,      // East - to col 17
            };

            pressure[dst_idx as usize] = pressure[src_idx as usize] * reflection_coeff;
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!("// Apply Boundary Reflection - applies boundary conditions for tiles at grid edges\n// edge: 0=North, 1=South, 2=West, 3=East\n{}", cuda),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

// ============================================================================
// Packed Tile Kernels (fdtd_packed.cu equivalent)
// ============================================================================

/// Generate the complete CUDA source for packed tile kernels.
///
/// This generates CUDA code equivalent to `shaders/fdtd_packed.cu`:
/// - `exchange_all_halos`: Copy halos between adjacent tiles
/// - `fdtd_all_tiles`: Batched FDTD for all tiles in parallel
/// - `upload_tile_data`: Upload initial state to a tile
/// - `read_all_interiors`: Read all tile interiors to output
/// - `inject_impulse`: Add impulse to specific cell
/// - `apply_boundary_conditions`: Apply boundary conditions to edge tiles
#[cfg(feature = "cuda-codegen")]
pub fn generate_packed_kernels() -> String {
    let mut output = String::new();

    output.push_str(PACKED_KERNELS_HEADER);
    output.push_str("\nextern \"C\" {\n\n");

    output.push_str(&generate_exchange_all_halos());
    output.push('\n');

    output.push_str(&generate_fdtd_all_tiles());
    output.push('\n');

    output.push_str(&generate_upload_tile_data());
    output.push('\n');

    output.push_str(&generate_read_all_interiors());
    output.push('\n');

    output.push_str(&generate_inject_impulse());
    output.push('\n');

    output.push_str(&generate_apply_boundary_conditions());

    output.push_str("\n}  // extern \"C\"\n");

    output
}

/// Header comment for generated packed kernels.
pub const PACKED_KERNELS_HEADER: &str = r#"// CUDA Kernels for Packed Tile-Based FDTD Wave Simulation
// Generated by ringkernel-cuda-codegen from Rust DSL
//
// All tiles packed contiguously: [Tile(0,0)][Tile(1,0)]...[Tile(n,m)]
// Each tile is 18x18 floats (16x16 interior + 1-cell halo)
//
// Benefits:
// - Zero host<->GPU transfers during simulation
// - All tiles computed in parallel
// - Halo exchange is just GPU memory copies
"#;

/// Generate exchange_all_halos kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_exchange_all_halos() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn exchange_all_halos(
            packed_buffer: &mut [f32],
            copies: &[u32],
            num_copies: i32,
        ) {
            let idx = block_idx_x() * block_dim_x() + thread_idx_x();
            if idx >= num_copies {
                return;
            }

            let src_idx = copies[(idx * 2) as usize];
            let dst_idx = copies[(idx * 2 + 1) as usize];

            packed_buffer[dst_idx as usize] = packed_buffer[src_idx as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Halo Exchange Kernel - copies all halo edges between adjacent tiles\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate fdtd_all_tiles kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_fdtd_all_tiles() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn fdtd_all_tiles(
            packed_curr: &[f32],
            packed_prev: &mut [f32],
            tiles_x: i32,
            tiles_y: i32,
            tile_size: i32,
            buffer_width: i32,
            c2: f32,
            damping: f32,
        ) {
            let tile_x = block_idx_x();
            let tile_y = block_idx_y();
            let lx = thread_idx_x();
            let ly = thread_idx_y();

            if tile_x >= tiles_x || tile_y >= tiles_y {
                return;
            }
            if lx >= tile_size || ly >= tile_size {
                return;
            }

            let tile_buffer_size = buffer_width * buffer_width;
            let tile_idx = tile_y * tiles_x + tile_x;
            let tile_offset = tile_idx * tile_buffer_size;

            let idx = tile_offset + (ly + 1) * buffer_width + (lx + 1);

            let p = packed_curr[idx as usize];
            let p_prev_val = packed_prev[idx as usize];

            let p_n = packed_curr[(idx - buffer_width) as usize];
            let p_s = packed_curr[(idx + buffer_width) as usize];
            let p_w = packed_curr[(idx - 1) as usize];
            let p_e = packed_curr[(idx + 1) as usize];

            let laplacian = p_n + p_s + p_e + p_w - 4.0 * p;
            let p_new = 2.0 * p - p_prev_val + c2 * laplacian;

            packed_prev[idx as usize] = p_new * damping;
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Batched FDTD Kernel - computes FDTD for ALL tiles in single launch\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate upload_tile_data kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_upload_tile_data() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn upload_tile_data(
            packed_buffer: &mut [f32],
            staging: &[f32],
            tile_x: i32,
            tile_y: i32,
            tiles_x: i32,
            buffer_width: i32,
        ) {
            let lx = thread_idx_x();
            let ly = thread_idx_y();

            if lx >= buffer_width || ly >= buffer_width {
                return;
            }

            let tile_buffer_size = buffer_width * buffer_width;
            let tile_idx = tile_y * tiles_x + tile_x;
            let tile_offset = tile_idx * tile_buffer_size;

            let local_idx = ly * buffer_width + lx;
            let global_idx = tile_offset + local_idx;

            packed_buffer[global_idx as usize] = staging[local_idx as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Upload Initial State - copies initial data to packed buffer\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate read_all_interiors kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_read_all_interiors() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn read_all_interiors(
            packed_buffer: &[f32],
            output: &mut [f32],
            tiles_x: i32,
            tiles_y: i32,
            tile_size: i32,
            buffer_width: i32,
            grid_width: i32,
            grid_height: i32,
        ) {
            let gx = block_idx_x() * block_dim_x() + thread_idx_x();
            let gy = block_idx_y() * block_dim_y() + thread_idx_y();

            if gx >= grid_width || gy >= grid_height {
                return;
            }

            let tile_x = gx / tile_size;
            let tile_y = gy / tile_size;

            let lx = gx % tile_size;
            let ly = gy % tile_size;

            let tile_buffer_size = buffer_width * buffer_width;
            let tile_idx = tile_y * tiles_x + tile_x;
            let tile_offset = tile_idx * tile_buffer_size;
            let src_idx = tile_offset + (ly + 1) * buffer_width + (lx + 1);

            let dst_idx = gy * grid_width + gx;

            output[dst_idx as usize] = packed_buffer[src_idx as usize];
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Read All Interiors - extracts all tile interiors for visualization\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate inject_impulse kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_inject_impulse() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn inject_impulse(
            packed_buffer: &mut [f32],
            tile_x: i32,
            tile_y: i32,
            local_x: i32,
            local_y: i32,
            tiles_x: i32,
            buffer_width: i32,
            amplitude: f32,
        ) {
            let tile_buffer_size = buffer_width * buffer_width;
            let tile_idx = tile_y * tiles_x + tile_x;
            let tile_offset = tile_idx * tile_buffer_size;
            let idx = tile_offset + (local_y + 1) * buffer_width + (local_x + 1);

            packed_buffer[idx as usize] = packed_buffer[idx as usize] + amplitude;
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!("// Inject Impulse - adds energy to specific cell\n{}", cuda),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

/// Generate apply_boundary_conditions kernel.
#[cfg(feature = "cuda-codegen")]
fn generate_apply_boundary_conditions() -> String {
    use syn::parse_quote;

    let kernel_fn: syn::ItemFn = parse_quote! {
        fn apply_boundary_conditions(
            packed_buffer: &mut [f32],
            tiles_x: i32,
            tiles_y: i32,
            tile_size: i32,
            buffer_width: i32,
            reflection_coeff: f32,
        ) {
            let edge = block_idx_x();
            let idx = thread_idx_x();

            let tile_buffer_size = buffer_width * buffer_width;

            if edge == 0 {
                // North boundary: tiles with tile_y == 0
                let tile_x = idx / tile_size;
                let cell_x = idx % tile_size;
                if tile_x >= tiles_x {
                    return;
                }

                let tile_idx = 0 * tiles_x + tile_x;
                let tile_offset = tile_idx * tile_buffer_size;
                let src_idx = tile_offset + 1 * buffer_width + (cell_x + 1);
                let dst_idx = tile_offset + 0 * buffer_width + (cell_x + 1);
                packed_buffer[dst_idx as usize] = packed_buffer[src_idx as usize] * reflection_coeff;
            } else if edge == 1 {
                // South boundary: tiles with tile_y == tiles_y - 1
                let tile_x = idx / tile_size;
                let cell_x = idx % tile_size;
                if tile_x >= tiles_x {
                    return;
                }

                let tile_idx = (tiles_y - 1) * tiles_x + tile_x;
                let tile_offset = tile_idx * tile_buffer_size;
                let src_idx = tile_offset + tile_size * buffer_width + (cell_x + 1);
                let dst_idx = tile_offset + (tile_size + 1) * buffer_width + (cell_x + 1);
                packed_buffer[dst_idx as usize] = packed_buffer[src_idx as usize] * reflection_coeff;
            } else if edge == 2 {
                // West boundary: tiles with tile_x == 0
                let tile_y = idx / tile_size;
                let cell_y = idx % tile_size;
                if tile_y >= tiles_y {
                    return;
                }

                let tile_idx = tile_y * tiles_x + 0;
                let tile_offset = tile_idx * tile_buffer_size;
                let src_idx = tile_offset + (cell_y + 1) * buffer_width + 1;
                let dst_idx = tile_offset + (cell_y + 1) * buffer_width + 0;
                packed_buffer[dst_idx as usize] = packed_buffer[src_idx as usize] * reflection_coeff;
            } else if edge == 3 {
                // East boundary: tiles with tile_x == tiles_x - 1
                let tile_y = idx / tile_size;
                let cell_y = idx % tile_size;
                if tile_y >= tiles_y {
                    return;
                }

                let tile_idx = tile_y * tiles_x + (tiles_x - 1);
                let tile_offset = tile_idx * tile_buffer_size;
                let src_idx = tile_offset + (cell_y + 1) * buffer_width + tile_size;
                let dst_idx = tile_offset + (cell_y + 1) * buffer_width + (tile_size + 1);
                packed_buffer[dst_idx as usize] = packed_buffer[src_idx as usize] * reflection_coeff;
            }
        }
    };

    match transpile_global_kernel(&kernel_fn) {
        Ok(cuda) => format!(
            "// Apply Boundary Conditions - handles domain edges for packed tiles\n{}",
            cuda
        ),
        Err(e) => format!("// Transpilation error: {}\n", e),
    }
}

// ============================================================================
// Fallback implementations (when cuda-codegen is not enabled)
// ============================================================================

#[cfg(not(feature = "cuda-codegen"))]
pub fn generate_tile_kernels() -> String {
    "// CUDA codegen not enabled - use handwritten shaders/fdtd_tile.cu".to_string()
}

#[cfg(not(feature = "cuda-codegen"))]
pub fn generate_packed_kernels() -> String {
    "// CUDA codegen not enabled - use handwritten shaders/fdtd_packed.cu".to_string()
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    #[allow(unused_imports)]
    use super::*;

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_tile_kernels_structure() {
        let source = generate_tile_kernels();

        // Check all kernels are present
        assert!(source.contains("fdtd_tile_step"), "Missing fdtd_tile_step");
        assert!(source.contains("extract_halo"), "Missing extract_halo");
        assert!(source.contains("inject_halo"), "Missing inject_halo");
        assert!(source.contains("read_interior"), "Missing read_interior");
        assert!(
            source.contains("apply_boundary_reflection"),
            "Missing apply_boundary_reflection"
        );

        // Check extern "C" wrapper
        assert!(source.contains("extern \"C\""), "Missing extern C");

        // Check FDTD kernel has correct structure
        assert!(source.contains("__global__ void fdtd_tile_step"));
        assert!(source.contains("threadIdx.x"));
        assert!(source.contains("threadIdx.y"));
        assert!(source.contains("buffer_width = 18") || source.contains("* 18"));
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_packed_kernels_structure() {
        let source = generate_packed_kernels();

        // Check all kernels are present
        assert!(
            source.contains("exchange_all_halos"),
            "Missing exchange_all_halos"
        );
        assert!(source.contains("fdtd_all_tiles"), "Missing fdtd_all_tiles");
        assert!(
            source.contains("upload_tile_data"),
            "Missing upload_tile_data"
        );
        assert!(
            source.contains("read_all_interiors"),
            "Missing read_all_interiors"
        );
        assert!(source.contains("inject_impulse"), "Missing inject_impulse");
        assert!(
            source.contains("apply_boundary_conditions"),
            "Missing apply_boundary_conditions"
        );

        // Check batched FDTD uses blockIdx
        assert!(source.contains("blockIdx.x"), "Missing blockIdx usage");
        assert!(source.contains("blockIdx.y"), "Missing blockIdx.y usage");
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_fdtd_tile_step_matches_handwritten() {
        let generated = generate_fdtd_tile_step();

        // Verify key structural elements match handwritten version
        assert!(generated.contains("const float* __restrict__ pressure"));
        assert!(generated.contains("float* __restrict__ pressure_prev"));
        assert!(generated.contains("float c2"));
        assert!(generated.contains("float damping"));
        assert!(generated.contains("if (lx >= 16 || ly >= 16) return;"));
        assert!(
            generated.contains("idx = (ly + 1) * buffer_width + (lx + 1)")
                || generated.contains("(ly + 1) * 18 + (lx + 1)")
        );
        assert!(generated.contains("laplacian"));
        assert!(generated.contains("* damping"));

        println!("Generated fdtd_tile_step:\n{}", generated);
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_generated_vs_handwritten_tile() {
        let generated = generate_tile_kernels();
        let handwritten = include_str!("../shaders/fdtd_tile.cu");

        // Count kernels in both
        let gen_kernel_count = generated.matches("__global__").count();
        let hw_kernel_count = handwritten.matches("__global__").count();

        assert_eq!(
            gen_kernel_count, hw_kernel_count,
            "Kernel count mismatch: generated={}, handwritten={}",
            gen_kernel_count, hw_kernel_count
        );
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_generated_vs_handwritten_packed() {
        let generated = generate_packed_kernels();
        let handwritten = include_str!("../shaders/fdtd_packed.cu");

        // Count kernels in both
        let gen_kernel_count = generated.matches("__global__").count();
        let hw_kernel_count = handwritten.matches("__global__").count();

        assert_eq!(
            gen_kernel_count, hw_kernel_count,
            "Kernel count mismatch: generated={}, handwritten={}",
            gen_kernel_count, hw_kernel_count
        );
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_match_expression_transpiles_to_switch() {
        // Test extract_halo uses switch for edge selection
        let extract = generate_extract_halo();
        assert!(
            extract.contains("switch (edge)"),
            "extract_halo should use switch: {}",
            extract
        );
        assert!(
            extract.contains("case 0:"),
            "extract_halo should have case 0"
        );
        assert!(
            extract.contains("case 1:"),
            "extract_halo should have case 1"
        );
        assert!(
            extract.contains("case 2:"),
            "extract_halo should have case 2"
        );
        assert!(
            extract.contains("default:"),
            "extract_halo should have default"
        );

        // Test inject_halo uses switch for edge selection
        let inject = generate_inject_halo();
        assert!(
            inject.contains("switch (edge)"),
            "inject_halo should use switch: {}",
            inject
        );

        // Test apply_boundary_reflection uses switch
        let boundary = generate_apply_boundary_reflection();
        assert!(
            boundary.contains("switch (edge)"),
            "apply_boundary_reflection should use switch: {}",
            boundary
        );

        println!("Generated extract_halo:\n{}", extract);
    }

    #[test]
    #[cfg(feature = "cuda-codegen")]
    fn test_all_kernels_transpile_successfully() {
        // Verify all tile kernels transpile without errors
        let tile_source = generate_tile_kernels();
        assert!(
            !tile_source.contains("Transpilation error"),
            "Tile kernels had transpilation errors:\n{}",
            tile_source
        );

        // Verify all packed kernels transpile without errors
        let packed_source = generate_packed_kernels();
        assert!(
            !packed_source.contains("Transpilation error"),
            "Packed kernels had transpilation errors:\n{}",
            packed_source
        );

        // Count __global__ functions to ensure all generated
        let tile_count = tile_source.matches("__global__").count();
        let packed_count = packed_source.matches("__global__").count();

        assert_eq!(tile_count, 5, "Expected 5 tile kernels, got {}", tile_count);
        assert_eq!(
            packed_count, 6,
            "Expected 6 packed kernels, got {}",
            packed_count
        );

        println!(
            "Successfully generated {} tile kernels and {} packed kernels",
            tile_count, packed_count
        );
    }
}