diff --git a/src/nouveau/compiler/nak/calc_instr_deps.rs b/src/nouveau/compiler/nak/calc_instr_deps.rs index cde8691f048..c5197e933af 100644 --- a/src/nouveau/compiler/nak/calc_instr_deps.rs +++ b/src/nouveau/compiler/nak/calc_instr_deps.rs @@ -4,7 +4,7 @@ use crate::api::{GetDebugFlags, DEBUG}; use crate::ir::*; use crate::opt_instr_sched_common::estimate_block_weight; -use crate::reg_tracker::RegTracker; +use crate::reg_tracker::{RegRefIterable, RegTracker}; use compiler::dataflow::ForwardDataflow; use rustc_hash::{FxHashMap, FxHashSet}; diff --git a/src/nouveau/compiler/nak/opt_instr_sched_postpass.rs b/src/nouveau/compiler/nak/opt_instr_sched_postpass.rs index 62b441f1801..5ce08c11f63 100644 --- a/src/nouveau/compiler/nak/opt_instr_sched_postpass.rs +++ b/src/nouveau/compiler/nak/opt_instr_sched_postpass.rs @@ -3,6 +3,7 @@ use crate::ir::*; use crate::opt_instr_sched_common::*; +use crate::reg_tracker::RegRefIterable; use crate::reg_tracker::RegTracker; use std::cmp::max; use std::cmp::Reverse; diff --git a/src/nouveau/compiler/nak/reg_tracker.rs b/src/nouveau/compiler/nak/reg_tracker.rs index 49341ca1e8f..3da71b87fd9 100644 --- a/src/nouveau/compiler/nak/reg_tracker.rs +++ b/src/nouveau/compiler/nak/reg_tracker.rs @@ -1,6 +1,8 @@ // Copyright © 2022 Collabora, Ltd. // SPDX-License-Identifier: MIT +use rustc_hash::FxHashMap; + use crate::ir::*; use std::ops::{Index, IndexMut, Range}; @@ -14,7 +16,7 @@ pub struct RegTracker { } fn new_array_with(f: &impl Fn() -> T) -> [T; N] { - let mut v = Vec::new(); + let mut v = Vec::with_capacity(N); for _ in 0..N { v.push(f()); } @@ -33,57 +35,6 @@ impl RegTracker { } } - pub 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 { - for i in &mut self[*reg] { - f(i); - } - } - } - - pub 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) => { - for t in &mut self[*reg] { - f(i, t); - } - } - SrcRef::CBuf(CBufRef { - buf: CBuf::BindlessUGPR(reg), - .. - }) => { - for t in &mut self[*reg] { - f(i, t); - } - } - _ => (), - } - } - } - - pub 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 { - for t in &mut self[*reg] { - f(i, t); - } - } - } - } - pub fn for_each_pred(&mut self, mut f: impl FnMut(&mut T)) { for p in &mut self.pred[..] { f(p); @@ -138,3 +89,110 @@ impl IndexMut for RegTracker { } } } + +/// Memory-light version of [RegTracker]. +/// +/// This version uses sparse hashmaps instead of dense arrays. +#[derive(Clone, PartialEq, Eq, Default)] +pub struct SparseRegTracker { + regs: FxHashMap, +} + +impl SparseRegTracker { + 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 { + 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 RegRefIterable for SparseRegTracker { + 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 RegRefIterable for RegTracker { + fn for_each_ref_mut(&mut self, reg: RegRef, mut f: impl FnMut(&mut T)) { + for entry in &mut self[reg] { + f(entry); + } + } +}