From af8b5ee30b61e6a8dff1aec4985bfa6fbb1c01d9 Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Mon, 1 Jun 2026 02:28:18 -0400
Subject: [PATCH] kraid: Make SrcRef::Imm32 explicitly non-zero

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/42154>
---
 src/panfrost/compiler/kraid/encode_v9.rs      |  6 ++--
 src/panfrost/compiler/kraid/ir.rs             | 35 +++++++++++++------
 .../compiler/kraid/small_constants.rs         |  2 ++
 3 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/src/panfrost/compiler/kraid/encode_v9.rs b/src/panfrost/compiler/kraid/encode_v9.rs
index cf6817daa6e..eb28461ba18 100644
--- a/src/panfrost/compiler/kraid/encode_v9.rs
+++ b/src/panfrost/compiler/kraid/encode_v9.rs
@@ -479,7 +479,7 @@ impl V9Instr for OpFAdd {
                 variant: self.dst_type.try_into().unwrap(),
                 dst: op_encode_dst(self, &self.dst),
                 src0: op_encode_src(self, &self.srcs[0]),
-                imm1w: *imm,
+                imm1w: (*imm).into(),
             })
         } else {
             e.encode(Fadd {
@@ -527,7 +527,7 @@ impl V9Instr for OpIAdd {
                 variant: self.dst_type.try_into().unwrap(),
                 dst: op_encode_dst(self, &self.dst),
                 src0: op_encode_src(self, &self.srcs[0]),
-                imm1w: *imm,
+                imm1w: (*imm).into(),
             })
         } else {
             e.encode(Iadd {
@@ -623,7 +623,7 @@ impl V9Instr for OpMov {
             e.encode(MovImm {
                 variant: self.dst_type.try_into().unwrap(),
                 dst: op_encode_dst(self, &self.dst),
-                imm1w: *imm,
+                imm1w: (*imm).into(),
             })
         } else {
             e.encode(Mov {
diff --git a/src/panfrost/compiler/kraid/ir.rs b/src/panfrost/compiler/kraid/ir.rs
index ba65e59b48c..050fca6334e 100644
--- a/src/panfrost/compiler/kraid/ir.rs
+++ b/src/panfrost/compiler/kraid/ir.rs
@@ -12,6 +12,7 @@ pub use crate::swizzle::Swizzle;
 use compiler::as_slice::*;
 
 use std::fmt;
+use std::num::NonZeroU32;
 use std::ops::{Deref, DerefMut};
 
 pub struct SmallConstant {
@@ -156,7 +157,7 @@ pub enum SrcRef {
     /// A zero value
     Zero,
     /// A 32-bit immediate
-    Imm32(u32),
+    Imm32(NonZeroU32),
     FAU(FAURef),
     SSA(SSARef),
     Reg(RegRef),
@@ -205,10 +206,22 @@ impl SrcRef {
 
 impl From<u32> for SrcRef {
     fn from(u: u32) -> SrcRef {
-        if u == 0 {
-            SrcRef::Zero
+        if let Some(nz) = NonZeroU32::new(u) {
+            SrcRef::Imm32(nz)
         } else {
-            SrcRef::Imm32(u)
+            SrcRef::Zero
+        }
+    }
+}
+
+impl TryFrom<&SrcRef> for u32 {
+    type Error = &'static str;
+
+    fn try_from(src_ref: &SrcRef) -> Result<u32, Self::Error> {
+        match src_ref {
+            SrcRef::Zero => Ok(0),
+            SrcRef::Imm32(nz) => Ok((*nz).into()),
+            _ => Err("Value not known at compile time"),
         }
     }
 }
@@ -365,16 +378,18 @@ impl Src {
     }
 
     pub fn is_zero(&self) -> bool {
-        matches!(self.src_ref, SrcRef::Zero | SrcRef::Imm32(0))
+        matches!(self.src_ref, SrcRef::Zero)
     }
 
     pub fn replicates_byte(&self) -> bool {
         match self.src_ref {
             SrcRef::Zero => true,
-            SrcRef::Imm32(u) => self.swizzle.fold_u32(u).is_some_and(|u| {
-                let b = u.to_le_bytes();
-                b[0] == b[1] && b[0] == b[2] && b[0] == b[3]
-            }),
+            SrcRef::Imm32(u) => {
+                self.swizzle.fold_u32(u.into()).is_some_and(|u| {
+                    let b = u.to_le_bytes();
+                    b[0] == b[1] && b[0] == b[2] && b[0] == b[3]
+                })
+            }
             _ => self.swizzle.replicates_byte(),
         }
     }
@@ -384,7 +399,7 @@ impl Src {
             SrcRef::Zero => true,
             SrcRef::Imm32(u) => self
                 .swizzle
-                .fold_u32(u)
+                .fold_u32(u.into())
                 .is_some_and(|u| (u & 0xffff) == (u >> 16)),
             _ => self.swizzle.replicates_half(),
         }
diff --git a/src/panfrost/compiler/kraid/small_constants.rs b/src/panfrost/compiler/kraid/small_constants.rs
index 6288ea2c7c8..4fa74749a8a 100644
--- a/src/panfrost/compiler/kraid/small_constants.rs
+++ b/src/panfrost/compiler/kraid/small_constants.rs
@@ -17,6 +17,8 @@ fn try_lower_src(src: &mut Src, sc_table: &[SmallConstant]) -> bool {
     let SrcRef::Imm32(imm32) = src.src_ref else {
         return false;
     };
+    let imm32 = u32::from(imm32);
+
     let imm_bytes_read = src.swizzle.bytes_read();
 
     if imm_bytes_read.count_ones() == 1 {