llvm-native-core 0.1.10

//! NVPTX Register Information — virtual SSA registers and special
//! registers for the NVIDIA PTX virtual ISA.
//!
//! Clean-room behavioral reconstruction from the NVIDIA PTX ISA
//! Reference Manual. Zero LLVM source code consultation.
//!
//! PTX uses unlimited virtual registers in SSA form. There are no
//! fixed physical registers. Instead, PTX defines special registers
//! that provide access to thread/block/grid identifiers, performance
//! monitors, and clock counters. Predicate registers are also virtual.

// ============================================================================
// Special Register IDs (base 10000)
// ============================================================================

/// Base offset for NVPTX register IDs.
pub const NVPTX_REG_BASE: u16 = 10000;

/// Thread ID within a block (x, y, z components).
pub const TID_X: u16 = 10000;
pub const TID_Y: u16 = 10001;
pub const TID_Z: u16 = 10002;
/// %tid combined alias.
pub const TID: u16 = 10000;

/// Number of threads in a block (block dimension).
pub const NTID_X: u16 = 10003;
pub const NTID_Y: u16 = 10004;
pub const NTID_Z: u16 = 10005;
/// %ntid combined alias.
pub const NTID: u16 = 10003;

/// Block ID within a grid (x, y, z components).
pub const CTAID_X: u16 = 10006;
pub const CTAID_Y: u16 = 10007;
pub const CTAID_Z: u16 = 10008;
/// %ctaid combined alias.
pub const CTAID: u16 = 10006;

/// Number of blocks in a grid (grid dimension).
pub const NCTAID_X: u16 = 10009;
pub const NCTAID_Y: u16 = 10010;
pub const NCTAID_Z: u16 = 10011;
/// %nctaid combined alias.
pub const NCTAID: u16 = 10009;

/// Grid-wide identifier.
pub const GRIDID: u16 = 10012;

/// Lane ID within a warp (0–31).
pub const LANEID: u16 = 10013;

/// Warp ID within a block.
pub const WARPID: u16 = 10014;

/// Number of streaming multiprocessors.
pub const NSMID: u16 = 10015;

/// Streaming multiprocessor ID.
pub const SMID: u16 = 10016;

/// Clock register (per-SM cycle counter).
pub const CLOCK: u16 = 10017;

/// Performance monitor registers 0–7.
pub const PM0: u16 = 10018;
pub const PM1: u16 = 10019;
pub const PM2: u16 = 10020;
pub const PM3: u16 = 10021;
pub const PM4: u16 = 10022;
pub const PM5: u16 = 10023;
pub const PM6: u16 = 10024;
pub const PM7: u16 = 10025;

/// Lane ID mask (full 32-bit mask of active lanes).
pub const LANEMASK_EQ: u16 = 10026;
pub const LANEMASK_LE: u16 = 10027;
pub const LANEMASK_LT: u16 = 10028;
pub const LANEMASK_GE: u16 = 10029;
pub const LANEMASK_GT: u16 = 10030;

/// Environment / driver-level registers.
pub const ENVREG0: u16 = 10031;
pub const ENVREG1: u16 = 10032;
pub const ENVREG2: u16 = 10033;
pub const ENVREG3: u16 = 10034;
pub const ENVREG4: u16 = 10035;
pub const ENVREG5: u16 = 10036;
pub const ENVREG6: u16 = 10037;
pub const ENVREG7: u16 = 10038;

/// Maximum special register ID.
pub const NVPTX_SPECIAL_REG_LAST: u16 = 10038;

// ============================================================================
// Predicate Register IDs (base 10100)
// ============================================================================

/// Base offset for predicate registers.
pub const PRED_BASE: u16 = 10100;

/// Maximum predicate register ID (virtual, up to 10299).
pub const PRED_MAX: u16 = 10299;

/// Helper to compute predicate register ID from index.
pub const fn pred_reg(n: u16) -> u16 {
    PRED_BASE + n
}

// ============================================================================
// Virtual Register Range
// ============================================================================

/// Base offset for virtual general-purpose registers.
pub const VREG_BASE: u16 = 10300;

