llvm-native-core 0.1.15

//! AMDGPU Register Information — complete GCN/RDNA register file.
//!
//! Defines every architectural register in the AMDGPU ISA:
//! - SGPRs (Scalar GPRs): s0–s103 (GFX9+), s0–s101 (GFX8)
//! - VGPRs (Vector GPRs): v0–v255
//! - AGPRs (Accumulation GPRs): a0–a255 (GFX908+)
//! - Special registers: VCC, VCC_LO, VCC_HI, EXEC, EXEC_LO, EXEC_HI, SCC
//! - Hardware registers: HW_REG_MODE, HW_REG_STATUS, HW_REG_TRAPSTS
//! - Trap temporaries: TTMP0–TTMP15
//! - FLAT scratch: FLAT_SCR_LO, FLAT_SCR_HI
//! - XNACK mask: XNACK_MASK_LO, XNACK_MASK_HI
//! - M0 (memory descriptor)
//!
//! Register encoding uses flat u16 IDs starting at 21000.
//!
//! Clean-room reconstruction from AMD GCN3/RDNA/GFX ISA manuals.
//! Zero LLVM source code consultation.

use std::collections::HashMap;
use std::fmt;

// ═══════════════════════════════════════════════════════════════════════════
// Register IDs — flat u16 namespace starting at 21000
// ═══════════════════════════════════════════════════════════════════════════

// SGPRs: scalar general-purpose registers (s0–s103)
pub const SGPR0: u16 = 21000;
pub const SGPR1: u16 = 21001;
pub const SGPR2: u16 = 21002;
pub const SGPR3: u16 = 21003;
pub const SGPR4: u16 = 21004;
pub const SGPR5: u16 = 21005;
pub const SGPR6: u16 = 21006;
pub const SGPR7: u16 = 21007;
pub const SGPR8: u16 = 21008;
pub const SGPR9: u16 = 21009;
pub const SGPR10: u16 = 21010;
pub const SGPR11: u16 = 21011;
pub const SGPR12: u16 = 21012;
pub const SGPR13: u16 = 21013;
pub const SGPR14: u16 = 21014;
pub const SGPR15: u16 = 21015;
pub const SGPR16: u16 = 21016;
pub const SGPR32: u16 = 21032;
pub const SGPR48: u16 = 21048;
pub const SGPR64: u16 = 21064;
pub const SGPR80: u16 = 21080;
pub const SGPR96: u16 = 21096;
pub const SGPR100: u16 = 21100;
pub const SGPR101: u16 = 21101;
pub const SGPR102: u16 = 21102;
pub const SGPR103: u16 = 21103;

// VGPRs: vector general-purpose registers (v0–v255)
pub const VGPR0: u16 = 21200;
pub const VGPR1: u16 = 21201;
pub const VGPR2: u16 = 21202;
pub const VGPR3: u16 = 21203;
pub const VGPR4: u16 = 21204;
pub const VGPR5: u16 = 21205;
pub const VGPR6: u16 = 21206;
pub const VGPR7: u16 = 21207;
pub const VGPR8: u16 = 21208;
pub const VGPR16: u16 = 21216;
pub const VGPR32: u16 = 21232;
pub const VGPR64: u16 = 21264;
pub const VGPR128: u16 = 21328;
pub const VGPR192: u16 = 21392;
pub const VGPR240: u16 = 21440;
pub const VGPR252: u16 = 21452;
pub const VGPR253: u16 = 21453;
pub const VGPR254: u16 = 21454;
pub const VGPR255: u16 = 21455;

// AGPRs: accumulation vector registers (GFX908+, a0–a255)
pub const AGPR0: u16 = 21500;
pub const AGPR1: u16 = 21501;
pub const AGPR32: u16 = 21532;
pub const AGPR255: u16 = 21755;

// Special registers
pub const VCC_LO: u16 = 21800;
pub const VCC_HI: u16 = 21801;
pub const VCC: u16 = 21802; // alias for VCC_LO:VCC_HI pair
pub const EXEC_LO: u16 = 21810;
pub const EXEC_HI: u16 = 21811;
pub const EXEC: u16 = 21812; // alias for EXEC_LO:EXEC_HI pair
pub const SCC: u16 = 21820; // scalar condition code (1-bit)

// Hardware registers (accessed via s_getreg/s_setreg)
pub const HW_REG_MODE: u16 = 21830;
pub const HW_REG_STATUS: u16 = 21831;
pub const HW_REG_TRAPSTS: u16 = 21832;
pub const HW_REG_HW_ID: u16 = 21833;
pub const HW_REG_GPR_ALLOC: u16 = 21834;
pub const HW_REG_LDS_ALLOC: u16 = 21835;
pub const HW_REG_IB_STS: u16 = 21836;

