init

src/nouveau/compiler/nak/legalize.rs

@@ -0,0 +1,620 @@
+// Copyright © 2022 Collabora, Ltd.
+// SPDX-License-Identifier: MIT
+
+use crate::api::{GetDebugFlags, DEBUG};
+use crate::const_tracker::ConstTracker;
+use crate::ir::*;
+use crate::liveness::{BlockLiveness, Liveness, SimpleLiveness};
+
+use rustc_hash::{FxHashMap, FxHashSet};
+
+pub fn src_is_upred_reg(src: &Src) -> bool {
+    match &src.src_ref {
+        SrcRef::True | SrcRef::False => false,
+        SrcRef::SSA(ssa) => {
+            assert!(ssa.comps() == 1);
+            match ssa[0].file() {
+                RegFile::Pred => false,
+                RegFile::UPred => true,
+                _ => panic!("Not a predicate source"),
+            }
+        }
+        SrcRef::Reg(_) => panic!("Not in SSA form"),
+        _ => panic!("Not a predicate source"),
+    }
+}
+
+pub fn src_is_reg(src: &Src, reg_file: RegFile) -> bool {
+    match &src.src_ref {
+        SrcRef::Zero | SrcRef::True | SrcRef::False => true,
+        SrcRef::SSA(ssa) => ssa.file() == reg_file,
+        SrcRef::Imm32(_) | SrcRef::CBuf(_) => false,
+        SrcRef::Reg(_) => panic!("Not in SSA form"),
+    }
+}
+
+pub fn swap_srcs_if_not_reg(
+    x: &mut Src,
+    y: &mut Src,
+    reg_file: RegFile,
+) -> bool {
+    if !src_is_reg(x, reg_file) && src_is_reg(y, reg_file) {
+        std::mem::swap(x, y);
+        true
+    } else {
+        false
+    }
+}
+
+fn src_is_imm(src: &Src) -> bool {
+    matches!(src.src_ref, SrcRef::Imm32(_))
+}
+
+pub enum PadValue {
+    Zero,
+    #[allow(dead_code)]
+    Undefined,
+}
+
+pub trait LegalizeBuildHelpers: SSABuilder {
+    fn copy_ssa(&mut self, ssa: &mut SSAValue, reg_file: RegFile) {
+        let tmp = self.alloc_ssa(reg_file);
+        self.copy_to(tmp.into(), (*ssa).into());
+        *ssa = tmp;
+    }
+
+    fn copy_ssa_ref(&mut self, vec: &mut SSARef, reg_file: RegFile) {
+        for ssa in &mut vec[..] {
+            self.copy_ssa(ssa, reg_file);
+        }
+    }
+
+    fn copy_pred_ssa_if_uniform(&mut self, ssa: &mut SSAValue) {
+        match ssa.file() {
+            RegFile::Pred => (),
+            RegFile::UPred => self.copy_ssa(ssa, RegFile::Pred),
+            _ => panic!("Not a predicate value"),
+        }
+    }
+
+    fn copy_pred_if_upred(&mut self, pred: &mut Pred) {
+        match &mut pred.pred_ref {
+            PredRef::None => (),
+            PredRef::SSA(ssa) => {
+                self.copy_pred_ssa_if_uniform(ssa);
+            }
+            PredRef::Reg(_) => panic!("Not in SSA form"),
+        }
+    }
+
+    fn copy_src_if_upred(&mut self, src: &mut Src) {
+        match &mut src.src_ref {
+            SrcRef::True | SrcRef::False => (),
+            SrcRef::SSA(ssa) => {
+                assert!(ssa.comps() == 1);
+                self.copy_pred_ssa_if_uniform(&mut ssa[0]);
+            }
+            SrcRef::Reg(_) => panic!("Not in SSA form"),
+            _ => panic!("Not a predicate source"),
+        }
+    }
+
+    fn copy_src_if_not_same_file(&mut self, src: &mut Src) {
+        let SrcRef::SSA(vec) = &mut src.src_ref else {
+            return;
+        };
+
+        if vec.comps() == 1 {
+            return;
+        }
+
+        let mut all_same = true;
+        let file = vec[0].file();
+        for i in 1..vec.comps() {
+            let c_file = vec[usize::from(i)].file();
+            if c_file != file {
+                debug_assert!(c_file.to_warp() == file.to_warp());
+                all_same = false;
+            }
+        }
+
+        if !all_same {
+            self.copy_ssa_ref(vec, file.to_warp());
+        }
+    }
+
+    fn align_reg(
+        &mut self,
+        src: &mut Src,
+        n_comps: usize,
+        pad_value: PadValue,
+    ) {
+        debug_assert!(!matches!(src.src_ref, SrcRef::Reg(_)));
+        let SrcRef::SSA(ref old_val) = src.src_ref else {
+            return;
+        };
+        assert!(old_val.len() <= n_comps);
+        assert!(src.is_unmodified());
+
+        let pad_fn = || {
+            Some(match pad_value {
+                PadValue::Zero => self.copy(0.into()),
+                PadValue::Undefined => self.undef(),
+            })
+        };
+
+        // Pad the given ssa_ref with either undefined or zero
+        let ssa_vals: Vec<_> = old_val
+            .iter()
+            .copied()
+            .chain(std::iter::from_fn(pad_fn))
+            .take(n_comps)
+            .collect();
+
+        // Collect it in a new ssa_ref and replace it with the original.
+        let val = SSARef::try_from(ssa_vals).expect("Cannot create SSARef");
+        src.src_ref = val.into();
+    }
+
+    fn copy_alu_src(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        let val = match src_type {
+            SrcType::GPR
+            | SrcType::ALU
+            | SrcType::F32
+            | SrcType::F16
+            | SrcType::F16v2
+            | SrcType::I32
+            | SrcType::B32 => self.alloc_ssa_vec(reg_file, 1),
+            SrcType::F64 => self.alloc_ssa_vec(reg_file, 2),
+            SrcType::Pred => self.alloc_ssa_vec(reg_file, 1),
+            _ => panic!("Unknown source type"),
+        };
+
+        if DEBUG.annotate() {
+            self.push_instr(Instr::new(OpAnnotate {
+                annotation: "copy generated by legalizer".into(),
+            }));
+        }
+
+        let old_src_ref =
+            std::mem::replace(&mut src.src_ref, val.clone().into());
+        if val.comps() == 1 {
+            self.copy_to(val[0].into(), old_src_ref.into());
+        } else {
+            match old_src_ref {
+                SrcRef::Imm32(u) => {
+                    // Immediates go in the top bits
+                    self.copy_to(val[0].into(), 0.into());
+                    self.copy_to(val[1].into(), u.into());
+                }
+                SrcRef::CBuf(cb) => {
+                    // CBufs load 8B
+                    self.copy_to(val[0].into(), cb.clone().into());
+                    self.copy_to(val[1].into(), cb.offset(4).into());
+                }
+                SrcRef::SSA(vec) => {
+                    assert!(vec.comps() == 2);
+                    self.copy_to(val[0].into(), vec[0].into());
+                    self.copy_to(val[1].into(), vec[1].into());
+                }
+                _ => panic!("Invalid 64-bit SrcRef"),
+            }
+        }
+    }
+
+    fn copy_alu_src_if_not_reg(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if !src_is_reg(src, reg_file) {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_not_reg_or_imm(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if !src_is_reg(src, reg_file)
+            && !matches!(&src.src_ref, SrcRef::Imm32(_))
+        {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_imm(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if src_is_imm(src) {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_ineg_imm(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        assert!(src_type == SrcType::I32);
+        if src_is_imm(src) && src.src_mod.is_ineg() {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_both_not_reg(
+        &mut self,
+        src1: &Src,
+        src2: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if !src_is_reg(src1, reg_file) && !src_is_reg(src2, reg_file) {
+            self.copy_alu_src(src2, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_and_lower_fmod(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        match src_type {
+            SrcType::F16 | SrcType::F16v2 => {
+                let val = self.alloc_ssa(reg_file);
+                let old_src = std::mem::replace(src, val.into());
+                self.push_op(OpHAdd2 {
+                    dst: val.into(),
+                    srcs: [Src::ZERO.fneg(), old_src],
+                    saturate: false,
+                    ftz: false,
+                    f32: false,
+                });
+            }
+            SrcType::F32 => {
+                let val = self.alloc_ssa(reg_file);
+                let old_src = std::mem::replace(src, val.into());
+                self.push_op(OpFAdd {
+                    dst: val.into(),
+                    srcs: [Src::ZERO.fneg(), old_src],
+                    saturate: false,
+                    rnd_mode: FRndMode::NearestEven,
+                    ftz: false,
+                });
+            }
+            SrcType::F64 => {
+                let val = self.alloc_ssa_vec(reg_file, 2);
+                let old_src = std::mem::replace(src, val.clone().into());
+                self.push_op(OpDAdd {
+                    dst: val.into(),
+                    srcs: [Src::ZERO.fneg(), old_src],
+                    rnd_mode: FRndMode::NearestEven,
+                });
+            }
+            _ => panic!("Invalid ffabs srouce type"),
+        }
+    }
+
+    fn copy_alu_src_and_lower_ineg(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        assert!(src_type == SrcType::I32);
+        let val = self.alloc_ssa(reg_file);
+        let old_src = std::mem::replace(src, val.into());
+        if self.sm() >= 70 {
+            self.push_op(OpIAdd3 {
+                srcs: [Src::ZERO, old_src, Src::ZERO],
+                overflow: [Dst::None, Dst::None],
+                dst: val.into(),
+            });
+        } else {
+            self.push_op(OpIAdd2 {
+                dst: val.into(),
+                carry_out: Dst::None,
+                srcs: [Src::ZERO, old_src],
+            });
+        }
+    }
+
+    fn copy_alu_src_if_fabs(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if src.src_mod.has_fabs() {
+            self.copy_alu_src_and_lower_fmod(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_i20_overflow(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if src.as_imm_not_i20().is_some() {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_alu_src_if_f20_overflow(
+        &mut self,
+        src: &mut Src,
+        reg_file: RegFile,
+        src_type: SrcType,
+    ) {
+        if src.as_imm_not_f20().is_some() {
+            self.copy_alu_src(src, reg_file, src_type);
+        }
+    }
+
+    fn copy_ssa_ref_if_uniform(&mut self, ssa_ref: &mut SSARef) {
+        for ssa in &mut ssa_ref[..] {
+            if ssa.is_uniform() {
+                let warp = self.alloc_ssa(ssa.file().to_warp());
+                self.copy_to(warp.into(), (*ssa).into());
+                *ssa = warp;
+            }
+        }
+    }
+
+    fn copy_src_if_uniform(&mut self, src: &mut Src) {
+        match &mut src.src_ref {
+            SrcRef::SSA(ssa) => self.copy_ssa_ref_if_uniform(ssa),
+            SrcRef::Imm32(_) | SrcRef::False | SrcRef::True | SrcRef::Zero => {}
+            _ => panic!("Unsupported source reference"),
+        }
+    }
+}
+
+pub struct LegalizeBuilder<'a> {
+    b: SSAInstrBuilder<'a>,
+    const_tracker: &'a mut ConstTracker,
+}
+
+impl<'a> LegalizeBuilder<'a> {
+    fn new(
+        sm: &'a ShaderModelInfo,
+        alloc: &'a mut SSAValueAllocator,
+        const_tracker: &'a mut ConstTracker,
+    ) -> Self {
+        LegalizeBuilder {
+            b: SSAInstrBuilder::new(sm, alloc),
+            const_tracker,
+        }
+    }
+
+    pub fn into_vec(self) -> Vec<Instr> {
+        self.b.into_vec()
+    }
+
+    #[allow(dead_code)]
+    pub fn into_mapped_instrs(self) -> MappedInstrs {
+        self.b.into_mapped_instrs()
+    }
+}
+
+impl<'a> Builder for LegalizeBuilder<'a> {
+    fn push_instr(&mut self, instr: Instr) -> &mut Instr {
+        self.b.push_instr(instr)
+    }
+
+    fn sm(&self) -> u8 {
+        self.b.sm()
+    }
+
+    fn copy_to(&mut self, dst: Dst, mut src: Src) {
+        if let Some(ssa_ref) = src.as_ssa() {
+            if let &[ssa_value] = &ssa_ref[..] {
+                if let Some(new_src) = self.const_tracker.get(&ssa_value) {
+                    src = new_src.clone().into();
+                }
+            }
+        };
+        self.b.copy_to(dst, src);
+    }
+}
+
+impl<'a> SSABuilder for LegalizeBuilder<'a> {
+    fn alloc_ssa(&mut self, file: RegFile) -> SSAValue {
+        self.b.alloc_ssa(file)
+    }
+
+    fn alloc_ssa_vec(&mut self, file: RegFile, comps: u8) -> SSARef {
+        self.b.alloc_ssa_vec(file, comps)
+    }
+}
+
+impl LegalizeBuildHelpers for LegalizeBuilder<'_> {}
+
+fn legalize_instr(
+    sm: &ShaderModelInfo,
+    b: &mut LegalizeBuilder,
+    bl: &impl BlockLiveness,
+    block_uniform: bool,
+    pinned: &FxHashSet<SSARef>,
+    ip: usize,
+    instr: &mut Instr,
+) {
+    // Handle a few no-op cases up-front
+    match &instr.op {
+        Op::Annotate(_) => {
+            // OpAnnotate does nothing.  There's nothing to legalize.
+            return;
+        }
+        Op::Undef(_)
+        | Op::PhiSrcs(_)
+        | Op::PhiDsts(_)
+        | Op::Pin(_)
+        | Op::Unpin(_)
+        | Op::RegOut(_) => {
+            // These are implemented by RA and can take pretty much anything
+            // you can throw at them.
+            debug_assert!(instr.pred.is_true());
+            return;
+        }
+        Op::Copy(_) => {
+            // OpCopy is implemented in a lowering pass and can handle anything
+            return;
+        }
+        Op::SrcBar(_) => {
+            // This is turned into a nop by calc_instr_deps
+            return;
+        }
+        Op::Swap(_) | Op::ParCopy(_) => {
+            // These are generated by RA and should not exist yet
+            panic!("Unsupported instruction");
+        }
+        _ => (),
+    }
+
+    if !instr.is_uniform() {
+        b.copy_pred_if_upred(&mut instr.pred);
+    }
+
+    let src_types = instr.src_types();
+    for (i, src) in instr.srcs_mut().iter_mut().enumerate() {
+        if matches!(src.src_ref, SrcRef::Imm32(_)) {
+            // Fold modifiers on Imm32 sources whenever possible.  Not all
+            // instructions suppport modifiers and immediates at the same time.
+            // But leave Zero sources alone as we don't want to make things
+            // immediates that could just be rZ.
+            if let Some(u) = src.as_u32(src_types[i]) {
+                *src = u.into();
+            }
+        }
+        b.copy_src_if_not_same_file(src);
+
+        if !block_uniform {
+            // In non-uniform control-flow, we can't collect uniform vectors so
+            // we need to insert copies to warp regs which we can collect.
+            match &mut src.src_ref {
+                SrcRef::SSA(vec) => {
+                    if vec.is_uniform()
+                        && vec.comps() > 1
+                        && !pinned.contains(vec)
+                    {
+                        b.copy_ssa_ref(vec, vec.file().to_warp());
+                    }
+                }
+                SrcRef::CBuf(CBufRef {
+                    buf: CBuf::BindlessSSA(handle),
+                    ..
+                }) => assert!(pinned.contains(&SSARef::new(handle))),
+                _ => (),
+            }
+        }
+    }
+
+    // OpBreak and OpBSsy impose additional RA constraints
+    let mut legalize_break_bssy = |bar_in: &mut Src, bar_out: &mut Dst| {
+        let bar_in_ssa = bar_in.src_ref.as_ssa().unwrap();
+        if !bar_out.is_none() && bl.is_live_after_ip(&bar_in_ssa[0], ip) {
+            let gpr = b.bmov_to_gpr(bar_in.clone());
+            let tmp = b.bmov_to_bar(gpr.into());
+            *bar_in = tmp.into();
+        }
+    };
+    match &mut instr.op {
+        Op::Break(op) => legalize_break_bssy(&mut op.bar_in, &mut op.bar_out),
+        Op::BSSy(op) => legalize_break_bssy(&mut op.bar_in, &mut op.bar_out),
+        _ => (),
+    }
+
+    sm.legalize_op(b, &mut instr.op);
+
+    let mut vec_src_map: FxHashMap<SSARef, SSARef> = Default::default();
+    let mut vec_comps: FxHashSet<_> = Default::default();
+    for src in instr.srcs_mut() {
+        if let SrcRef::SSA(vec) = &src.src_ref {
+            if vec.comps() == 1 {
+                continue;
+            }
+
+            // If the same vector shows up twice in one instruction, that's
+            // okay. Just make it look the same as the previous source we
+            // fixed up.
+            if let Some(new_vec) = vec_src_map.get(vec) {
+                src.src_ref = new_vec.clone().into();
+                continue;
+            }
+
+            let mut new_vec = vec.clone();
+            for c in 0..vec.comps() {
+                let ssa = vec[usize::from(c)];
+                // If the same SSA value shows up in multiple non-identical
+                // vector sources or as multiple components in the same
+                // source, we need to make a copy so it can get assigned to
+                // multiple different registers.
+                if vec_comps.contains(&ssa) {
+                    let copy = b.alloc_ssa(ssa.file());
+                    b.copy_to(copy.into(), ssa.into());
+                    new_vec[usize::from(c)] = copy;
+                } else {
+                    vec_comps.insert(ssa);
+                }
+            }
+
+            vec_src_map.insert(vec.clone(), new_vec.clone());
+            src.src_ref = new_vec.into();
+        }
+    }
+}
+
+impl Shader<'_> {
+    pub fn legalize(&mut self) {
+        let sm = self.sm;
+        for f in &mut self.functions {
+            let live = SimpleLiveness::for_function(f);
+            let mut pinned: FxHashSet<_> = Default::default();
+            let mut const_tracker = ConstTracker::new();
+
+            for (bi, b) in f.blocks.iter_mut().enumerate() {
+                let bl = live.block_live(bi);
+                let bu = b.uniform;
+
+                let mut instrs = Vec::new();
+                for (ip, mut instr) in b.instrs.drain(..).enumerate() {
+                    match &instr.op {
+                        Op::Pin(pin) => {
+                            if let Dst::SSA(ssa) = &pin.dst {
+                                pinned.insert(ssa.clone());
+                            }
+                        }
+                        Op::Copy(copy) => {
+                            const_tracker.add_copy(copy);
+                        }
+                        _ => (),
+                    }
+
+                    let mut b = LegalizeBuilder::new(
+                        sm,
+                        &mut f.ssa_alloc,
+                        &mut const_tracker,
+                    );
+                    legalize_instr(sm, &mut b, bl, bu, &pinned, ip, &mut instr);
+                    b.push_instr(instr);
+                    instrs.append(&mut b.into_vec());
+                }
+                b.instrs = instrs;
+            }
+        }
+    }
+}