nak: Plumb a ShaderModel trait through everywhere

Instead of scattering number checks everywhere, this lets us actually start splitting code paths. This commit just adds the shader model trait. Later commits will add more methods. Part-of: <https://gitlab.freedesktop.org/mesa/mesa/-/merge_requests/30141>
2024-07-09 10:31:11 -05:00
parent 69bea2b49f
commit e6b8da5427
19 changed files with 133 additions and 73 deletions
--- a/src/nouveau/compiler/nak/api.rs
+++ b/src/nouveau/compiler/nak/api.rs
@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: MIT

 use crate::from_nir::*;
-use crate::ir::{ShaderIoInfo, ShaderStageInfo};
+use crate::ir::{ShaderIoInfo, ShaderModel, ShaderStageInfo};
+use crate::sm50::ShaderModel50;
+use crate::sm70::ShaderModel70;
 use crate::sph;

 use nak_bindings::*;
@@ -264,7 +266,15 @@ pub extern "C" fn nak_compile_shader(
        Some(unsafe { &*fs_key })
    };

-    let mut s = nak_shader_from_nir(nir, nak.sm);
+    let sm: Box<dyn ShaderModel> = if nak.sm >= 70 {
+        Box::new(ShaderModel70::new(nak.sm))
+    } else if nak.sm >= 50 {
+        Box::new(ShaderModel50::new(nak.sm))
+    } else {
+        panic!("Unsupported shader model");
+    };
+
+    let mut s = nak_shader_from_nir(nir, sm.as_ref());

    if DEBUG.print() {
        eprintln!("NAK IR:\n{}", &s);
@@ -323,8 +333,8 @@ pub extern "C" fn nak_compile_shader(

    let info = nak_shader_info {
        stage: nir.info.stage(),
-        sm: s.info.sm,
-        num_gprs: if s.info.sm >= 70 {
+        sm: s.sm.sm(),
+        num_gprs: if s.sm.sm() >= 70 {
            max(4, s.info.num_gprs + 2)
        } else {
            max(4, s.info.num_gprs)
@@ -429,7 +439,7 @@ pub extern "C" fn nak_compile_shader(
            }
            _ => unsafe { std::mem::zeroed() },
        },
-        hdr: sph::encode_header(&s.info, fs_key),
+        hdr: sph::encode_header(sm.as_ref(), &s.info, fs_key),
    };

    let mut asm = String::new();
--- a/src/nouveau/compiler/nak/assign_regs.rs
+++ b/src/nouveau/compiler/nak/assign_regs.rs
@@ -1299,7 +1299,7 @@ impl AssignRegsBlock {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn assign_regs(&mut self) {
        assert!(self.functions.len() == 1);
        let f = &mut self.functions[0];
@@ -1316,7 +1316,7 @@ impl Shader {
        let spill_files =
            [RegFile::UPred, RegFile::Pred, RegFile::UGPR, RegFile::Bar];
        for file in spill_files {
-            let num_regs = file.num_regs(self.info.sm);
+            let num_regs = file.num_regs(self.sm.sm());
            if max_live[file] > num_regs {
                f.spill_values(file, num_regs);

@@ -1336,7 +1336,7 @@ impl Shader {
        let mut gpr_limit = max(max_live[RegFile::GPR], 16);
        let mut total_gprs = gpr_limit + u32::from(tmp_gprs);

-        let max_gprs = RegFile::GPR.num_regs(self.info.sm);
+        let max_gprs = RegFile::GPR.num_regs(self.sm.sm());
        if total_gprs > max_gprs {
            // If we're spilling GPRs, we need to reserve 2 GPRs for OpParCopy
            // lowering because it needs to be able lower Mem copies which
@@ -1364,7 +1364,7 @@ impl Shader {
            if file == RegFile::GPR {
                gpr_limit
            } else {
-                file.num_regs(self.info.sm)
+                file.num_regs(self.sm.sm())
            }
        });

--- a/src/nouveau/compiler/nak/builder.rs
+++ b/src/nouveau/compiler/nak/builder.rs
@@ -701,19 +701,21 @@ pub trait SSABuilder: Builder {
    }
 }

-pub struct InstrBuilder {
+pub struct InstrBuilder<'a> {
    instrs: MappedInstrs,
-    sm: u8,
+    sm: &'a dyn ShaderModel,
 }

-impl InstrBuilder {
-    pub fn new(sm: u8) -> Self {
+impl<'a> InstrBuilder<'a> {
+    pub fn new(sm: &'a dyn ShaderModel) -> Self {
        Self {
            instrs: MappedInstrs::None,
            sm,
        }
    }
+}

+impl InstrBuilder<'_> {
    pub fn as_vec(self) -> Vec<Box<Instr>> {
        match self.instrs {
            MappedInstrs::None => Vec::new(),
@@ -727,30 +729,35 @@ impl InstrBuilder {
    }
 }

-impl Builder for InstrBuilder {
+impl Builder for InstrBuilder<'_> {
    fn push_instr(&mut self, instr: Box<Instr>) -> &mut Instr {
        self.instrs.push(instr);
        self.instrs.last_mut().unwrap().as_mut()
    }

    fn sm(&self) -> u8 {
-        self.sm
+        self.sm.sm()
    }
 }

 pub struct SSAInstrBuilder<'a> {
-    b: InstrBuilder,
+    b: InstrBuilder<'a>,
    alloc: &'a mut SSAValueAllocator,
 }

 impl<'a> SSAInstrBuilder<'a> {
-    pub fn new(sm: u8, alloc: &'a mut SSAValueAllocator) -> Self {
+    pub fn new(
+        sm: &'a dyn ShaderModel,
+        alloc: &'a mut SSAValueAllocator,
+    ) -> Self {
        Self {
            b: InstrBuilder::new(sm),
            alloc: alloc,
        }
    }
+}

+impl SSAInstrBuilder<'_> {
    pub fn as_vec(self) -> Vec<Box<Instr>> {
        self.b.as_vec()
    }
--- a/src/nouveau/compiler/nak/calc_instr_deps.rs
+++ b/src/nouveau/compiler/nak/calc_instr_deps.rs
@@ -361,7 +361,7 @@ impl BarAlloc {
    }
 }

-fn assign_barriers(f: &mut Function, sm: u8) {
+fn assign_barriers(f: &mut Function, sm: &dyn ShaderModel) {
    let mut uses = RegTracker::new_with(&|| RegUse::None);
    let mut deps = DepGraph::new();

@@ -379,7 +379,7 @@ fn assign_barriers(f: &mut Function, sm: u8) {
                    waits.extend_from_slice(u.deps());
                });

-                if instr.has_fixed_latency(sm) {
+                if instr.has_fixed_latency(sm.sm()) {
                    // Delays will cover us here.  We just need to make sure
                    // that we wait on any uses that we consume.
                    uses.for_each_instr_src_mut(instr, |_, u| {
@@ -436,7 +436,7 @@ fn assign_barriers(f: &mut Function, sm: u8) {
                instr.deps.set_yield(true);
            }

-            if instr.has_fixed_latency(sm) {
+            if instr.has_fixed_latency(sm.sm()) {
                continue;
            }

@@ -542,7 +542,7 @@ fn paw_latency(_sm: u8, _write: &Op, _dst_idx: usize) -> u32 {
    13
 }

-fn calc_delays(f: &mut Function, sm: u8) {
+fn calc_delays(f: &mut Function, sm: &dyn ShaderModel) {
    for b in f.blocks.iter_mut().rev() {
        let mut cycle = 0_u32;

@@ -560,7 +560,7 @@ fn calc_delays(f: &mut Function, sm: u8) {

        for ip in (0..b.instrs.len()).rev() {
            let instr = &b.instrs[ip];
-            let mut min_start = cycle + exec_latency(sm, &instr.op);
+            let mut min_start = cycle + exec_latency(sm.sm(), &instr.op);
            if let Some(bar) = instr.deps.rd_bar() {
                min_start = max(min_start, bars[usize::from(bar)] + 2);
            }
@@ -578,7 +578,7 @@ fn calc_delays(f: &mut Function, sm: u8) {
                RegUse::Write((w_ip, w_dst_idx)) => {
                    let s = instr_cycle[*w_ip]
                        + waw_latency(
-                            sm,
+                            sm.sm(),
                            &instr.op,
                            i,
                            &b.instrs[*w_ip].op,
@@ -590,10 +590,10 @@ fn calc_delays(f: &mut Function, sm: u8) {
                    for (r_ip, r_src_idx) in reads {
                        let c = instr_cycle[*r_ip];
                        let s = if *r_src_idx == usize::MAX {
-                            c + paw_latency(sm, &instr.op, i)
+                            c + paw_latency(sm.sm(), &instr.op, i)
                        } else {
                            c + raw_latency(
-                                sm,
+                                sm.sm(),
                                &instr.op,
                                i,
                                &b.instrs[*r_ip].op,
@@ -609,7 +609,7 @@ fn calc_delays(f: &mut Function, sm: u8) {
                RegUse::Write((w_ip, w_dst_idx)) => {
                    let s = instr_cycle[*w_ip]
                        + war_latency(
-                            sm,
+                            sm.sm(),
                            &instr.op,
                            i,
                            &b.instrs[*w_ip].op,
@@ -657,7 +657,7 @@ fn calc_delays(f: &mut Function, sm: u8) {
        if matches!(instr.op, Op::SrcBar(_)) {
            instr.op = Op::Nop(OpNop { label: None });
            MappedInstrs::One(instr)
-        } else if exec_latency(sm, &instr.op) > 1 {
+        } else if exec_latency(sm.sm(), &instr.op) > 1 {
            let mut nop = Instr::new_boxed(OpNop { label: None });
            nop.deps.set_delay(2);
            MappedInstrs::Many(vec![instr, nop])
@@ -667,7 +667,7 @@ fn calc_delays(f: &mut Function, sm: u8) {
    });
 }

-impl Shader {
+impl Shader<'_> {
    pub fn assign_deps_serial(&mut self) {
        for f in &mut self.functions {
            for b in &mut f.blocks.iter_mut().rev() {
@@ -704,8 +704,8 @@ impl Shader {
            self.assign_deps_serial();
        } else {
            for f in &mut self.functions {
-                assign_barriers(f, self.info.sm);
-                calc_delays(f, self.info.sm);
+                assign_barriers(f, self.sm);
+                calc_delays(f, self.sm);
            }
        }
    }
--- a/src/nouveau/compiler/nak/from_nir.rs
+++ b/src/nouveau/compiler/nak/from_nir.rs
@@ -18,9 +18,8 @@ use std::cmp::max;
 use std::collections::{HashMap, HashSet};
 use std::ops::Index;

-fn init_info_from_nir(nir: &nir_shader, sm: u8) -> ShaderInfo {
+fn init_info_from_nir(nir: &nir_shader) -> ShaderInfo {
    ShaderInfo {
-        sm: sm,
        num_gprs: 0,
        num_barriers: 0,
        slm_size: nir.scratch_size,
@@ -233,6 +232,7 @@ impl Index<FloatType> for ShaderFloatControls {

 struct ShaderFromNir<'a> {
    nir: &'a nir_shader,
+    sm: &'a dyn ShaderModel,
    info: ShaderInfo,
    float_ctl: ShaderFloatControls,
    cfg: CFGBuilder<u32, BasicBlock>,
@@ -247,10 +247,11 @@ struct ShaderFromNir<'a> {
 }

 impl<'a> ShaderFromNir<'a> {
-    fn new(nir: &'a nir_shader, sm: u8) -> Self {
+    fn new(nir: &'a nir_shader, sm: &'a dyn ShaderModel) -> Self {
        Self {
            nir: nir,
-            info: init_info_from_nir(nir, sm),
+            sm: sm,
+            info: init_info_from_nir(nir),
            float_ctl: ShaderFloatControls::from_nir(nir),
            cfg: CFGBuilder::new(),
            label_alloc: LabelAllocator::new(),
@@ -1899,7 +1900,7 @@ impl<'a> ShaderFromNir<'a> {
        &mut self,
        access: gl_access_qualifier,
    ) -> MemEvictionPriority {
-        if self.info.sm >= 70 && access & ACCESS_NON_TEMPORAL != 0 {
+        if self.sm.sm() >= 70 && access & ACCESS_NON_TEMPORAL != 0 {
            MemEvictionPriority::First
        } else {
            MemEvictionPriority::Normal
@@ -2985,7 +2986,7 @@ impl<'a> ShaderFromNir<'a> {
            nir_intrinsic_final_primitive_nv => {
                let handle = self.get_src(&srcs[0]);

-                if self.info.sm >= 70 {
+                if self.sm.sm() >= 70 {
                    b.push_op(OpOutFinal { handle: handle });
                }
            }
@@ -3139,7 +3140,8 @@ impl<'a> ShaderFromNir<'a> {
        phi_map: &mut PhiAllocMap,
        nb: &nir_block,
    ) {
-        let mut b = SSAInstrBuilder::new(self.info.sm, ssa_alloc);
+        let sm = self.sm;
+        let mut b = SSAInstrBuilder::new(sm, ssa_alloc);

        if nb.index == 0 && self.nir.info.shared_size > 0 {
            // The blob seems to always do a BSYNC before accessing shared
@@ -3187,7 +3189,7 @@ impl<'a> ShaderFromNir<'a> {
            }

            let uniform = !nb.divergent
-                && self.info.sm >= 75
+                && self.sm.sm() >= 75
                && !DEBUG.no_ugpr()
                && !np.def.divergent;

@@ -3226,7 +3228,7 @@ impl<'a> ShaderFromNir<'a> {
            }

            let uniform = !nb.divergent
-                && self.info.sm >= 75
+                && self.sm.sm() >= 75
                && !DEBUG.no_ugpr()
                && ni.def().is_some_and(|d| !d.divergent);
            let mut b = UniformBuilder::new(&mut b, uniform);
@@ -3436,7 +3438,7 @@ impl<'a> ShaderFromNir<'a> {
        f
    }

-    pub fn parse_shader(mut self) -> Shader {
+    pub fn parse_shader(mut self) -> Shader<'a> {
        let mut functions = Vec::new();
        for nf in self.nir.iter_functions() {
            if let Some(nfi) = nf.get_impl() {
@@ -3458,12 +3460,16 @@ impl<'a> ShaderFromNir<'a> {
        }

        Shader {
+            sm: self.sm,
            info: self.info,
            functions: functions,
        }
    }
 }

-pub fn nak_shader_from_nir(ns: &nir_shader, sm: u8) -> Shader {
+pub fn nak_shader_from_nir<'a>(
+    ns: &'a nir_shader,
+    sm: &'a dyn ShaderModel,
+) -> Shader<'a> {
    ShaderFromNir::new(ns, sm).parse_shader()
 }
--- a/src/nouveau/compiler/nak/ir.rs
+++ b/src/nouveau/compiler/nak/ir.rs
@@ -6316,7 +6316,6 @@ pub enum ShaderIoInfo {

 #[derive(Debug)]
 pub struct ShaderInfo {
-    pub sm: u8,
    pub num_gprs: u8,
    pub num_barriers: u8,
    pub slm_size: u32,
@@ -6327,12 +6326,17 @@ pub struct ShaderInfo {
    pub io: ShaderIoInfo,
 }

-pub struct Shader {
+pub trait ShaderModel {
+    fn sm(&self) -> u8;
+}
+
+pub struct Shader<'a> {
+    pub sm: &'a dyn ShaderModel,
    pub info: ShaderInfo,
    pub functions: Vec<Function>,
 }

-impl Shader {
+impl Shader<'_> {
    pub fn for_each_instr(&self, f: &mut impl FnMut(&Instr)) {
        for func in &self.functions {
            for b in &func.blocks {
@@ -6386,7 +6390,7 @@ impl Shader {
    }
 }

-impl fmt::Display for Shader {
+impl fmt::Display for Shader<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for func in &self.functions {
            write!(f, "{}", func)?;
--- a/src/nouveau/compiler/nak/legalize.rs
+++ b/src/nouveau/compiler/nak/legalize.rs
@@ -1052,9 +1052,9 @@ fn legalize_instr(
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn legalize(&mut self) {
-        let sm = self.info.sm;
+        let sm = self.sm;
        for f in &mut self.functions {
            let live = SimpleLiveness::for_function(f);
            let mut pinned = HashSet::new();
--- a/src/nouveau/compiler/nak/lower_copy_swap.rs
+++ b/src/nouveau/compiler/nak/lower_copy_swap.rs
@@ -238,7 +238,7 @@ impl LowerCopySwap {
    }

    fn run(&mut self, s: &mut Shader) {
-        let sm = s.info.sm;
+        let sm = s.sm;
        s.map_instrs(|instr: Box<Instr>, _| -> MappedInstrs {
            match instr.op {
                Op::R2UR(r2ur) => {
@@ -283,7 +283,7 @@ impl LowerCopySwap {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn lower_copy_swap(&mut self) {
        let mut pass = LowerCopySwap::new(self.info.slm_size);
        pass.run(self);
--- a/src/nouveau/compiler/nak/lower_par_copies.rs
+++ b/src/nouveau/compiler/nak/lower_par_copies.rs
@@ -89,7 +89,7 @@ fn cycle_use_swap(pc: &OpParCopy, file: RegFile) -> bool {
    }
 }

-fn lower_par_copy(pc: OpParCopy, sm: u8) -> MappedInstrs {
+fn lower_par_copy(pc: OpParCopy, sm: &dyn ShaderModel) -> MappedInstrs {
    let mut graph = CopyGraph::new();
    let mut vals = Vec::new();
    let mut reg_to_idx = HashMap::new();
@@ -250,9 +250,9 @@ fn lower_par_copy(pc: OpParCopy, sm: u8) -> MappedInstrs {
    b.as_mapped_instrs()
 }

-impl Shader {
+impl Shader<'_> {
    pub fn lower_par_copies(&mut self) {
-        let sm = self.info.sm;
+        let sm = self.sm;
        self.map_instrs(|instr, _| -> MappedInstrs {
            match instr.op {
                Op::ParCopy(pc) => {
--- a/src/nouveau/compiler/nak/opt_bar_prop.rs
+++ b/src/nouveau/compiler/nak/opt_bar_prop.rs
@@ -269,7 +269,7 @@ impl BarPropPass {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_bar_prop(&mut self) {
        for f in &mut self.functions {
            BarPropPass::new().run(f);
--- a/src/nouveau/compiler/nak/opt_copy_prop.rs
+++ b/src/nouveau/compiler/nak/opt_copy_prop.rs
@@ -722,7 +722,7 @@ impl CopyPropPass {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_copy_prop(&mut self) {
        for f in &mut self.functions {
            CopyPropPass::new().run(f);
--- a/src/nouveau/compiler/nak/opt_dce.rs
+++ b/src/nouveau/compiler/nak/opt_dce.rs
@@ -182,7 +182,7 @@ impl Function {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_dce(&mut self) {
        for f in &mut self.functions {
            f.opt_dce();
--- a/src/nouveau/compiler/nak/opt_jump_thread.rs
+++ b/src/nouveau/compiler/nak/opt_jump_thread.rs
@@ -143,7 +143,7 @@ impl Function {
    }
 }

-impl Shader {
+impl Shader<'_> {
    /// A simple jump threading pass
    ///
    /// Note that this can introduce critical edges, so it cannot be run before RA
--- a/src/nouveau/compiler/nak/opt_lop.rs
+++ b/src/nouveau/compiler/nak/opt_lop.rs
@@ -266,7 +266,7 @@ impl LopPass {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_lop(&mut self) {
        for f in &mut self.functions {
            let mut pass = LopPass::new(f);
--- a/src/nouveau/compiler/nak/opt_out.rs
+++ b/src/nouveau/compiler/nak/opt_out.rs
@@ -37,7 +37,7 @@ fn try_combine_outs(emit: &mut Instr, cut: &Instr) -> bool {
    true
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_out(&mut self) {
        if !matches!(self.info.stage, ShaderStageInfo::Geometry(_)) {
            return;
--- a/src/nouveau/compiler/nak/opt_uniform_instrs.rs
+++ b/src/nouveau/compiler/nak/opt_uniform_instrs.rs
@@ -5,11 +5,11 @@ use crate::ir::*;
 use std::collections::HashMap;

 fn should_lower_to_warp(
-    sm: u8,
+    sm: &dyn ShaderModel,
    instr: &Instr,
    r2ur: &HashMap<SSAValue, SSAValue>,
 ) -> bool {
-    if !instr.can_be_uniform(sm) {
+    if !instr.can_be_uniform(sm.sm()) {
        return true;
    }

@@ -49,9 +49,9 @@ fn propagate_r2ur(
    progress
 }

-impl Shader {
+impl Shader<'_> {
    pub fn opt_uniform_instrs(&mut self) {
-        let sm = self.info.sm;
+        let sm = self.sm;
        let mut r2ur = HashMap::new();
        let mut propagated_r2ur = false;
        self.map_instrs(|mut instr, alloc| {
--- a/src/nouveau/compiler/nak/sm50.rs
+++ b/src/nouveau/compiler/nak/sm50.rs
@@ -7,6 +7,23 @@ use bitview::*;
 use std::collections::HashMap;
 use std::ops::Range;

+pub struct ShaderModel50 {
+    sm: u8,
+}
+
+impl ShaderModel50 {
+    pub fn new(sm: u8) -> Self {
+        assert!(sm >= 50 && sm < 70);
+        Self { sm }
+    }
+}
+
+impl ShaderModel for ShaderModel50 {
+    fn sm(&self) -> u8 {
+        self.sm
+    }
+}
+
 impl Src {
    fn is_reg_or_zero(&self) -> bool {
        matches!(self.src_ref, SrcRef::Zero | SrcRef::Reg(_))
@@ -2179,7 +2196,7 @@ fn encode_instr(
    res.inst
 }

-impl Shader {
+impl Shader<'_> {
    pub fn encode_sm50(&self) -> Vec<u32> {
        assert!(self.functions.len() == 1);
        let func = &self.functions[0];
@@ -2214,7 +2231,7 @@ impl Shader {
                let instr0 = encode_instr(
                    0,
                    instrs_iter.next(),
-                    self.info.sm,
+                    self.sm.sm(),
                    &labels,
                    &mut ip,
                    &mut sched_instr,
@@ -2222,7 +2239,7 @@ impl Shader {
                let instr1 = encode_instr(
                    1,
                    instrs_iter.next(),
-                    self.info.sm,
+                    self.sm.sm(),
                    &labels,
                    &mut ip,
                    &mut sched_instr,
@@ -2230,7 +2247,7 @@ impl Shader {
                let instr2 = encode_instr(
                    2,
                    instrs_iter.next(),
-                    self.info.sm,
+                    self.sm.sm(),
                    &labels,
                    &mut ip,
                    &mut sched_instr,
--- a/src/nouveau/compiler/nak/sm70.rs
+++ b/src/nouveau/compiler/nak/sm70.rs
@@ -7,6 +7,23 @@ use bitview::*;
 use std::collections::HashMap;
 use std::ops::Range;

+pub struct ShaderModel70 {
+    sm: u8,
+}
+
+impl ShaderModel70 {
+    pub fn new(sm: u8) -> Self {
+        assert!(sm >= 70);
+        Self { sm }
+    }
+}
+
+impl ShaderModel for ShaderModel70 {
+    fn sm(&self) -> u8 {
+        self.sm
+    }
+}
+
 struct ALURegRef {
    pub reg: RegRef,
    pub abs: bool,
@@ -2570,7 +2587,7 @@ impl SM70Instr {
    }
 }

-impl Shader {
+impl Shader<'_> {
    pub fn encode_sm70(&self) -> Vec<u32> {
        assert!(self.functions.len() == 1);
        let func = &self.functions[0];
@@ -2594,7 +2611,7 @@ impl Shader {
            for instr in &b.instrs {
                let e = SM70Instr::encode(
                    instr,
-                    self.info.sm,
+                    self.sm.sm(),
                    encoded.len(),
                    &labels,
                );
--- a/src/nouveau/compiler/nak/sph.rs
+++ b/src/nouveau/compiler/nak/sph.rs
@@ -4,7 +4,7 @@
 extern crate bitview;
 extern crate nvidia_headers;

-use crate::ir::{ShaderInfo, ShaderIoInfo, ShaderStageInfo};
+use crate::ir::{ShaderInfo, ShaderIoInfo, ShaderModel, ShaderStageInfo};
 use bitview::{
    BitMutView, BitMutViewable, BitView, BitViewable, SetBit, SetField,
    SetFieldU64,
@@ -462,6 +462,7 @@ impl ShaderProgramHeader {
 }

 pub fn encode_header(
+    sm: &dyn ShaderModel,
    shader_info: &ShaderInfo,
    fs_key: Option<&nak_fs_key>,
 ) -> [u32; CURRENT_MAX_SHADER_HEADER_SIZE] {
@@ -469,10 +470,8 @@ pub fn encode_header(
        return [0_u32; CURRENT_MAX_SHADER_HEADER_SIZE];
    }

-    let mut sph = ShaderProgramHeader::new(
-        ShaderType::from(&shader_info.stage),
-        shader_info.sm,
-    );
+    let mut sph =
+        ShaderProgramHeader::new(ShaderType::from(&shader_info.stage), sm.sm());

    sph.set_sass_version(1);
    sph.set_does_load_or_store(shader_info.uses_global_mem);