kraid: Add a lower_small_constants() pass

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41841>

src/panfrost/compiler/kraid/compile.rs

@@ -73,6 +73,10 @@ pub extern "C" fn kraid_compile_nir(
     dump_shader(&s, "after translation from NIR");
     s.validate();
 
+    s.lower_small_constants();
+    dump_shader(&s, "after lowering small constants");
+    s.validate();
+
     s.assign_registers();
     dump_shader(&s, "after register assignment");
     s.validate();

src/panfrost/compiler/kraid/lib.rs

@@ -13,6 +13,7 @@ mod model;
 mod nir;
 mod ops;
 mod ra;
+mod small_constants;
 mod ssa_value;
 mod swizzle;
 mod validate;

src/panfrost/compiler/kraid/small_constants.rs

@@ -0,0 +1,67 @@
+// Copyright © 2026 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+use crate::ir::*;
+
+fn try_fold_src(src: &mut Src, sc: &SmallConstant, sc_swz: Swizzle) -> bool {
+    let Some(swz) = sc_swz.swizzle(src.swizzle) else {
+        return false;
+    };
+
+    src.swizzle = swz;
+    src.src_ref = SrcRef::FAU(FAURef::from(sc).into());
+    true
+}
+
+fn try_lower_src(src: &mut Src, sc_table: &[SmallConstant]) -> bool {
+    let SrcRef::Imm32(imm32) = src.src_ref else {
+        return false;
+    };
+    let imm_bytes_read = src.swizzle.bytes_read();
+
+    if imm_bytes_read.count_ones() == 1 {
+        let imm_byte = imm_bytes_read.trailing_zeros();
+        let imm8 = (imm32 >> (imm_byte * 8)) as u8;
+        for sc in sc_table {
+            for b in 0..4 {
+                if imm8 == ((sc.imm32 >> (b * 8)) as u8)
+                    && try_fold_src(src, sc, Swizzle::replicate_byte(b))
+                {
+                    return true;
+                }
+            }
+        }
+    } else if imm_bytes_read == 0b1100 || imm_bytes_read == 0b0011 {
+        let imm_half = imm_bytes_read.trailing_zeros() / 2;
+        let imm16 = (imm32 >> (imm_half * 16)) as u16;
+        for sc in sc_table {
+            for h in 0..2 {
+                if imm16 == ((sc.imm32 >> (h * 16)) as u16)
+                    && try_fold_src(src, sc, Swizzle::replicate_half(h))
+                {
+                    return true;
+                }
+            }
+        }
+    } else {
+        for sc in sc_table {
+            if imm32 == sc.imm32 && try_fold_src(src, sc, Swizzle::NONE) {
+                return true;
+            }
+        }
+    }
+    false
+}
+
+impl Shader<'_> {
+    pub fn lower_small_constants(&mut self) {
+        let sc_table = self.model.small_constants();
+        for b in self.blocks.iter_mut() {
+            for i in b.instrs.iter_mut() {
+                for src in i.srcs_mut() {
+                    try_lower_src(src, sc_table);
+                }
+            }
+        }
+    }
+}