/// Maximum register ID used in this backend.
pub const NVPTX_MAX_REG_ID: u16 = 65535;

/// Number of special registers (non-predicate, non-vreg).
pub const NVPTX_SPECIAL_REG_COUNT: usize = 39;

// ============================================================================
// Register Class Enum
// ============================================================================

/// NVPTX register class.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum NvptxRegClass {
    /// Special register (thread id, block id, etc.).
    SpecialReg,
    /// Predicate register.
    PredReg,
    /// Virtual general-purpose register (integer or float).
    VirtReg,
}

impl std::fmt::Display for NvptxRegClass {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            NvptxRegClass::SpecialReg => write!(f, "SpecialReg"),
            NvptxRegClass::PredReg => write!(f, "PredReg"),
            NvptxRegClass::VirtReg => write!(f, "VirtReg"),
        }
    }
}

// ============================================================================
// NvptxRegisterInfo
// ============================================================================

/// NVPTX register information provider.
///
/// PTX uses unlimited virtual registers. Physical register assignment
/// is performed by the PTX-to-SASS JIT compiler in the NVIDIA driver.
pub struct NvptxRegisterInfo;

impl NvptxRegisterInfo {
    /// Get the PTX assembly name for a register ID.
    pub fn get_asm_name(reg_id: u16) -> String {
        match reg_id {
            TID_X => "%tid.x".to_string(),
            TID_Y => "%tid.y".to_string(),
            TID_Z => "%tid.z".to_string(),
            NTID_X => "%ntid.x".to_string(),
            NTID_Y => "%ntid.y".to_string(),
            NTID_Z => "%ntid.z".to_string(),
            CTAID_X => "%ctaid.x".to_string(),
            CTAID_Y => "%ctaid.y".to_string(),
            CTAID_Z => "%ctaid.z".to_string(),
            NCTAID_X => "%nctaid.x".to_string(),
            NCTAID_Y => "%nctaid.y".to_string(),
            NCTAID_Z => "%nctaid.z".to_string(),
            GRIDID => "%gridid".to_string(),
            LANEID => "%laneid".to_string(),
            WARPID => "%warpid".to_string(),
            NSMID => "%nsmid".to_string(),
            SMID => "%smid".to_string(),
            CLOCK => "%clock".to_string(),
            PM0 => "%pm0".to_string(),
            PM1 => "%pm1".to_string(),
            PM2 => "%pm2".to_string(),
            PM3 => "%pm3".to_string(),
            PM4 => "%pm4".to_string(),
            PM5 => "%pm5".to_string(),
            PM6 => "%pm6".to_string(),
            PM7 => "%pm7".to_string(),
            LANEMASK_EQ => "%lanemask_eq".to_string(),
            LANEMASK_LE => "%lanemask_le".to_string(),
            LANEMASK_LT => "%lanemask_lt".to_string(),
            LANEMASK_GE => "%lanemask_ge".to_string(),
            LANEMASK_GT => "%lanemask_gt".to_string(),
            ENVREG0 => "%envreg0".to_string(),
            ENVREG1 => "%envreg1".to_string(),
            ENVREG2 => "%envreg2".to_string(),
            ENVREG3 => "%envreg3".to_string(),
            ENVREG4 => "%envreg4".to_string(),
            ENVREG5 => "%envreg5".to_string(),
            ENVREG6 => "%envreg6".to_string(),
            ENVREG7 => "%envreg7".to_string(),
            id if id >= PRED_BASE && id <= PRED_MAX => format!("%p{}", id - PRED_BASE),
            id if id >= VREG_BASE => format!("%r{}", id - VREG_BASE),
            _ => format!("%r{}", reg_id),
        }
    }

    /// Get the register class for a register ID.
    pub fn get_reg_class(reg_id: u16) -> NvptxRegClass {
        if reg_id <= NVPTX_SPECIAL_REG_LAST {
            NvptxRegClass::SpecialReg
        } else if reg_id >= PRED_BASE && reg_id <= PRED_MAX {
            NvptxRegClass::PredReg
        } else {
            NvptxRegClass::VirtReg
        }
    }

