如学

wgpu-mine/main.wgsl

@@ -0,0 +1,38 @@
+@group(0) @binding(0) var<uniform> first_step: array<vec4<u32>, 64>;
+
+// 处理第二步的两次 rc4 名称, 输入名字 bytes, 从第一步得到的 val
+@compute
+@workgroup_size(16, 16)
+fn rc4_name(name_bytes: array<u32, 256>, name_len: u32) -> array<u32, 256> {
+    // var val: array<u32, 256> = first_step;
+    var val: array<u32, 256> = array<u32, 256>();
+    // 把 first_step 的值复制到 val
+    // first_step 内当成连续的内存即可
+    for (var i: u32 = 0; i < 64; i = i + 1) {
+        val[i * 4] = first_step[i].x;
+        val[i * 4 + 1] = first_step[i].y;
+        val[i * 4 + 2] = first_step[i].z;
+        val[i * 4 + 3] = first_step[i].w;
+    }
+    // 上面的重复两次
+    for (var n: u32 = 0; n < 2; n++) {    
+        var s: u32 = 0u;
+        var k: u32 = 0u;
+        for (var i: u32 = 0; i < 256; i = i + 1) {
+            if (k != 0) {
+                s = s + val[k - 1];
+            }
+            s = s + val[i];
+            s = s % 256u;
+            var tmp = val[i];
+            val[i] = val[s];
+            val[s] = tmp;
+            if (k == name_len - 1) {
+                k = 0u;
+            } else {
+                k++;
+            }
+        }
+    }
+    return val;
+}

wgpu-mine/src/main.rs

@@ -1,51 +1,30 @@
-// use wgpu::util::DeviceExt;
+use std::{borrow::Cow, io::Read};
-
-// async fn run() {
-
-// }
-
-// fn main() {
-//     println!("Hello, world!");
-// }
-
-use std::{borrow::Cow, str::FromStr};
 use wgpu::util::DeviceExt;
 
-pub struct Team {
+/// 工作信息
-    pub val: [u8; 256],
+pub struct Works {
+    /// 队伍名(统一)
+    pub team: String,
+    /// 队伍名(每个任务)
+    pub names: Vec<String>,
 }
 
-// Indicates a u32 overflow in an intermediate Collatz value
-const OVERFLOW: u32 = 0xffffffff;
-
 #[cfg_attr(test, allow(dead_code))]
 async fn run() {
-    let numbers = if std::env::args().len() <= 2 {
+    let works = Works {
-        let default = vec![1, 2, 3, 4];
+        team: "team".to_string(),
-        println!("No numbers were provided, defaulting to {default:?}");
+        names: vec!["name1".to_string(), "name2".to_string()],
-        default
-    } else {
-        std::env::args()
-            .skip(2)
-            .map(|s| u32::from_str(&s).expect("You must pass a list of positive integers!"))
-            .collect()
     };
 
-    let steps = execute_gpu(&numbers).await.unwrap();
+    let steps = execute_gpu(works).await.unwrap();
 
-    let disp_steps: Vec<String> = steps
+    let disp_steps: Vec<String> = steps.iter().map(|&n| n.to_string()).collect();
-        .iter()
-        .map(|&n| match n {
-            OVERFLOW => "OVERFLOW".to_string(),
-            _ => n.to_string(),
-        })
-        .collect();
 
     println!("Steps: [{}]", disp_steps.join(", "));
 }
 
 #[cfg_attr(test, allow(dead_code))]
-async fn execute_gpu(numbers: &[u32]) -> Option<Vec<u32>> {
+async fn execute_gpu(works: Works) -> Option<Vec<u32>> {
     // Instantiates instance of WebGPU
     let instance = wgpu::Instance::default();
 
@@ -66,18 +45,41 @@ async fn execute_gpu(numbers: &[u32]) -> Option<Vec<u32>> {
         .await
         .unwrap();
 
-    execute_gpu_inner(&device, &queue, numbers).await
+    execute_gpu_inner(&device, &queue, works).await
 }
 
-async fn execute_gpu_inner(device: &wgpu::Device, queue: &wgpu::Queue, numbers: &[u32]) -> Option<Vec<u32>> {
+async fn execute_gpu_inner(device: &wgpu::Device, queue: &wgpu::Queue, works: Works) -> Option<Vec<u32>> {
     // Loads the shader from WGSL
+    let source = {
+        let mut file = std::fs::File::open("main.wgsl").unwrap();
+        let mut source = String::new();
+        file.read_to_string(&mut source).unwrap();
+        source
+    };
+
     let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
         label: None,
-        source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("main.wgsl"))),
+        source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(&source)),
     });
 
+    // 初始化一下 team bytes 计算出的 val
+    let mut val: Vec<u8> = (0..=255).collect();
+    let team_bytes = works.team.as_bytes();
+    let t_len = works.team.len() + 1;
+    let mut s = 0_u8;
+    for i in 0..256 {
+        if (i % t_len) != 0 {
+            s = s.wrapping_add(team_bytes[(i % t_len) - 1]);
+        }
+        s = s.wrapping_add(val[i]);
+        val.swap(i, s as usize);
+    }
+
+    // 满足 uniform 的对其要求, 把 vec<u8; 256> 转成 vec<[u8; 4]; 64>
+    let uniform_val = val.chunks_exact(4).map(|c| [c[0], c[1], c[2], c[3]]).collect::<Vec<[u8; 4]>>();
+
     // Gets the size in bytes of the buffer.
-    let size = std::mem::size_of_val(numbers) as wgpu::BufferAddress;
+    let size = std::mem::size_of_val(&uniform_val) as wgpu::BufferAddress;
 
     // Instantiates buffer without data.
     // `usage` of buffer specifies how it can be used:
@@ -95,13 +97,11 @@ async fn execute_gpu_inner(device: &wgpu::Device, queue: &wgpu::Queue, numbers:
     //   A storage buffer (can be bound within a bind group and thus available to a shader).
     //   The destination of a copy.
     //   The source of a copy.
-    // let storage_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
+    let storage_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
-    //     label: Some("Storage Buffer"),
+        label: Some("team bytes input buffer"),
-    //     contents: bytemuck::cast_slice(numbers),
+        contents: bytemuck::cast_slice(&uniform_val),
-    //     usage: wgpu::BufferUsages::STORAGE
+        usage: wgpu::BufferUsages::STORAGE,
-    //         | wgpu::BufferUsages::COPY_DST
+    });
-    //         | wgpu::BufferUsages::COPY_SRC,
-    // });
 
     // A bind group defines how buffers are accessed by shaders.
     // It is to WebGPU what a descriptor set is to Vulkan.
