Fold OpLoad from Private variables with constant initializers

Disclosure: I had AI write this and I only superficially understand the code.

Replace OpLoad from Private/Function
storage class variables with their constant initializers. This eliminates
pointer-to-pointer patterns like `&&123` which would otherwise generate
invalid SPIR-V in Logical addressing mode (pointer-to-pointer is only
allowed for StorageBuffer/Workgroup, not Private).

The optimization runs after inlining (which may expose such patterns)
and before DCE (so the now-unused variables can be removed).

Also extend DCE to treat Private storage class variables as pure,
allowing them to be removed when unused.

Newer versions of spirv-val fire catch this.

Author: LegNeato

crates/rustc_codegen_spirv/src/linker/dce.rs

@@ -282,7 +282,13 @@ fn instruction_is_pure(inst: &Instruction) -> bool {
         | PtrEqual
         | PtrNotEqual
         | PtrDiff => true,
-        Variable => inst.operands.first() == Some(&Operand::StorageClass(StorageClass::Function)),
+        // Variables with Function or Private storage class are pure and can be DCE'd if unused.
+        // Other storage classes (Input, Output, Uniform, etc.) are part of the shader interface
+        // and must be kept.
+        Variable => matches!(
+            inst.operands.first(),
+            Some(&Operand::StorageClass(StorageClass::Function | StorageClass::Private))
+        ),
         _ => false,
     }
 }

crates/rustc_codegen_spirv/src/linker/mod.rs

@@ -415,6 +415,16 @@ pub fn link(
         inline::inline(sess, &mut output)?;
     }
 
+    // Fold OpLoad from Private/Function variables with constant initializers.
+    // This must run after inlining (which may expose such patterns) and before DCE
+    // (so that the now-unused variables can be removed).
+    // This is critical for pointer-to-pointer patterns like `&&123` which would
+    // otherwise generate invalid SPIR-V in Logical addressing mode.
+    {
+        let _timer = sess.timer("link_fold_load_from_constant_variable");
+        peephole_opts::fold_load_from_constant_variable(&mut output);
+    }
+
     {
         let _timer = sess.timer("link_dce-after-inlining");
         dce::dce(&mut output);

crates/rustc_codegen_spirv/src/linker/peephole_opts.rs

@@ -606,3 +606,76 @@ pub fn bool_fusion(
     }
     super::apply_rewrite_rules(&rewrite_rules, &mut function.blocks);
 }
+
+/// Fold `OpLoad` from Private/Function storage class variables that have constant initializers.
+///
+/// This optimization handles patterns like:
+/// ```text
+/// %ptr = OpVariable %_ptr_Private_T Private %initializer
+/// %val = OpLoad %T %ptr
+/// ```
+/// After this optimization, uses of `%val` are replaced with `%initializer`.
+///
+/// This is particularly important for pointer-to-pointer constants (e.g., `&&123`)
+/// where the outer pointer variable has a known initializer (the inner pointer).
+/// Without this optimization, such patterns generate invalid SPIR-V in Logical
+/// addressing mode (pointer-to-pointer with Private storage class is not allowed).
+pub fn fold_load_from_constant_variable(
+    module: &mut Module,
+) {
+    use rspirv::spirv::StorageClass;
+
+    // Build a map of variable ID -> initializer ID for Private/Function variables
+    // that have constant initializers.
+    let var_initializers: FxHashMap<Word, Word> = module
+        .types_global_values
+        .iter()
+        .filter_map(|inst| {
+            if inst.class.opcode != Op::Variable {
+                return None;
+            }
+            // Check storage class - only Private and Function can be folded
+            let storage_class = inst.operands.first()?.unwrap_storage_class();
+            if !matches!(storage_class, StorageClass::Private | StorageClass::Function) {
+                return None;
+            }
+            // Check for initializer (second operand after storage class)
+            let initializer = inst.operands.get(1)?.id_ref_any()?;
+            Some((inst.result_id?, initializer))
+        })
+        .collect();
+
+    if var_initializers.is_empty() {
+        return;
+    }
+
+    // Rewrite OpLoad instructions that load from variables with known initializers
+    let mut rewrite_rules: FxHashMap<Word, Word> = FxHashMap::default();
+
+    for func in &mut module.functions {
+        for block in &mut func.blocks {
+            for inst in &mut block.instructions {
+                if inst.class.opcode != Op::Load {
+                    continue;
+                }
+                // OpLoad has pointer as first operand
+                let ptr_id = inst.operands[0].unwrap_id_ref();
+                if let Some(&initializer) = var_initializers.get(&ptr_id) {
+                    // This load can be replaced with the initializer value
+                    if let Some(result_id) = inst.result_id {
+                        rewrite_rules.insert(result_id, initializer);
+                        // Turn this instruction into a Nop
+                        *inst = Instruction::new(Op::Nop, None, None, Vec::new());
+                    }
+                }
+            }
+        }
+    }
+
+    if !rewrite_rules.is_empty() {
+        // Apply rewrite rules to all functions
+        for func in &mut module.functions {
+            super::apply_rewrite_rules(&rewrite_rules, &mut func.blocks);
+        }
+    }
+}