    /// Return the bit width of a register. Virtual regs are 32 or 64-bit.
    pub fn get_reg_width(_reg_id: u16, is_64bit: bool) -> u8 {
        if is_64bit {
            64
        } else {
            32
        }
    }

    /// Returns true if this register is a special register.
    pub fn is_special_reg(reg_id: u16) -> bool {
        reg_id <= NVPTX_SPECIAL_REG_LAST
    }

    /// Returns true if this register is a predicate register.
    pub fn is_predicate(reg_id: u16) -> bool {
        reg_id >= PRED_BASE && reg_id <= PRED_MAX
    }

    /// Returns true if this register is a virtual register.
    pub fn is_virtual(reg_id: u16) -> bool {
        reg_id >= VREG_BASE
    }

    /// Get the index for a virtual register.
    pub fn get_vreg_index(reg_id: u16) -> u16 {
        if reg_id >= VREG_BASE {
            reg_id - VREG_BASE
        } else {
            0
        }
    }

    /// Get the list of all special register IDs.
    pub fn get_special_regs() -> Vec<u16> {
        vec![
            TID_X,
            TID_Y,
            TID_Z,
            NTID_X,
            NTID_Y,
            NTID_Z,
            CTAID_X,
            CTAID_Y,
            CTAID_Z,
            NCTAID_X,
            NCTAID_Y,
            NCTAID_Z,
            GRIDID,
            LANEID,
            WARPID,
            NSMID,
            SMID,
            CLOCK,
            PM0,
            PM1,
            PM2,
            PM3,
            PM4,
            PM5,
            PM6,
            PM7,
            LANEMASK_EQ,
            LANEMASK_LE,
            LANEMASK_LT,
            LANEMASK_GE,
            LANEMASK_GT,
            ENVREG0,
            ENVREG1,
            ENVREG2,
            ENVREG3,
            ENVREG4,
            ENVREG5,
            ENVREG6,
            ENVREG7,
        ]
    }
}

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

    #[test]
    fn test_special_reg_names() {
        assert_eq!(NvptxRegisterInfo::get_asm_name(TID_X), "%tid.x");
        assert_eq!(NvptxRegisterInfo::get_asm_name(NTID_X), "%ntid.x");
        assert_eq!(NvptxRegisterInfo::get_asm_name(CTAID_X), "%ctaid.x");
        assert_eq!(NvptxRegisterInfo::get_asm_name(LANEID), "%laneid");
        assert_eq!(NvptxRegisterInfo::get_asm_name(CLOCK), "%clock");
    }

    #[test]
    fn test_predicate_reg_names() {
        assert_eq!(NvptxRegisterInfo::get_asm_name(pred_reg(0)), "%p0");
        assert_eq!(NvptxRegisterInfo::get_asm_name(pred_reg(42)), "%p42");
    }

    #[test]
    fn test_vreg_names() {
        assert_eq!(NvptxRegisterInfo::get_asm_name(VREG_BASE), "%r0");
        assert_eq!(NvptxRegisterInfo::get_asm_name(VREG_BASE + 15), "%r15");
    }

    #[test]
    fn test_reg_class() {
        assert_eq!(
            NvptxRegisterInfo::get_reg_class(TID_X),
            NvptxRegClass::SpecialReg
        );
        assert_eq!(
            NvptxRegisterInfo::get_reg_class(pred_reg(0)),
            NvptxRegClass::PredReg
        );
        assert_eq!(
            NvptxRegisterInfo::get_reg_class(VREG_BASE + 5),
            NvptxRegClass::VirtReg
        );
    }

    #[test]
    fn test_is_special_predicate_virtual() {
        assert!(NvptxRegisterInfo::is_special_reg(TID_X));
        assert!(!NvptxRegisterInfo::is_special_reg(VREG_BASE));
        assert!(NvptxRegisterInfo::is_predicate(pred_reg(0)));
        assert!(!NvptxRegisterInfo::is_predicate(VREG_BASE));
        assert!(NvptxRegisterInfo::is_virtual(VREG_BASE + 10));
    }
}