nak: Add OpLea

Reviewed-by: Faith Ekstrand <faith.ekstrand@collabora.com>
Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/32517>

src/nouveau/compiler/nak/hw_tests.rs

@@ -725,6 +725,37 @@ fn test_op_isetp() {
     }
 }
 
+#[test]
+fn test_op_lea() {
+    if RunSingleton::get().sm.sm() >= 70 {
+        let src_mods = [
+            (SrcMod::None, SrcMod::None),
+            (SrcMod::INeg, SrcMod::None),
+            (SrcMod::None, SrcMod::INeg),
+        ];
+
+        for (intermediate_mod, b_mod) in src_mods {
+            for shift in 0..32 {
+                for dst_high in [false, true] {
+                    let mut op = OpLea {
+                        dst: Dst::None,
+                        overflow: Dst::None,
+                        a: 0.into(),
+                        b: 0.into(),
+                        a_high: 0.into(),
+                        shift,
+                        dst_high,
+                        intermediate_mod,
+                    };
+                    op.b.src_mod = b_mod;
+
+                    test_foldable_op(op);
+                }
+            }
+        }
+    }
+}
+
 #[test]
 fn test_op_lop2() {
     if RunSingleton::get().sm.sm() < 70 {

src/nouveau/compiler/nak/ir.rs

@@ -3783,6 +3783,82 @@ impl DisplayOp for OpISetP {
 }
 impl_display_for_op!(OpISetP);
 
+#[repr(C)]
+#[derive(Clone, SrcsAsSlice, DstsAsSlice)]
+pub struct OpLea {
+    #[dst_type(GPR)]
+    pub dst: Dst,
+
+    #[dst_type(Pred)]
+    pub overflow: Dst,
+
+    #[src_type(ALU)]
+    pub a: Src,
+
+    #[src_type(I32)]
+    pub b: Src,
+
+    #[src_type(ALU)]
+    pub a_high: Src, // High 32-bits of a if .dst_high is set
+
+    pub shift: u8,
+    pub dst_high: bool,
+    pub intermediate_mod: SrcMod, // Modifier for shifted temporary (a << shift)
+}
+
+impl Foldable for OpLea {
+    fn fold(&self, _sm: &dyn ShaderModel, f: &mut OpFoldData<'_>) {
+        let a = f.get_u32_src(self, &self.a);
+        let mut b = f.get_u32_src(self, &self.b);
+        let a_high = f.get_u32_src(self, &self.a_high);
+
+        let mut overflow = false;
+
+        let mut shift_result = if self.dst_high {
+            let a = a as u64;
+            let a_high = a_high as u64;
+            let a = (a_high << 32) | a;
+
+            (a >> (32 - self.shift)) as u32
+        } else {
+            a << self.shift
+        };
+
+        if self.intermediate_mod.is_ineg() {
+            let o;
+            (shift_result, o) = u32::overflowing_add(!shift_result, 1);
+            overflow |= o;
+        }
+
+        if self.b.src_mod.is_ineg() {
+            let o;
+            (b, o) = u32::overflowing_add(!b, 1);
+            overflow |= o;
+        }
+
+        let (dst, o) = u32::overflowing_add(shift_result, b);
+        overflow |= o;
+
+        f.set_u32_dst(self, &self.dst, dst as u32);
+        f.set_pred_dst(self, &self.overflow, overflow);
+    }
+}
+
+impl DisplayOp for OpLea {
+    fn fmt_op(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "lea")?;
+        if self.dst_high {
+            write!(f, ".hi")?;
+        }
+        write!(f, " {} {} {}", self.a, self.shift, self.b)?;
+        if self.dst_high {
+            write!(f, " {}", self.a_high)?;
+        }
+        Ok(())
+    }
+}
+impl_display_for_op!(OpLea);
+
 #[repr(C)]
 #[derive(Clone, SrcsAsSlice, DstsAsSlice)]
 pub struct OpLop2 {
@@ -6234,6 +6310,7 @@ pub enum Op {
     IMul(OpIMul),
     IMnMx(OpIMnMx),
     ISetP(OpISetP),
+    Lea(OpLea),
     Lop2(OpLop2),
     Lop3(OpLop3),
     PopC(OpPopC),
@@ -6727,6 +6804,7 @@ impl Instr {
             | Op::IMad64(_)
             | Op::IMnMx(_)
             | Op::ISetP(_)
+            | Op::Lea(_)
             | Op::Lop2(_)
             | Op::Lop3(_)
             | Op::Shf(_)

src/nouveau/compiler/nak/opt_copy_prop.rs

@@ -650,6 +650,7 @@ impl CopyPropPass {
                     Op::IAdd3X(add) => {
                         !add.overflow[0].is_none() || !add.overflow[1].is_none()
                     }
+                    Op::Lea(lea) => !lea.overflow.is_none(),
                     _ => false,
                 };
 
@@ -658,7 +659,9 @@ impl CopyPropPass {
                     let mut src_type = src_types[i];
                     if force_alu_src_type {
                         src_type = match src_type {
-                            SrcType::B32 | SrcType::I32 => SrcType::ALU,
+                            SrcType::ALU | SrcType::B32 | SrcType::I32 => {
+                                SrcType::ALU
+                            }
                             SrcType::Carry | SrcType::Pred => src_type,
                             _ => panic!("Unhandled src_type"),
                         };

src/nouveau/compiler/nak/sm70.rs

@@ -83,6 +83,7 @@ impl ShaderModel for ShaderModel70 {
             | Op::IMad(_)
             | Op::IMad64(_)
             | Op::ISetP(_)
+            | Op::Lea(_)
             | Op::Lop3(_)
             | Op::Mov(_)
             | Op::PLop3(_)
@@ -1747,6 +1748,59 @@ impl SM70Op for OpISetP {
     }
 }
 
+impl SM70Op for OpLea {
+    fn legalize(&mut self, b: &mut LegalizeBuilder) {
+        let gpr = op_gpr(self);
+        b.copy_alu_src_if_not_reg(&mut self.a, gpr, SrcType::ALU);
+        if self.dst_high {
+            b.copy_alu_src_if_both_not_reg(
+                &self.b,
+                &mut self.a_high,
+                gpr,
+                SrcType::ALU,
+            );
+        }
+    }
+
+    fn encode(&self, e: &mut SM70Encoder<'_>) {
+        assert!(self.a.src_mod == SrcMod::None);
+        assert!(
+            self.intermediate_mod == SrcMod::None
+                || self.b.src_mod == SrcMod::None
+        );
+
+        let c = if self.dst_high {
+            Some(&self.a_high)
+        } else {
+            None
+        };
+
+        if self.is_uniform() {
+            e.encode_ualu(
+                0x091,
+                Some(&self.dst),
+                Some(&self.a),
+                Some(&self.b),
+                c,
+            );
+        } else {
+            e.encode_alu(
+                0x011,
+                Some(&self.dst),
+                Some(&self.a),
+                Some(&self.b),
+                c,
+            );
+        }
+
+        e.set_bit(72, self.intermediate_mod.is_ineg());
+        e.set_field(75..80, self.shift);
+        e.set_bit(80, self.dst_high);
+        e.set_pred_dst(81..84, self.overflow);
+        e.set_bit(74, false); // .X
+    }
+}
+
 fn src_as_lop_imm(src: &Src) -> Option<bool> {
     let x = match src.src_ref {
         SrcRef::Zero => false,
@@ -3501,6 +3555,7 @@ macro_rules! as_sm70_op_match {
             Op::IMad64(op) => op,
             Op::IMnMx(op) => op,
             Op::ISetP(op) => op,
+            Op::Lea(op) => op,
             Op::Lop3(op) => op,
             Op::PopC(op) => op,
             Op::Shf(op) => op,