add loop unrolling optimization pass

examples/spv-unroll-test/lib.rs

@@ -0,0 +1,70 @@
+#![no_std]
+
+use spirv_std::spirv;
+
+pub use spirv_std::glam;
+
+// Basic while loop: i < 4
+#[spirv(compute(threads(1)))]
+pub fn test_unroll(
+    #[spirv(global_invocation_id)] _id: glam::UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] data: &mut [u32; 64],
+) {
+    let mut i: u32 = 0;
+    while i < 4 {
+        data[i as usize] = i;
+        i += 1;
+    }
+}
+
+// for i in 0..4
+#[spirv(compute(threads(1)))]
+pub fn test_for_range(
+    #[spirv(global_invocation_id)] _id: glam::UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] data: &mut [u32; 64],
+) {
+    for i in 0u32..4 {
+        data[i as usize] = i * 2;
+    }
+}
+
+// Accumulator: sum of 0..4
+#[spirv(compute(threads(1)))]
+pub fn test_accumulate(
+    #[spirv(global_invocation_id)] _id: glam::UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] data: &mut [u32; 64],
+) {
+    let mut sum: u32 = 0;
+    for i in 0u32..4 {
+        sum += i;
+    }
+    data[0] = sum;
+}
+
+// Two independent state variables
+#[spirv(compute(threads(1)))]
+pub fn test_two_state_vars(
+    #[spirv(global_invocation_id)] _id: glam::UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] data: &mut [u32; 64],
+) {
+    let mut i: u32 = 0;
+    let mut val: u32 = 10;
+    while i < 4 {
+        data[i as usize] = val;
+        i += 1;
+        val += 3;
+    }
+}
+
+// u64 for-range: for i in 0u64..4
+#[spirv(compute(threads(1)))]
+pub fn test_u64_for_range(
+    #[spirv(global_invocation_id)] _id: glam::UVec3,
+    #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] data: &mut [u32; 64],
+) {
+    for i in 0u64..4 {
+        data[i as usize] = i as u32;
+    }
+}
+
+

examples/spv-unroll-test/main.rs

@@ -0,0 +1,27 @@
+use std::path::Path;
+use std::rc::Rc;
+
+fn main() -> std::io::Result<()> {
+    let args: Vec<_> = std::env::args().collect();
+    let in_file = match args.as_slice() {
+        [_, f] => f.clone(),
+        _ => {
+            eprintln!("usage: {} FILE.spv", args[0]);
+            std::process::exit(1);
+        }
+    };
+    let in_path = Path::new(&in_file);
+
+    let cx = Rc::new(spirt::Context::new());
+    let mut module = spirt::Module::lower_from_spv_file(cx, in_path)?;
+
+    spirt::passes::legalize::structurize_func_cfgs(&mut module);
+
+    println!("{}", spirt::print::Plan::for_module(&module).pretty_print());
+    println!("new optimzied --------");
+
+    spirt::passes::unroll::unroll_loops(&mut module);
+    println!("{}", spirt::print::Plan::for_module(&module).pretty_print());
+
+    Ok(())
+}

src/context.rs

@@ -348,19 +348,14 @@ pub struct EntityOrientedDenseMap<K: EntityOrientedMapKey<V>, V> {
 // since the ideal state is one chunk per map, the slow case might never be hit,
 // unless one `EntityOrientedDenseMap` is used with more than one `EntityDefs`,
 // which could still maybe be implemented more efficiently than `FxHashMap`.
-#[derive(Clone)]
+#[derive(Clone, Default)]
 enum SmallFxHashMap<K, V> {
+    #[default]
     Empty,
     One(K, V),
     More(FxHashMap<K, V>),
 }
 
-impl<K, V> Default for SmallFxHashMap<K, V> {
-    fn default() -> Self {
-        Self::Empty
-    }
-}
-
 impl<K: Copy + Eq + Hash, V: Default> SmallFxHashMap<K, V> {
     fn get_mut_or_insert_default(&mut self, k: K) -> &mut V {
         // HACK(eddyb) to avoid borrowing issues, this is done in two stages:

src/lib.rs

@@ -71,7 +71,7 @@
     clippy::dbg_macro,
     clippy::debug_assert_with_mut_call,
     clippy::doc_markdown,
-    clippy::empty_enum,
+    clippy::empty_enums,
     clippy::enum_glob_use,
     clippy::exit,
     clippy::expl_impl_clone_on_copy,
@@ -168,6 +168,7 @@ pub mod passes {
     pub mod legalize;
     pub mod link;
     pub mod qptr;
+    pub mod unroll;
 }
 pub mod qptr;
 pub mod spv;