kraid: Add OpRegIn and OpRegOut

These will be implemented entirely in the register allocator as register
assignment constraings (and possibly a copy in the case of OpRegOut).
Only OpRegIn is implemented in the trivial RA.  OpRegOut will have to
wait for the real RA.

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

src/panfrost/compiler/kraid/ir.rs

@@ -149,6 +149,20 @@ pub enum RegRange {
     Regs(u8),
 }
 
+impl From<RegRange> for Swizzle {
+    fn from(range: RegRange) -> Swizzle {
+        match range {
+            RegRange::Byte0 => Swizzle::B0000,
+            RegRange::Byte1 => Swizzle::B1111,
+            RegRange::Byte2 => Swizzle::B2222,
+            RegRange::Byte3 => Swizzle::B3333,
+            RegRange::Half0 => Swizzle::H00,
+            RegRange::Half1 => Swizzle::H11,
+            RegRange::Regs(_) => Swizzle::NONE,
+        }
+    }
+}
+
 #[derive(Clone, Copy, PartialEq)]
 pub struct RegRef {
     pub idx: u8,

src/panfrost/compiler/kraid/ops.rs

@@ -673,6 +673,54 @@ impl fmt::Display for OpMov {
     }
 }
 
+#[repr(C)]
+#[derive(Clone, Opcode)]
+#[variants(dst_type in [I8, I16, I32, I64])]
+pub struct OpRegIn {
+    pub dst: Dst,
+    pub dst_type: DataType,
+    pub reg: RegRef,
+}
+
+impl fmt::Display for OpRegIn {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{} = REG_IN.{} {}", &self.dst, self.dst_type, &self.reg)
+    }
+}
+
+impl VirtualOpcode for OpRegIn {
+    fn dst_supports_lanes(&self, lanes: DstLanes) -> bool {
+        lanes == DstLanes::from(self.reg.range)
+    }
+}
+
+#[repr(C)]
+#[derive(Clone, Opcode)]
+#[variants(src_type in [I8, I16, I32, I64])]
+pub struct OpRegOut {
+    pub reg: RegRef,
+    pub src_type: DataType,
+    pub src: Src,
+}
+
+impl fmt::Display for OpRegOut {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{} = REG_OUT.{} {}",
+            &self.reg,
+            self.src_type,
+            self.fmt_src(&self.src),
+        )
+    }
+}
+
+impl VirtualOpcode for OpRegOut {
+    fn src_supports_swizzle(&self, _src: &Src, swizzle: Swizzle) -> bool {
+        swizzle == Swizzle::from(self.reg.range)
+    }
+}
+
 #[repr(C)]
 #[derive(Clone, Opcode)]
 pub struct OpNop {}
@@ -922,6 +970,8 @@ pub enum Op {
     MkVecV4I8(Box<OpMkVecV4I8>),
     Nop(OpNop),
     Mov(Box<OpMov>),
+    RegIn(Box<OpRegIn>),
+    RegOut(Box<OpRegOut>),
     ShiftLop(Box<OpShiftLop>),
     Swz(Box<OpSwz>),
     Store(Box<OpStore>),
@@ -938,6 +988,8 @@ impl Op {
             Op::Copy(op) => Some(op.as_ref()),
             Op::MkVecV2I8(op) => Some(op.as_ref()),
             Op::MkVecV4I8(op) => Some(op.as_ref()),
+            Op::RegIn(op) => Some(op.as_ref()),
+            Op::RegOut(op) => Some(op.as_ref()),
             Op::Swz(op) => Some(op.as_ref()),
             _ => None,
         }

src/panfrost/compiler/kraid/ra.rs

@@ -129,7 +129,31 @@ fn ra_trivial(s: &mut Shader) {
     let mut ssa_b: FxHashMap<SSAValue, u8> = Default::default();
 
     for (bi, block) in s.blocks.iter_mut().enumerate() {
-        for (ip, instr) in block.instrs.iter_mut().enumerate() {
+        for (ip, mut instr) in
+            std::mem::take(&mut block.instrs).into_iter().enumerate()
+        {
+            if let Op::RegIn(op) = instr.op {
+                let DstRef::SSA(vec) = op.dst.dst_ref else {
+                    panic!("We must have SSA destinations");
+                };
+
+                let b = op.reg.idx * 4 + op.reg.byte_offset();
+                let bytes = vec.bytes();
+                debug_assert_eq!(bytes, op.reg.bytes());
+
+                for (i, ssa) in vec.iter().enumerate() {
+                    ssa_b.insert(*ssa, b + u8::try_from(i * 4).unwrap());
+                }
+                for i in 0..bytes {
+                    let b = usize::from(b) + usize::from(i);
+                    assert!(!byte_used.contains(b));
+                    byte_used.insert(b);
+                }
+
+                // Drop the actual instruction on the floor
+                continue;
+            }
+
             for src in instr.srcs_mut() {
                 let SrcRef::SSA(vec) = &mut src.src_ref else {
                     continue;
@@ -155,15 +179,7 @@ fn ra_trivial(s: &mut Shader) {
                 }
 
                 let reg = reg_ref_for_byte(vec_b, vec.bytes());
-                let swz = match reg.range {
-                    RegRange::Byte0 => Swizzle::B0000,
-                    RegRange::Byte1 => Swizzle::B1111,
-                    RegRange::Byte2 => Swizzle::B2222,
-                    RegRange::Byte3 => Swizzle::B3333,
-                    RegRange::Half0 => Swizzle::H00,
-                    RegRange::Half1 => Swizzle::H11,
-                    RegRange::Regs(_) => Swizzle::NONE,
-                };
+                let swz = Swizzle::from(reg.range);
                 src.swizzle = swz
                     .swizzle(src.swizzle)
                     .expect("16-bit and smaller sources have to swizzle");
@@ -268,6 +284,8 @@ fn ra_trivial(s: &mut Shader) {
             {
                 *dst = reg.into();
             }
+
+            block.instrs.push(instr);
         }
     }
 }

src/panfrost/compiler/kraid/validate.rs

@@ -96,9 +96,24 @@ impl Shader<'_> {
             blocks.insert(bb.label);
         }
 
+        let mut allow_reg_in = true;
+        let mut allow_non_reg_out = true;
         let mut ssa_vals: FxHashSet<SSAValue> = Default::default();
-        for bb in &self.blocks {
+        for (bi, bb) in self.blocks.iter().enumerate() {
             for i in &bb.instrs {
+                if matches!(&i.op, Op::RegIn(_)) {
+                    assert!(bi == 0);
+                    assert!(allow_reg_in);
+                } else if !matches!(&i.op, Op::Nop(_)) {
+                    allow_reg_in = false;
+                }
+
+                if matches!(&i.op, Op::RegOut(_)) {
+                    allow_non_reg_out = false;
+                } else if !matches!(&i.op, Op::Nop(_)) {
+                    assert!(allow_non_reg_out);
+                }
+
                 validate_instr(&i, &mut ssa_vals);
             }
         }