// Memory descriptor / addressing registers
pub const M0: u16 = 21840;
pub const FLAT_SCR_LO: u16 = 21841;
pub const FLAT_SCR_HI: u16 = 21842;
pub const XNACK_MASK_LO: u16 = 21843;
pub const XNACK_MASK_HI: u16 = 21844;

// Trap temporaries (TTMP0–TTMP15)
pub const TTMP0: u16 = 21850;
pub const TTMP1: u16 = 21851;
pub const TTMP2: u16 = 21852;
pub const TTMP3: u16 = 21853;
pub const TTMP4: u16 = 21854;
pub const TTMP5: u16 = 21855;
pub const TTMP6: u16 = 21856;
pub const TTMP7: u16 = 21857;
pub const TTMP8: u16 = 21858;
pub const TTMP9: u16 = 21859;
pub const TTMP10: u16 = 21860;
pub const TTMP11: u16 = 21861;
pub const TTMP12: u16 = 21862;
pub const TTMP13: u16 = 21863;
pub const TTMP14: u16 = 21864;
pub const TTMP15: u16 = 21865;

// Null register (writes discarded, reads return 0)
pub const SGPR_NULL: u16 = 21870;

// ═══════════════════════════════════════════════════════════════════════════
// Register Class
// ═══════════════════════════════════════════════════════════════════════════

/// AMDGPU register classes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum AmdgpuRegClass {
    /// Scalar general-purpose registers (SGPRs): 32-bit each
    SGPR32,
    /// SGPR pair (64-bit): s[N:N+1]
    SGPR64,
    /// SGPR quad (128-bit): s[N:N+3]
    SGPR128,
    /// Vector general-purpose registers (VGPRs): 32-bit each
    VGPR32,
    /// VGPR pair (64-bit): v[N:N+1]
    VGPR64,
    /// VGPR quad (128-bit): v[N:N+3]
    VGPR128,
    /// Accumulation vector registers (AGPRs): 32-bit
    AGPR32,
    /// Special registers: VCC, EXEC, SCC
    SpecialReg,
    /// Hardware registers
    HwReg,
    /// Trap temporaries
    TtmpReg,
}

impl AmdgpuRegClass {
    /// Size of one register in this class (in bits).
    pub fn reg_size_bits(&self) -> u32 {
        match self {
            AmdgpuRegClass::SGPR32 | AmdgpuRegClass::VGPR32 | AmdgpuRegClass::AGPR32 => 32,
            AmdgpuRegClass::SGPR64 | AmdgpuRegClass::VGPR64 => 64,
            AmdgpuRegClass::SGPR128 | AmdgpuRegClass::VGPR128 => 128,
            AmdgpuRegClass::SpecialReg => 64,
            AmdgpuRegClass::HwReg => 32,
            AmdgpuRegClass::TtmpReg => 32,
        }
    }

