init

src/nouveau/compiler/nak/reg_tracker.rs

@@ -0,0 +1,186 @@
+// Copyright © 2022 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+use rustc_hash::FxHashMap;
+
+use crate::ir::*;
+
+use std::ops::{Index, IndexMut, Range};
+
+pub struct RegTracker<T> {
+    reg: [T; 255],
+    ureg: [T; 80],
+    pred: [T; 7],
+    upred: [T; 7],
+    carry: [T; 1],
+}
+
+fn new_array_with<T, const N: usize>(f: &impl Fn() -> T) -> [T; N] {
+    let mut v = Vec::with_capacity(N);
+    for _ in 0..N {
+        v.push(f());
+    }
+    v.try_into()
+        .unwrap_or_else(|_| panic!("Array size mismatch"))
+}
+
+impl<T> RegTracker<T> {
+    pub fn new_with(f: &impl Fn() -> T) -> Self {
+        Self {
+            reg: new_array_with(f),
+            ureg: new_array_with(f),
+            pred: new_array_with(f),
+            upred: new_array_with(f),
+            carry: new_array_with(f),
+        }
+    }
+}
+
+impl<T> Index<RegRef> for RegTracker<T> {
+    type Output = [T];
+
+    fn index(&self, reg: RegRef) -> &[T] {
+        let range = reg.idx_range();
+        let range = Range {
+            start: usize::try_from(range.start).unwrap(),
+            end: usize::try_from(range.end).unwrap(),
+        };
+
+        match reg.file() {
+            RegFile::GPR => &self.reg[range],
+            RegFile::UGPR => &self.ureg[range],
+            RegFile::Pred => &self.pred[range],
+            RegFile::UPred => &self.upred[range],
+            RegFile::Carry => &self.carry[range],
+            RegFile::Bar => &[], // Barriers have a HW scoreboard
+            RegFile::Mem => panic!("Not a register"),
+        }
+    }
+}
+
+impl<T> IndexMut<RegRef> for RegTracker<T> {
+    fn index_mut(&mut self, reg: RegRef) -> &mut [T] {
+        let range = reg.idx_range();
+        let range = Range {
+            start: usize::try_from(range.start).unwrap(),
+            end: usize::try_from(range.end).unwrap(),
+        };
+
+        match reg.file() {
+            RegFile::GPR => &mut self.reg[range],
+            RegFile::UGPR => &mut self.ureg[range],
+            RegFile::Pred => &mut self.pred[range],
+            RegFile::UPred => &mut self.upred[range],
+            RegFile::Carry => &mut self.carry[range],
+            RegFile::Bar => &mut [], // Barriers have a HW scoreboard
+            RegFile::Mem => panic!("Not a register"),
+        }
+    }
+}
+
+/// Memory-light version of [RegTracker].
+///
+/// This version uses sparse hashmaps instead of dense arrays.
+#[derive(Clone, PartialEq, Eq, Default)]
+pub struct SparseRegTracker<T: Default> {
+    regs: FxHashMap<RegRef, T>,
+}
+
+impl<T: Default + Clone + Eq> SparseRegTracker<T> {
+    pub fn for_each_pred(&mut self, f: impl FnMut(&mut T)) {
+        self.for_each_ref_mut(RegRef::new(RegFile::Pred, 0, 7), f);
+    }
+
+    pub fn for_each_carry(&mut self, f: impl FnMut(&mut T)) {
+        self.for_each_ref_mut(RegRef::new(RegFile::Carry, 0, 1), f);
+    }
+
+    pub fn merge_with(&mut self, other: &Self, mut f: impl FnMut(&mut T, &T)) {
+        use std::collections::hash_map::Entry;
+
+        for (k, v) in other.regs.iter() {
+            match self.regs.entry(*k) {
+                Entry::Occupied(mut occupied_entry) => {
+                    f(occupied_entry.get_mut(), v);
+                }
+                Entry::Vacant(vacant_entry) => {
+                    vacant_entry.insert((*v).clone());
+                }
+            }
+        }
+    }
+
+    pub fn retain(&mut self, mut f: impl FnMut(&mut T) -> bool) {
+        self.regs.retain(|_k, v| f(v));
+    }
+}
+
+/// Common behavior for [RegTracker] and [SparseRegTracker]
+pub trait RegRefIterable<T> {
+    fn for_each_ref_mut(&mut self, reg: RegRef, f: impl FnMut(&mut T));
+
+    fn for_each_instr_pred_mut(
+        &mut self,
+        instr: &Instr,
+        mut f: impl FnMut(&mut T),
+    ) {
+        if let PredRef::Reg(reg) = &instr.pred.pred_ref {
+            self.for_each_ref_mut(*reg, |t| f(t));
+        }
+    }
+
+    fn for_each_instr_src_mut(
+        &mut self,
+        instr: &Instr,
+        mut f: impl FnMut(usize, &mut T),
+    ) {
+        for (i, src) in instr.srcs().iter().enumerate() {
+            match &src.src_ref {
+                SrcRef::Reg(reg) => {
+                    self.for_each_ref_mut(*reg, |t| f(i, t));
+                }
+                SrcRef::CBuf(CBufRef {
+                    buf: CBuf::BindlessUGPR(reg),
+                    ..
+                }) => {
+                    self.for_each_ref_mut(*reg, |t| f(i, t));
+                }
+                _ => (),
+            }
+        }
+    }
+
+    fn for_each_instr_dst_mut(
+        &mut self,
+        instr: &Instr,
+        mut f: impl FnMut(usize, &mut T),
+    ) {
+        for (i, dst) in instr.dsts().iter().enumerate() {
+            if let Dst::Reg(reg) = dst {
+                self.for_each_ref_mut(*reg, |t| f(i, t));
+            }
+        }
+    }
+}
+
+impl<T: Default> RegRefIterable<T> for SparseRegTracker<T> {
+    fn for_each_ref_mut(&mut self, reg: RegRef, mut f: impl FnMut(&mut T)) {
+        match reg.file() {
+            RegFile::Bar => return, // Barriers have a HW scoreboard
+            RegFile::Mem => panic!("Not a register"),
+            _ => {}
+        }
+
+        for i in 0..reg.comps() {
+            f(self.regs.entry(reg.comp(i)).or_default());
+        }
+    }
+}
+
+impl<T> RegRefIterable<T> for RegTracker<T> {
+    fn for_each_ref_mut(&mut self, reg: RegRef, mut f: impl FnMut(&mut T)) {
+        for entry in &mut self[reg] {
+            f(entry);
+        }
+    }
+}