From 60c3690165d152cfc0b78ae050a4dfa027d97afe Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Wed, 27 May 2026 13:53:34 -0400
Subject: [PATCH] kraid: Add a lower_small_constants() pass

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/41841>
---
 src/panfrost/compiler/kraid/compile.rs        |  4 ++
 src/panfrost/compiler/kraid/lib.rs            |  1 +
 .../compiler/kraid/small_constants.rs         | 67 +++++++++++++++++++
 3 files changed, 72 insertions(+)
 create mode 100644 src/panfrost/compiler/kraid/small_constants.rs

diff --git a/src/panfrost/compiler/kraid/compile.rs b/src/panfrost/compiler/kraid/compile.rs
index b32cfc515d3..7896eee59a1 100644
--- a/src/panfrost/compiler/kraid/compile.rs
+++ b/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();
diff --git a/src/panfrost/compiler/kraid/lib.rs b/src/panfrost/compiler/kraid/lib.rs
index 90214a0c82d..a15bbcf2c27 100644
--- a/src/panfrost/compiler/kraid/lib.rs
+++ b/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;
diff --git a/src/panfrost/compiler/kraid/small_constants.rs b/src/panfrost/compiler/kraid/small_constants.rs
new file mode 100644
index 00000000000..6288ea2c7c8
--- /dev/null
+++ b/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);
+                }
+            }
+        }
+    }
+}