    /// String name for this register class.
    pub fn name(&self) -> &'static str {
        match self {
            AmdgpuRegClass::SGPR32 => "SGPR_32",
            AmdgpuRegClass::SGPR64 => "SGPR_64",
            AmdgpuRegClass::SGPR128 => "SGPR_128",
            AmdgpuRegClass::VGPR32 => "VGPR_32",
            AmdgpuRegClass::VGPR64 => "VGPR_64",
            AmdgpuRegClass::VGPR128 => "VGPR_128",
            AmdgpuRegClass::AGPR32 => "AGPR_32",
            AmdgpuRegClass::SpecialReg => "SPECIAL",
            AmdgpuRegClass::HwReg => "HW_REG",
            AmdgpuRegClass::TtmpReg => "TTMP",
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Register Descriptor
// ═══════════════════════════════════════════════════════════════════════════

/// Describes a single AMDGPU register.
#[derive(Debug, Clone)]
pub struct AmdgpuRegisterDesc {
    /// Unique register ID.
    pub reg_id: u16,
    /// Assembly name (e.g., "s0", "v1", "vcc").
    pub name: String,
    /// Register class.
    pub class: AmdgpuRegClass,
    /// Index within the register class (0-based).
    pub index: u32,
    /// Whether this register is a 64-bit pair.
    pub is_pair: bool,
    /// Whether this register is a 128-bit quad.
    pub is_quad: bool,
    /// Whether this register is reserved (not allocatable).
    pub is_reserved: bool,
}

impl AmdgpuRegisterDesc {
    pub fn new(reg_id: u16, name: &str, class: AmdgpuRegClass, index: u32) -> Self {
        Self {
            reg_id,
            name: name.to_string(),
            class,
            index,
            is_pair: false,
            is_quad: false,
            is_reserved: false,
        }
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Register Info Database
// ═══════════════════════════════════════════════════════════════════════════

/// Complete AMDGPU register information database.
pub struct AmdgpuRegisterInfo {
    /// All registers by ID.
    registers: HashMap<u16, AmdgpuRegisterDesc>,
    /// Lookup by name.
    by_name: HashMap<String, u16>,
}

impl AmdgpuRegisterInfo {
    pub fn new() -> Self {
        let mut info = Self {
            registers: HashMap::new(),
            by_name: HashMap::new(),
        };
        info.populate();
        info
    }

    fn populate(&mut self) {
        // SGPRs (s0–s103)
        for i in 0..104u32 {
            let id = SGPR0 + i as u16;
            let name = format!("s{}", i);
            let desc = AmdgpuRegisterDesc::new(id, &name, AmdgpuRegClass::SGPR32, i);
            self.registers.insert(id, desc);
            self.by_name.insert(name, id);
        }

        // VGPRs (v0–v255)
        for i in 0..256u32 {
            let id = VGPR0 + i as u16;
            let name = format!("v{}", i);
            let desc = AmdgpuRegisterDesc::new(id, &name, AmdgpuRegClass::VGPR32, i);
            self.registers.insert(id, desc);
            self.by_name.insert(name, id);
        }

        // AGPRs (a0–a255)
        for i in 0..256u32 {
            let id = AGPR0 + i as u16;
            let name = format!("a{}", i);
            let desc = AmdgpuRegisterDesc::new(id, &name, AmdgpuRegClass::AGPR32, i);
            self.registers.insert(id, desc);
            self.by_name.insert(name, id);
        }

        // Special registers
        self.add_special(VCC_LO, "vcc_lo", AmdgpuRegClass::SpecialReg);
        self.add_special(VCC_HI, "vcc_hi", AmdgpuRegClass::SpecialReg);
        self.add_special(VCC, "vcc", AmdgpuRegClass::SpecialReg);
        self.add_special(EXEC_LO, "exec_lo", AmdgpuRegClass::SpecialReg);
        self.add_special(EXEC_HI, "exec_hi", AmdgpuRegClass::SpecialReg);
        self.add_special(EXEC, "exec", AmdgpuRegClass::SpecialReg);
        self.add_special(SCC, "scc", AmdgpuRegClass::SpecialReg);
        self.add_special(M0, "m0", AmdgpuRegClass::SpecialReg);
        self.add_special(FLAT_SCR_LO, "flat_scr_lo", AmdgpuRegClass::SpecialReg);
        self.add_special(FLAT_SCR_HI, "flat_scr_hi", AmdgpuRegClass::SpecialReg);
        self.add_special(XNACK_MASK_LO, "xnack_mask_lo", AmdgpuRegClass::SpecialReg);
        self.add_special(XNACK_MASK_HI, "xnack_mask_hi", AmdgpuRegClass::SpecialReg);

        // HW registers
        self.add_special(HW_REG_MODE, "hw_reg_mode", AmdgpuRegClass::HwReg);
        self.add_special(HW_REG_STATUS, "hw_reg_status", AmdgpuRegClass::HwReg);
        self.add_special(HW_REG_TRAPSTS, "hw_reg_trapsts", AmdgpuRegClass::HwReg);
        self.add_special(HW_REG_HW_ID, "hw_reg_hw_id", AmdgpuRegClass::HwReg);
        self.add_special(HW_REG_GPR_ALLOC, "hw_reg_gpr_alloc", AmdgpuRegClass::HwReg);
        self.add_special(HW_REG_LDS_ALLOC, "hw_reg_lds_alloc", AmdgpuRegClass::HwReg);

        // Trap temporaries
        for i in 0..16u32 {
            let id = TTMP0 + i as u16;
            let name = format!("ttmp{}", i);
            let desc = AmdgpuRegisterDesc::new(id, &name, AmdgpuRegClass::TtmpReg, i);
            self.registers.insert(id, desc);
            self.by_name.insert(name, id);
        }

        // Null register
        self.add_special(SGPR_NULL, "null", AmdgpuRegClass::SGPR32);
    }

    fn add_special(&mut self, id: u16, name: &str, class: AmdgpuRegClass) {
        let desc = AmdgpuRegisterDesc::new(id, name, class, 0);
        self.registers.insert(id, desc);
        self.by_name.insert(name.to_string(), id);
    }

    /// Look up a register by ID.
    pub fn lookup(&self, id: u16) -> Option<&AmdgpuRegisterDesc> {
        self.registers.get(&id)
    }

    /// Look up a register by assembly name.
    pub fn lookup_by_name(&self, name: &str) -> Option<&AmdgpuRegisterDesc> {
        self.by_name.get(name).and_then(|id| self.registers.get(id))
    }

    /// Get the 7-bit encoding of an SGPR for instruction fields.
    pub fn encode_sgpr(&self, id: u16) -> u8 {
        if (SGPR0..=SGPR103).contains(&id) {
            (id - SGPR0) as u8
        } else {
            0
        }
    }

    /// Get the 8-bit encoding of a VGPR for instruction fields.
    pub fn encode_vgpr(&self, id: u16) -> u8 {
        if (VGPR0..=VGPR255).contains(&id) {
            (id - VGPR0) as u8
        } else {
            0
        }
    }

    /// Get the total number of registers.
    pub fn count(&self) -> usize {
        self.registers.len()
    }

    /// Get all SGPR IDs.
    pub fn sgpr_ids(&self) -> Vec<u16> {
        (SGPR0..=SGPR103).collect()
    }

    /// Get all VGPR IDs.
    pub fn vgpr_ids(&self) -> Vec<u16> {
        (VGPR0..=VGPR255).collect()
    }
}

impl Default for AmdgpuRegisterInfo {
    fn default() -> Self {
        Self::new()
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Register Display
// ═══════════════════════════════════════════════════════════════════════════

impl fmt::Display for AmdgpuRegClass {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

// ═══════════════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sgpr_count() {
        let info = AmdgpuRegisterInfo::new();
        let sgprs = info.sgpr_ids();
        assert_eq!(sgprs.len(), 104);
    }

    #[test]
    fn test_vgpr_count() {
        let info = AmdgpuRegisterInfo::new();
        let vgprs = info.vgpr_ids();
        assert_eq!(vgprs.len(), 256);
    }

    #[test]
    fn test_lookup_sgpr() {
        let info = AmdgpuRegisterInfo::new();
        let desc = info.lookup(SGPR0).unwrap();
        assert_eq!(desc.name, "s0");
        assert_eq!(desc.class, AmdgpuRegClass::SGPR32);
    }

    #[test]
    fn test_lookup_vgpr() {
        let info = AmdgpuRegisterInfo::new();
        let desc = info.lookup(VGPR255).unwrap();
        assert_eq!(desc.name, "v255");
    }

    #[test]
    fn test_lookup_by_name() {
        let info = AmdgpuRegisterInfo::new();
        let desc = info.lookup_by_name("vcc").unwrap();
        assert_eq!(desc.class, AmdgpuRegClass::SpecialReg);
    }

    #[test]
    fn test_encode_sgpr() {
        let info = AmdgpuRegisterInfo::new();
        assert_eq!(info.encode_sgpr(SGPR0), 0);
        assert_eq!(info.encode_sgpr(SGPR7), 7);
    }

    #[test]
    fn test_encode_vgpr() {
        let info = AmdgpuRegisterInfo::new();
        assert_eq!(info.encode_vgpr(VGPR0), 0);
        assert_eq!(info.encode_vgpr(VGPR255), 255);
    }

    #[test]
    fn test_total_register_count() {
        let info = AmdgpuRegisterInfo::new();
        // 104 SGPRs + 256 VGPRs + 256 AGPRs + 13 special + 6 hw + 16 ttmp + 1 null = 652
        assert_eq!(info.count(), 651);
    }

    #[test]
    fn test_reg_class_sizes() {
        assert_eq!(AmdgpuRegClass::SGPR32.reg_size_bits(), 32);
        assert_eq!(AmdgpuRegClass::SGPR64.reg_size_bits(), 64);
        assert_eq!(AmdgpuRegClass::VGPR32.reg_size_bits(), 32);
        assert_eq!(AmdgpuRegClass::AGPR32.reg_size_bits(), 32);
    }

    #[test]
    fn test_special_register_names() {
        let info = AmdgpuRegisterInfo::new();
        assert!(info.lookup_by_name("exec").is_some());
        assert!(info.lookup_by_name("m0").is_some());
        assert!(info.lookup_by_name("scc").is_some());
        assert!(info.lookup_by_name("null").is_some());
    }

    #[test]
    fn test_ttmp_count() {
        let info = AmdgpuRegisterInfo::new();
        let mut count = 0;
        for i in 0..16 {
            if info.lookup_by_name(&format!("ttmp{}", i)).is_some() {
                count += 1;
            }
        }
        assert_eq!(count, 16);
    }
}