From 718ef00ca4d617680fdb666b10d0d307afc9ad3d Mon Sep 17 00:00:00 2001
From: Faith Ekstrand <faith.ekstrand@collabora.com>
Date: Tue, 11 Jun 2024 13:44:45 -0500
Subject: [PATCH] nak/ra: Add a concept of pinned registers to RegAllocator

Unlike the pinned set in VecRegAllocator which exists for the duration
of an instruction, registers which are pinned in the main allocator are
pinned until the register is freed.  The pinned set in VecRegAllocator
is initialized to a copy of the one in the main register allocator.

Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/29591>
---
 src/nouveau/compiler/nak/assign_regs.rs | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/nouveau/compiler/nak/assign_regs.rs b/src/nouveau/compiler/nak/assign_regs.rs
index a5501401c73..cb768d5324c 100644
--- a/src/nouveau/compiler/nak/assign_regs.rs
+++ b/src/nouveau/compiler/nak/assign_regs.rs
@@ -184,6 +184,7 @@ struct RegAllocator {
     file: RegFile,
     num_regs: u32,
     used: BitSet,
+    pinned: BitSet,
     reg_ssa: Vec<SSAValue>,
     ssa_reg: HashMap<SSAValue, u32>,
 }
@@ -194,6 +195,7 @@ impl RegAllocator {
             file: file,
             num_regs: num_regs,
             used: BitSet::new(),
+            pinned: BitSet::new(),
             reg_ssa: Vec::new(),
             ssa_reg: HashMap::new(),
         }
@@ -211,6 +213,10 @@ impl RegAllocator {
         self.used.get(reg.try_into().unwrap())
     }
 
+    pub fn reg_is_pinned(&self, reg: u32) -> bool {
+        self.pinned.get(reg.try_into().unwrap())
+    }
+
     pub fn try_get_reg(&self, ssa: SSAValue) -> Option<u32> {
         self.ssa_reg.get(&ssa).cloned()
     }
@@ -249,6 +255,7 @@ impl RegAllocator {
         let reg_usize = usize::try_from(reg).unwrap();
         assert!(self.reg_ssa[reg_usize] == ssa);
         self.used.remove(reg_usize);
+        self.pinned.remove(reg_usize);
         reg
     }
 
@@ -267,6 +274,11 @@ impl RegAllocator {
         self.used.insert(reg_usize);
     }
 
+    pub fn pin_reg(&mut self, reg: u32) {
+        assert!(self.reg_is_used(reg));
+        self.pinned.insert(reg.try_into().unwrap());
+    }
+
     fn reg_range_is_unset(set: &BitSet, reg: u32, comps: u8) -> bool {
         for c in 0..u32::from(comps) {
             if set.get((reg + c).try_into().unwrap()) {
@@ -394,10 +406,11 @@ struct VecRegAllocator<'a> {
 
 impl<'a> VecRegAllocator<'a> {
     fn new(ra: &'a mut RegAllocator) -> Self {
+        let pinned = ra.pinned.clone();
         VecRegAllocator {
-            ra: ra,
+            ra,
             pcopy: OpParCopy::new(),
-            pinned: Default::default(),
+            pinned,
             evicted: HashMap::new(),
         }
     }
@@ -1104,6 +1117,9 @@ impl AssignRegsBlock {
                 for (ssa, reg) in &pred_raf.ssa_reg {
                     if bl.is_live_in(ssa) {
                         raf.assign_reg(*ssa, *reg);
+                        if pred_raf.reg_is_pinned(*reg) {
+                            raf.pin_reg(*reg);
+                        }
                         self.live_in.push(LiveValue {
                             live_ref: LiveRef::SSA(*ssa),
                             reg_ref: RegRef::new(raf.file(), *reg, 1),