nak: Force RA to allocate bar_in/out to the same register

OpBreak and OpBSsy aren't very SSA friendly as they require bar_in and
bar_out to be assigned the same register.  We need to encure that RA
knows about this restriction.  For now, we just special-case these two
instructions.  In the future we may want a more generic mechanism for
this but it's not worth it for just two instructions.

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

src/nouveau/compiler/nak_assign_regs.rs

@@ -801,6 +801,12 @@ fn instr_assign_regs_file(
 }
 
 impl PerRegFile<RegAllocator> {
+    pub fn assign_reg(&mut self, ssa: SSAValue, reg: RegRef) {
+        assert!(reg.file() == ssa.file());
+        assert!(reg.comps() == 1);
+        self[ssa.file()].assign_reg(ssa, reg.base_idx());
+    }
+
     pub fn free_killed(&mut self, killed: &KillSet) {
         for ssa in killed.iter() {
             self[ssa.file()].free_ssa(*ssa);
@@ -927,6 +933,48 @@ impl AssignRegsBlock {
 
                 None
             }
+            Op::Break(op) => {
+                for src in op.srcs_as_mut_slice() {
+                    if let SrcRef::SSA(ssa) = src.src_ref {
+                        assert!(ssa.comps() == 1);
+                        let reg = self.get_scalar(ssa[0]);
+                        src.src_ref = reg.into();
+                    }
+                }
+
+                self.ra.free_killed(srcs_killed);
+
+                if let Dst::SSA(ssa) = &op.bar_out {
+                    let reg = *op.bar_in.src_ref.as_reg().unwrap();
+                    self.ra.assign_reg(ssa[0], reg);
+                    op.bar_out = reg.into();
+                }
+
+                self.ra.free_killed(dsts_killed);
+
+                Some(instr)
+            }
+            Op::BSSy(op) => {
+                for src in op.srcs_as_mut_slice() {
+                    if let SrcRef::SSA(ssa) = src.src_ref {
+                        assert!(ssa.comps() == 1);
+                        let reg = self.get_scalar(ssa[0]);
+                        src.src_ref = reg.into();
+                    }
+                }
+
+                self.ra.free_killed(srcs_killed);
+
+                if let Dst::SSA(ssa) = &op.bar_out {
+                    let reg = *op.bar_in.src_ref.as_reg().unwrap();
+                    self.ra.assign_reg(ssa[0], reg);
+                    op.bar_out = reg.into();
+                }
+
+                self.ra.free_killed(dsts_killed);
+
+                Some(instr)
+            }
             Op::Copy(copy) => {
                 if let SrcRef::SSA(src_vec) = &copy.src.src_ref {
                     debug_assert!(src_vec.comps() == 1);

src/nouveau/compiler/nak_builder.rs

@@ -426,6 +426,7 @@ impl<'a> SSAInstrBuilder<'a> {
         self.b.as_vec()
     }
 
+    #[allow(dead_code)]
     pub fn as_mapped_instrs(self) -> MappedInstrs {
         self.b.as_mapped_instrs()
     }

src/nouveau/compiler/nak_legalize.rs

@@ -4,6 +4,7 @@
  */
 
 use crate::nak_ir::*;
+use crate::nak_liveness::{BlockLiveness, Liveness, SimpleLiveness};
 
 use std::collections::{HashMap, HashSet};
 
@@ -68,7 +69,12 @@ fn swap_srcs_if_not_reg(x: &mut Src, y: &mut Src) {
     }
 }
 
-fn legalize_instr(b: &mut impl SSABuilder, instr: &mut Instr) {
+fn legalize_instr(
+    b: &mut impl SSABuilder,
+    bl: &impl BlockLiveness,
+    ip: usize,
+    instr: &mut Instr,
+) {
     match &mut instr.op {
         Op::FAdd(op) => {
             let [ref mut src0, ref mut src1] = op.srcs;
@@ -255,11 +261,27 @@ fn legalize_instr(b: &mut impl SSABuilder, instr: &mut Instr) {
             copy_src_if_not_reg(b, &mut op.handle, RegFile::GPR);
             copy_src_if_cbuf(b, &mut op.stream, RegFile::GPR);
         }
+        Op::Break(op) => {
+            let bar_in = op.bar_in.src_ref.as_ssa().unwrap();
+            if !op.bar_out.is_none() && bl.is_live_after_ip(&bar_in[0], ip) {
+                let gpr = b.bmov_to_gpr(op.bar_in);
+                let tmp = b.bmov_to_bar(gpr.into());
+                op.bar_in = tmp.into();
+            }
+        }
+        Op::BSSy(op) => {
+            let bar_in = op.bar_in.src_ref.as_ssa().unwrap();
+            if !op.bar_out.is_none() && bl.is_live_after_ip(&bar_in[0], ip) {
+                let gpr = b.bmov_to_gpr(op.bar_in);
+                let tmp = b.bmov_to_bar(gpr.into());
+                op.bar_in = tmp.into();
+            }
+        }
         Op::OutFinal(op) => {
             copy_src_if_not_reg(b, &mut op.handle, RegFile::GPR);
         }
         Op::Ldc(_) => (), // Nothing to do
-        Op::Break(_) | Op::BSSy(_) | Op::BSync(_) => (),
+        Op::BSync(_) => (),
         Op::Vote(_) => (), // Nothing to do
         Op::Copy(_) => (), // Nothing to do
         _ => {
@@ -331,11 +353,21 @@ fn legalize_instr(b: &mut impl SSABuilder, instr: &mut Instr) {
 
 impl Shader {
     pub fn legalize(&mut self) {
-        self.map_instrs(|mut instr, ssa_alloc| -> MappedInstrs {
-            let mut b = SSAInstrBuilder::new(ssa_alloc);
-            legalize_instr(&mut b, &mut instr);
-            b.push_instr(instr);
-            b.as_mapped_instrs()
-        });
+        for f in &mut self.functions {
+            let live = SimpleLiveness::for_function(f);
+
+            for (bi, b) in f.blocks.iter_mut().enumerate() {
+                let bl = live.block_live(bi);
+
+                let mut instrs = Vec::new();
+                for (ip, mut instr) in b.instrs.drain(..).enumerate() {
+                    let mut b = SSAInstrBuilder::new(&mut f.ssa_alloc);
+                    legalize_instr(&mut b, bl, ip, &mut instr);
+                    b.push_instr(instr);
+                    instrs.append(&mut b.as_vec());
+                }
+                b.instrs = instrs;
+            }
+        }
     }
 }