@@ -119,15 +119,15 @@ async fn execute_gpu_inner(device: &wgpu::Device, queue: &wgpu::Queue, numbers:
     });
 
     // Instantiates the bind group, once again specifying the binding of buffers.
-    // let bind_group_layout = compute_pipeline.get_bind_group_layout(0);
+    let bind_group_layout = compute_pipeline.get_bind_group_layout(0);
-    // let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
+    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
-    //     label: None,
+        label: None,
-    //     layout: &bind_group_layout,
+        layout: &bind_group_layout,
-    //     entries: &[wgpu::BindGroupEntry {
+        entries: &[wgpu::BindGroupEntry {
-    //         binding: 0,
+            binding: 0,
-    //         resource: storage_buffer.as_entire_binding(),
+            resource: storage_buffer.as_entire_binding(),
-    //     }],
+        }],
-    // });
+    });
 
     // A command encoder executes one or many pipelines.
     // It is to WebGPU what a command buffer is to Vulkan.
@@ -138,9 +138,9 @@ async fn execute_gpu_inner(device: &wgpu::Device, queue: &wgpu::Queue, numbers:
             timestamp_writes: None,
         });
         cpass.set_pipeline(&compute_pipeline);
-        // cpass.set_bind_group(0, &bind_group, &[]);
+        cpass.set_bind_group(0, &bind_group, &[]);
         cpass.insert_debug_marker("compute collatz iterations");
-        // cpass.dispatch_workgroups(numbers.len() as u32, 1, 1); // Number of cells to run, the (x,y,z) size of item being processed
+        cpass.dispatch_workgroups(numbers.len() as u32, 1, 1); // Number of cells to run, the (x,y,z) size of item being processed
     }
     // Sets adds copy operation to command encoder.
     // Will copy data from storage buffer on GPU to staging buffer on CPU.

wgpu-mine/src/main.wgsl

@@ -1,91 +0,0 @@
-@group(0)
-@binding(0)
-var<storage, read_write> v_indices: array<u32>; // this is used as both input and output for convenience
-
-fn collatz_iterations(n_base: u32) -> u32{
-    var n: u32 = n_base;
-    var i: u32 = 0u;
-    loop {
-        if (n <= 1u) {
-            break;
-        }
-        if (n % 2u == 0u) {
-            n = n / 2u;
-        }
-        else {
-            // Overflow? (i.e. 3*n + 1 > 0xffffffffu?)
-            if (n >= 1431655765u) {   // 0x55555555u
-                return 4294967295u;   // 0xffffffffu
-            }
-
-            n = 3u * n + 1u;
-        }
-        i = i + 1u;
-    }
-    return i;
-}
-
-// @compute
-@workgroup_size(1)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-    v_indices[global_id.x] = collatz_iterations(v_indices[global_id.x]);
-}
-
-@compute
-@workgroup_size(1)
-// 输入 team_bytes: 队伍名转化成的字节数组
-fn team_bytes(team_bytes: array<u32, 256>, team_len: u32) -> array<u32, 256> {
-    var val: array<u32, 256> = team_bytes;
-    var s: u32 = 0u;
-    for (var i: u32 = 0; i < 256; i = i + 1) {
-        if (i % team_len != 0) {
-            s = s + val[(i % team_len) - 1];
-        }
-        s = s + val[i];
-        // team_bytes.swap(i, s);
-        var n = val[i];
-        val[i] = val[s];
-        val[s] = n;
-    }
-    return val;
-}
-
-/*
-#[derive(Debug, Clone)]
-pub struct TeamNamer {
-    pub team: String,
-    pub val: [u8; 256],
-}
-
-impl TeamNamer {
-    /// 方便使用的 new
-    /// 会检查长度是否超过 256
-    #[inline(always)]
-    pub fn new(team: &str) -> Option<Self> {
-        if team.len() > 256 {
-            warn!("Team too long({}): {}", team.len(), team);
-            return None;
-        }
-        Some(Self::new_unchecked(team))
-    }
-    #[inline(always)]
-    pub fn new_unchecked(team: &str) -> Self {
-        let team_bytes = team.as_bytes();
-        let mut val: [u8; 256] = (0..=255).collect::<Vec<u8>>().try_into().unwrap();
-        let mut s = 0_u8;
-        let t_len = team.len() + 1;
-        for i in 0..256 {
-            if (i % t_len) != 0 {
-                s = s.wrapping_add(team_bytes[(i % t_len) - 1]);
-            }
-            s = s.wrapping_add(val[i]);
-            val.swap(i, s as usize);
-        }
-        Self {
-            team: team.to_string(),
-            val,
-        }
-    }
-    #[inline(always)]
-    pub fn clone_vals(&self) -> [u8; 256] { self.val.clone() }
-}