src/gpu_engine.rs

use wgpu::util::DeviceExt;
use bytemuck::{Pod, Zeroable};

#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
pub struct PackedPhaseU32 {
    pub data: u32,
}

pub async fn execute_gpu_sync(tensor_data: &mut [u32], iterations: usize) {
    let instance = wgpu::Instance::default();
    let adapter = instance.request_adapter(&wgpu::RequestAdapterOptions::default()).await.unwrap();
    
    let supported_limits = adapter.limits();
    let (device, queue) = adapter.request_device(&wgpu::DeviceDescriptor {
        label: None,
        required_features: wgpu::Features::empty(),
        required_limits: supported_limits.clone(),
    }, None).await.unwrap();

    let data_byte_size = std::mem::size_of_val(tensor_data);
    let max_buffer = supported_limits.max_storage_buffer_binding_size as usize;
    
    let chunk_size = if data_byte_size > max_buffer {
        max_buffer / 4 
    } else {
        tensor_data.len()
    };

    // [패치 완료] 100만 개가 넘는 워크그룹을 2D 평면(X, Y)으로 접어서 처리하는 WGSL 셰이더
    let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: None,
        source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(r#"
            @group(0) @binding(0) var<storage, read_write> phases: array<u32>;
            
            @compute @workgroup_size(256, 1, 1)
            fn main(
                @builtin(workgroup_id) group_id: vec3<u32>,
                @builtin(local_invocation_id) local_id: vec3<u32>,
                @builtin(num_workgroups) num_groups: vec3<u32>
            ) {
                // 2D 디스패치 그룹 좌표를 1D 선형 인덱스로 변환
                let group_idx = group_id.y * num_groups.x + group_id.x;
                let idx = group_idx * 256u + local_id.x;
                
                if (idx >= arrayLength(&phases)) { return; }
                
                let phase = phases[idx];
                let context_signal = 0x0F0F0F0Fu;
                let mask = 0x55555555u;
                
                let diff = phase ^ context_signal;
                let sync_pull = (diff & mask) >> 1u;
                let pull_11 = sync_pull | (sync_pull << 1u);
                
                phases[idx] = (phase & ~pull_11) | pull_11;
            }
        "#)),
    });

    let compute_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: None,
        layout: None,
        module: &shader,
        entry_point: "main",
        
    });

    let bind_group_layout = compute_pipeline.get_bind_group_layout(0);

    for chunk in tensor_data.chunks_mut(chunk_size) {
        let buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
            label: Some("Full Tensor Buffer"),
            contents: bytemuck::cast_slice(chunk),
            usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST,
        });

        let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: &bind_group_layout,
            entries: &[wgpu::BindGroupEntry {
                binding: 0,
                resource: buffer.as_entire_binding(),
            }],
        });

        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        
        let total_threads = chunk.len() as u32;
        let total_groups = (total_threads + 255) / 256;
        
        // Vulkan 한계(65535) 우회를 위한 2D 분할 (Folding)
        let max_groups_x = 65535;
        let groups_x = total_groups.min(max_groups_x);
        let groups_y = (total_groups + max_groups_x - 1) / max_groups_x;

        for _ in 0..iterations {
            let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            cpass.set_pipeline(&compute_pipeline);
            cpass.set_bind_group(0, &bind_group, &[]);
            cpass.dispatch_workgroups(groups_x, groups_y, 1);
        }
        queue.submit(Some(encoder.finish()));
        device.poll(wgpu::Maintain::Wait);
    }
}