llvm-native-core 0.1.10

//! NVPTX MC Encoder — emits PTX assembly text.
//!
//! Unlike traditional ISAs, PTX is a virtual text-based ISA. The
//! NVIDIA driver JIT-compiles PTX to SASS (the actual machine code).
//! Therefore, the "encoder" emits PTX assembly source lines rather
//! than binary bytes.
//!
//! Clean-room behavioral reconstruction from the NVIDIA PTX ISA
//! Reference Manual. Zero LLVM source code consultation.

use super::nvptx_instr_info::{NvptxInstrInfo, NvptxOpcode};
use super::nvptx_register_info::NvptxRegisterInfo;
use super::nvptx_register_info::VREG_BASE;
use crate::codegen::*;

// ============================================================================
// NvptxMCEncoder
// ============================================================================

/// NVPTX machine code encoder: emits PTX assembly text.
///
/// Accumulates PTX assembly instructions into a string buffer.
pub struct NvptxMCEncoder {
    /// Accumulated PTX assembly output lines.
    pub output: String,
    /// Instruction info for mnemonic lookup.
    pub instr_info: NvptxInstrInfo,
    /// Whether targeting sm_35+ (affects supported instructions).
    pub sm_version: u32,
    /// Whether using 64-bit addressing.
    pub is_64bit: bool,
    /// Current indentation level (tab count).
    indent: usize,
}

impl NvptxMCEncoder {
    /// Create a new PTX encoder.
    pub fn new(sm_version: u32, is_64bit: bool) -> Self {
        NvptxMCEncoder {
            output: String::with_capacity(4096),
            instr_info: NvptxInstrInfo::new(),
            sm_version,
            is_64bit,
            indent: 0,
        }
    }

    // ------------------------------------------------------------------
    // Public API
    // ------------------------------------------------------------------

    /// Encode a single machine instruction as a PTX assembly line.
    pub fn encode_instruction(&mut self, mi: &MachineInstr) {
        let line = self.emit_instruction(mi);
        self.output.push_str(&line);
    }

    /// Encode all instructions from a MachineFunction as PTX assembly.
    pub fn encode_function(&mut self, mf: &MachineFunction) {
        // Emit function declaration
        let entry_attr = self.function_attributes();
        self.output.push_str(&format!("\t.entry {}(\n", mf.name));

        if self.is_64bit {
            self.output.push_str("\t\t.param .u64 __retval0\n");
        } else {
            self.output.push_str("\t\t.param .u32 __retval0\n");
        }
        self.output.push_str("\t)\n");

        // Emit function body
        self.output.push_str("{\n");

        for block in &mf.blocks {
            if !block.name.is_empty() {
                self.output.push_str(&format!(".LBB_{}:\n", block.name));
            }
            for instr in &block.instructions {
                let line = self.emit_instruction(instr);
                self.output.push_str(&line);
            }
        }

        self.output.push_str("}\n");
    }

    /// Emit a complete PTX module with version directive, target,
    /// address size, and all functions.
    pub fn encode_module(&mut self, mf: &MachineFunction) -> String {
        // Version and target
        self.output
            .push_str(&format!("\t.version\t{}\n", self.ptx_version()));
        self.output
            .push_str(&format!("\t.target\tsm_{}\n", self.sm_version));
        if self.is_64bit {
            self.output.push_str("\t.address_size\t64\n");
        } else {
            self.output.push_str("\t.address_size\t32\n");
        }

        self.encode_function(mf);
        self.output.clone()
    }

    // ------------------------------------------------------------------
    // Instruction emission
    // ------------------------------------------------------------------

    /// Emit a single PTX instruction line (with indentation).
    fn emit_instruction(&self, mi: &MachineInstr) -> String {
        let opcode_val = mi.opcode;
        // Try to find the NvptxOpcode by value
        let mnemonic = match self.numeric_to_opcode(opcode_val) {
            Some(op) => self.instr_info.get_mnemonic(op).to_string(),
            None => "// unknown".to_string(),
        };

        if mnemonic == "INVALID" {
            return format!("\t// unknown opcode {}\n", opcode_val);
        }

        // Special handling for terminators and nop-like instructions
        match mnemonic.as_str() {
            "ret" => return "\tret;\n".to_string(),
            "exit" => return "\texit;\n".to_string(),
            "brkpt" => return "\tbrkpt;\n".to_string(),
            "membar" => {
                let level = self.format_operand_text(&mi.operands, 0);
                return format!("\tmembar.{};\n", level);
            }
            "bar" | "bar.sync" | "bar.arrive" => {
                let arg = self.format_operand_text(&mi.operands, 0);
                return format!("\t{}.sync\t{};\n", mnemonic, arg);
            }
            "bra" => {
                let target = self.format_operand_text(&mi.operands, 0);
                return format!("\tbra\t.LLabel_{};\n", target);
            }
            "call" => {
                let func = self.format_operand_text(&mi.operands, 0);
                return format!("\tcall\t{};\n", func);
            }
            _ => {}
        }

        // General instruction: format operands
        let mut ops_str = String::new();
        for (i, op) in mi.operands.iter().enumerate() {
            if i > 0 {
                ops_str.push_str(", ");
            }
            ops_str.push_str(&self.format_operand_text(&mi.operands, i));
        }

        // Determine type suffix based on instruction kind
        let type_suffix = self.get_type_suffix(mi);

        format!("\t{}{}\t{};\n", mnemonic, type_suffix, ops_str)
    }

    /// Get the PTX type suffix for an instruction.
    fn get_type_suffix(&self, mi: &MachineInstr) -> String {
        let op = self.numeric_to_opcode(mi.opcode);
        match op {
            // Integer ops default to .s32 (or .u32 for unsigned)
            Some(NvptxOpcode::Add)
            | Some(NvptxOpcode::Sub)
            | Some(NvptxOpcode::Mul)
            | Some(NvptxOpcode::Mad)
            | Some(NvptxOpcode::Div)
            | Some(NvptxOpcode::Rem) => ".s32".to_string(),
            Some(NvptxOpcode::Abs) | Some(NvptxOpcode::Neg) => ".s32".to_string(),
            Some(NvptxOpcode::Min) | Some(NvptxOpcode::Max) => ".s32".to_string(),
            Some(NvptxOpcode::Popc) | Some(NvptxOpcode::Clz) => ".b32".to_string(),
            Some(NvptxOpcode::Brev) | Some(NvptxOpcode::Bfind) => ".b32".to_string(),
            // Float ops default to .f32
            Some(NvptxOpcode::FAdd)
            | Some(NvptxOpcode::FSub)
            | Some(NvptxOpcode::FMul)
            | Some(NvptxOpcode::FDiv)
            | Some(NvptxOpcode::Fma) => ".f32".to_string(),
            Some(NvptxOpcode::FAbs) | Some(NvptxOpcode::FNeg) => ".f32".to_string(),
            Some(NvptxOpcode::FMin) | Some(NvptxOpcode::FMax) => ".f32".to_string(),
            Some(NvptxOpcode::Rcp) | Some(NvptxOpcode::Sqrt) | Some(NvptxOpcode::Rsqrt) => {
                ".f32".to_string()
            }
            Some(NvptxOpcode::Sin)
            | Some(NvptxOpcode::Cos)
            | Some(NvptxOpcode::Lg2)
            | Some(NvptxOpcode::Ex2)
            | Some(NvptxOpcode::Tanh) => ".f32".to_string(),
            // Logic ops default to .b32
            Some(NvptxOpcode::And)
            | Some(NvptxOpcode::Or)
            | Some(NvptxOpcode::Xor)
            | Some(NvptxOpcode::Not) => ".b32".to_string(),
            Some(NvptxOpcode::Shl) | Some(NvptxOpcode::Shr) => ".b32".to_string(),
            Some(NvptxOpcode::Cnot) => ".b32".to_string(),
            // Data: loads/stores
            Some(NvptxOpcode::Ld)
            | Some(NvptxOpcode::LdGlobal)
            | Some(NvptxOpcode::LdShared)
            | Some(NvptxOpcode::LdLocal)
            | Some(NvptxOpcode::LdParam)
            | Some(NvptxOpcode::LdConst)
            | Some(NvptxOpcode::LdGlobalNc)
            | Some(NvptxOpcode::LdGlobalCg)
            | Some(NvptxOpcode::LdGlobalCa) => ".b32".to_string(),
            Some(NvptxOpcode::St)
            | Some(NvptxOpcode::StGlobal)
            | Some(NvptxOpcode::StShared)
            | Some(NvptxOpcode::StLocal)
            | Some(NvptxOpcode::StParam) => ".b32".to_string(),
            // Movement
            Some(NvptxOpcode::Mov) => ".b32".to_string(),
            Some(NvptxOpcode::Cvta) => ".to".to_string(),
            // Conversion
            Some(NvptxOpcode::Cvt)
            | Some(NvptxOpcode::CvtRp)
            | Some(NvptxOpcode::CvtRn)
            | Some(NvptxOpcode::CvtRz)
            | Some(NvptxOpcode::CvtRm) => ".f32.s32".to_string(),
            // Compare
            op if op
                .map(|o| {
                    matches!(
                        o,
                        NvptxOpcode::SetpEq
                            | NvptxOpcode::SetpNe
                            | NvptxOpcode::SetpLt
                            | NvptxOpcode::SetpLe
                            | NvptxOpcode::SetpGt
                            | NvptxOpcode::SetpGe
                            | NvptxOpcode::SetpLo
                            | NvptxOpcode::SetpLs
                            | NvptxOpcode::SetpHi
                            | NvptxOpcode::SetpHs
                            | NvptxOpcode::SetpEqu
                            | NvptxOpcode::SetpNeu
                            | NvptxOpcode::SetpLtu
                            | NvptxOpcode::SetpLeu
                            | NvptxOpcode::SetpGtu
                            | NvptxOpcode::SetpGeu
                            | NvptxOpcode::SetpNum
                            | NvptxOpcode::SetpNan
                    )
                })
                .unwrap_or(false) =>
            {
                ".".to_string()
            }
            // Atomics
            Some(NvptxOpcode::AtomAdd)
            | Some(NvptxOpcode::AtomSub)
            | Some(NvptxOpcode::AtomExch)
            | Some(NvptxOpcode::AtomCas)
            | Some(NvptxOpcode::AtomMin)
            | Some(NvptxOpcode::AtomMax)
            | Some(NvptxOpcode::AtomAnd)
            | Some(NvptxOpcode::AtomOr)
            | Some(NvptxOpcode::AtomXor) => ".global.b32".to_string(),
            // Reductions
            Some(NvptxOpcode::RedAdd)
            | Some(NvptxOpcode::RedSub)
            | Some(NvptxOpcode::RedMin)
            | Some(NvptxOpcode::RedMax)
            | Some(NvptxOpcode::RedAnd)
            | Some(NvptxOpcode::RedOr)
            | Some(NvptxOpcode::RedXor) => ".shared.b32".to_string(),
            _ => "".to_string(),
        }
    }

    /// Format a single operand as PTX text.
    fn format_operand_text(&self, ops: &[MachineOperand], index: usize) -> String {
        if index >= ops.len() {
            return "0".to_string();
        }
        match &ops[index] {
            MachineOperand::Reg(vr) => {
                NvptxRegisterInfo::get_asm_name((*vr + VREG_BASE as u32) as u16)
            }
            MachineOperand::PhysReg(reg) => NvptxRegisterInfo::get_asm_name(*reg as u16),
            MachineOperand::Imm(imm) => format!("{}", imm),
            MachineOperand::Label(label) => format!(".LLabel_{}", label),
            MachineOperand::Global(name) => name.clone(),
        }
    }

    // ------------------------------------------------------------------
    // Helpers
    // ------------------------------------------------------------------

    fn numeric_to_opcode(&self, val: u32) -> Option<NvptxOpcode> {
        match val {
            20000 => Some(NvptxOpcode::Add),
            20001 => Some(NvptxOpcode::Sub),
            20002 => Some(NvptxOpcode::Mul),
            20003 => Some(NvptxOpcode::Mad),
            20004 => Some(NvptxOpcode::Div),
            20005 => Some(NvptxOpcode::Rem),
            20006 => Some(NvptxOpcode::Abs),
            20007 => Some(NvptxOpcode::Neg),
            20008 => Some(NvptxOpcode::Min),
            20009 => Some(NvptxOpcode::Max),
            20010 => Some(NvptxOpcode::Popc),
            20011 => Some(NvptxOpcode::Clz),
            20012 => Some(NvptxOpcode::Bfind),
            20013 => Some(NvptxOpcode::Brev),
            20014 => Some(NvptxOpcode::Bfe),
            20015 => Some(NvptxOpcode::Bfi),
            20016 => Some(NvptxOpcode::Sad),
            20017 => Some(NvptxOpcode::Selp),
            20020 => Some(NvptxOpcode::FAdd),
            20021 => Some(NvptxOpcode::FSub),
            20022 => Some(NvptxOpcode::FMul),
            20023 => Some(NvptxOpcode::FMad),
            20024 => Some(NvptxOpcode::Fma),
            20025 => Some(NvptxOpcode::FDiv),
            20026 => Some(NvptxOpcode::FAbs),
            20027 => Some(NvptxOpcode::FNeg),
            20028 => Some(NvptxOpcode::FMin),
            20029 => Some(NvptxOpcode::FMax),
            20030 => Some(NvptxOpcode::Rcp),
            20031 => Some(NvptxOpcode::Sqrt),
            20032 => Some(NvptxOpcode::Rsqrt),
            20033 => Some(NvptxOpcode::Sin),
            20034 => Some(NvptxOpcode::Cos),
            20035 => Some(NvptxOpcode::Lg2),
            20036 => Some(NvptxOpcode::Ex2),
            20037 => Some(NvptxOpcode::Tanh),
            20040 => Some(NvptxOpcode::SetpEq),
            20041 => Some(NvptxOpcode::SetpNe),
            20042 => Some(NvptxOpcode::SetpLt),
            20043 => Some(NvptxOpcode::SetpLe),
            20044 => Some(NvptxOpcode::SetpGt),
            20060 => Some(NvptxOpcode::And),
            20061 => Some(NvptxOpcode::Or),
            20062 => Some(NvptxOpcode::Xor),
            20063 => Some(NvptxOpcode::Not),
            20064 => Some(NvptxOpcode::Shl),
            20065 => Some(NvptxOpcode::Shr),
            20066 => Some(NvptxOpcode::Cnot),
            20070 => Some(NvptxOpcode::Ld),
            20071 => Some(NvptxOpcode::St),
            20072 => Some(NvptxOpcode::LdGlobal),
            20073 => Some(NvptxOpcode::StGlobal),
            20074 => Some(NvptxOpcode::LdShared),
            20075 => Some(NvptxOpcode::StShared),
            20076 => Some(NvptxOpcode::LdLocal),
            20077 => Some(NvptxOpcode::StLocal),
            20078 => Some(NvptxOpcode::LdParam),
            20079 => Some(NvptxOpcode::StParam),
            20080 => Some(NvptxOpcode::LdConst),
            20084 => Some(NvptxOpcode::Cvta),
            20085 => Some(NvptxOpcode::Mov),
            20090 => Some(NvptxOpcode::Cvt),
            20100 => Some(NvptxOpcode::Bra),
            20101 => Some(NvptxOpcode::Call),
            20102 => Some(NvptxOpcode::Ret),
            20103 => Some(NvptxOpcode::Exit),
            20104 => Some(NvptxOpcode::Brkpt),
            20105 => Some(NvptxOpcode::Bar),
            20106 => Some(NvptxOpcode::BarSync),
            20107 => Some(NvptxOpcode::BarArrive),
            20108 => Some(NvptxOpcode::Membar),
            20110 => Some(NvptxOpcode::AtomAdd),
            20111 => Some(NvptxOpcode::AtomSub),
            20112 => Some(NvptxOpcode::AtomExch),
            20113 => Some(NvptxOpcode::AtomCas),
            20114 => Some(NvptxOpcode::AtomMin),
            20115 => Some(NvptxOpcode::AtomMax),
            20116 => Some(NvptxOpcode::AtomAnd),
            20117 => Some(NvptxOpcode::AtomOr),
            20118 => Some(NvptxOpcode::AtomXor),
            20120 => Some(NvptxOpcode::RedAdd),
            20121 => Some(NvptxOpcode::RedSub),
            20122 => Some(NvptxOpcode::RedMin),
            20123 => Some(NvptxOpcode::RedMax),
            20124 => Some(NvptxOpcode::RedAnd),
            20125 => Some(NvptxOpcode::RedOr),
            20126 => Some(NvptxOpcode::RedXor),
            20130 => Some(NvptxOpcode::VoteAll),
            20131 => Some(NvptxOpcode::VoteAny),
            20132 => Some(NvptxOpcode::VoteUni),
            20133 => Some(NvptxOpcode::VoteBallot),
            20140 => Some(NvptxOpcode::Tex),
            20141 => Some(NvptxOpcode::Tld4),
            20142 => Some(NvptxOpcode::Txq),
            20143 => Some(NvptxOpcode::Suld),
            20144 => Some(NvptxOpcode::Sust),
            20145 => Some(NvptxOpcode::Sured),
            20146 => Some(NvptxOpcode::Suq),
            20150 => Some(NvptxOpcode::Isspacep),
            _ => None,
        }
    }

    fn ptx_version(&self) -> &str {
        if self.sm_version >= 70 {
            "7.0"
        } else if self.sm_version >= 60 {
            "6.0"
        } else if self.sm_version >= 50 {
            "5.0"
        } else {
            "4.3"
        }
    }

    fn function_attributes(&self) -> &str {
        ".visible"
    }
}

// ============================================================================
// Compute Capability Support
// ============================================================================

/// PTX compute capability mapping: sm version to PTX ISA version,
/// supported features, and hardware generation.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ComputeCapability {
    /// Maxwell generation (sm_50, sm_52, sm_53)
    Sm50,
    /// Pascal generation (sm_60, sm_61, sm_62)
    Sm60,
    /// Volta generation (sm_70, sm_72)
    Sm70,
    /// Turing generation (sm_75)
    Sm75,
    /// Ampere generation (sm_80, sm_86)
    Sm80,
    /// Ada Lovelace generation (sm_89)
    Sm89,
    /// Hopper generation (sm_90)
    Sm90,
}

impl ComputeCapability {
    /// Determine compute capability from SM version number.
    pub fn from_sm(sm: u32) -> Self {
        match sm {
            50..=53 => ComputeCapability::Sm50,
            60..=62 => ComputeCapability::Sm60,
            70..=72 => ComputeCapability::Sm70,
            75 => ComputeCapability::Sm75,
            80..=86 => ComputeCapability::Sm80,
            89 => ComputeCapability::Sm89,
            90..=99 => ComputeCapability::Sm90,
            _ => ComputeCapability::Sm50,
        }
    }

    /// Return the PTX ISA version string for this compute capability.
    pub fn ptx_version(&self) -> &str {
        match self {
            ComputeCapability::Sm50 => "5.0",
            ComputeCapability::Sm60 => "6.0",
            ComputeCapability::Sm70 => "7.0",
            ComputeCapability::Sm75 => "7.5",
            ComputeCapability::Sm80 => "7.8",
            ComputeCapability::Sm89 => "8.0",
            ComputeCapability::Sm90 => "8.2",
        }
    }

    /// Check if a feature is supported at this compute capability.
    pub fn supports(&self, feature: ComputeFeature) -> bool {
        match feature {
            ComputeFeature::UnifiedMemory => matches!(self, ComputeCapability::Sm60 | ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::TensorCores => matches!(self, ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::CooperativeGroups => matches!(self, ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::IndependentThreadScheduling => matches!(self, ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::WarpMatrix => matches!(self, ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::AsyncCopy => matches!(self, ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::DPX => matches!(self, ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::ThreadBlockCluster => matches!(self, ComputeCapability::Sm90),
            ComputeFeature::FP64 => matches!(self, ComputeCapability::Sm60 | ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::FP16 => matches!(self, ComputeCapability::Sm50 | ComputeCapability::Sm60 | ComputeCapability::Sm70 | ComputeCapability::Sm75 | ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
            ComputeFeature::BF16 => matches!(self, ComputeCapability::Sm80 | ComputeCapability::Sm89 | ComputeCapability::Sm90),
        }
    }

    /// Number of registers per SM.
    pub fn registers_per_sm(&self) -> u32 {
        match self {
            ComputeCapability::Sm50 => 65536,
            ComputeCapability::Sm60 => 65536,
            ComputeCapability::Sm70 => 65536,
            ComputeCapability::Sm75 => 65536,
            ComputeCapability::Sm80 => 65536,
            ComputeCapability::Sm89 => 65536,
            ComputeCapability::Sm90 => 65536,
        }
    }

    /// Max shared memory per block (bytes).
    pub fn max_shared_memory(&self) -> u32 {
        match self {
            ComputeCapability::Sm50 => 49152,
            ComputeCapability::Sm60 => 49152,
            ComputeCapability::Sm70 => 98304,
            ComputeCapability::Sm75 => 65536,
            ComputeCapability::Sm80 => 163840,
            ComputeCapability::Sm89 => 49152,
            ComputeCapability::Sm90 => 227328,
        }
    }
}

/// Compute features gated by SM version.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ComputeFeature {
    UnifiedMemory,
    TensorCores,
    CooperativeGroups,
    IndependentThreadScheduling,
    WarpMatrix,
    AsyncCopy,
    DPX,
    ThreadBlockCluster,
    FP64,
    FP16,
    BF16,
}

// ============================================================================
// Predicate Execution Support
// ============================================================================

/// PTX predicate guard: @%p (execute when true) or @!%p (execute when false).
#[derive(Debug, Clone, PartialEq)]
pub struct PredicateGuard {
    /// Predicate register number (0-7 typically).
    pub reg: u8,
    /// If true, negate the predicate (@!%p form).
    pub negated: bool,
}

impl PredicateGuard {
    pub fn new(reg: u8, negated: bool) -> Self {
        PredicateGuard { reg, negated }
    }

    /// Format predicate as PTX text: "@%p0" or "@!%p0".
    pub fn to_ptx(&self) -> String {
        if self.negated {
            format!("@!%p{}", self.reg)
        } else {
            format!("@%p{}", self.reg)
        }
    }
}

/// Builder for constructing predicated PTX instructions.
pub struct PredicatedInstruction {
    /// Optional predicate guard.
    pub guard: Option<PredicateGuard>,
    /// The base PTX opcode.
    pub opcode: NvptxOpcode,
    /// Type suffix (e.g., ".s32").
    pub type_suffix: PtxType,
    /// Vector specifier if applicable.
    pub vector_spec: Option<PtxVector>,
    /// Cache operator if applicable.
    pub cache_op: Option<PtxCacheOp>,
    /// Operands as PTX text.
    pub operands: Vec<String>,
}

impl PredicatedInstruction {
    pub fn new(opcode: NvptxOpcode) -> Self {
        PredicatedInstruction {
            guard: None,
            opcode,
            type_suffix: PtxType::B32,
            vector_spec: None,
            cache_op: None,
            operands: Vec::new(),
        }
    }

    pub fn with_predicate(mut self, reg: u8, negated: bool) -> Self {
        self.guard = Some(PredicateGuard::new(reg, negated));
        self
    }

    pub fn with_type(mut self, ty: PtxType) -> Self {
        self.type_suffix = ty;
        self
    }

    pub fn with_vector(mut self, v: PtxVector) -> Self {
        self.vector_spec = Some(v);
        self
    }

    pub fn with_cache(mut self, c: PtxCacheOp) -> Self {
        self.cache_op = Some(c);
        self
    }

    pub fn add_operand(mut self, op: &str) -> Self {
        self.operands.push(op.to_string());
        self
    }

    /// Emit the full PTX instruction text line.
    pub fn emit(&self) -> String {
        let mut line = String::new();
        line.push('\t');
        if let Some(ref g) = self.guard {
            line.push_str(&g.to_ptx());
            line.push(' ');
        }
        // Opcode part
        let mnemonic = NvptxInstrInfo::get_mnemonic_static(self.opcode);
        line.push_str(mnemonic);
        // Type suffix
        line.push_str(&self.type_suffix.to_ptx());
        // Vector spec
        if let Some(ref v) = self.vector_spec {
            line.push_str(&v.to_ptx());
        }
        // Cache op
        if let Some(ref c) = self.cache_op {
            line.push_str(&c.to_ptx());
        }
        // Operands
        if !self.operands.is_empty() {
            line.push('\t');
            line.push_str(&self.operands.join(", "));
        }
        line.push_str(";\n");
        line
    }
}

// ============================================================================
// PTX Type System
// ============================================================================

/// PTX scalar type suffixes.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxType {
    /// Unsigned 8-bit integer.
    U8,
    /// Signed 8-bit integer.
    S8,
    /// Bit-size 8-bit.
    B8,
    /// Unsigned 16-bit integer.
    U16,
    /// Signed 16-bit integer.
    S16,
    /// Bit-size 16-bit (also used for bf16).
    B16,
    /// Unsigned 32-bit integer.
    U32,
    /// Signed 32-bit integer.
    S32,
    /// Bit-size 32-bit.
    B32,
    /// Unsigned 64-bit integer.
    U64,
    /// Signed 64-bit integer.
    S64,
    /// Bit-size 64-bit.
    B64,
    /// 16-bit floating point (IEEE 754 binary16).
    F16,
    /// 16-bit floating point x2 (packed).
    F16x2,
    /// 32-bit floating point.
    F32,
    /// 64-bit floating point.
    F64,
    /// Brain floating point 16.
    Bf16,
    /// Predicate (1-bit).
    Pred,
}

impl PtxType {
    /// Convert to PTX type suffix string, e.g., ".u32".
    pub fn to_ptx(&self) -> String {
        match self {
            PtxType::U8 => ".u8".to_string(),
            PtxType::S8 => ".s8".to_string(),
            PtxType::B8 => ".b8".to_string(),
            PtxType::U16 => ".u16".to_string(),
            PtxType::S16 => ".s16".to_string(),
            PtxType::B16 => ".b16".to_string(),
            PtxType::U32 => ".u32".to_string(),
            PtxType::S32 => ".s32".to_string(),
            PtxType::B32 => ".b32".to_string(),
            PtxType::U64 => ".u64".to_string(),
            PtxType::S64 => ".s64".to_string(),
            PtxType::B64 => ".b64".to_string(),
            PtxType::F16 => ".f16".to_string(),
            PtxType::F16x2 => ".f16x2".to_string(),
            PtxType::F32 => ".f32".to_string(),
            PtxType::F64 => ".f64".to_string(),
            PtxType::Bf16 => ".bf16".to_string(),
            PtxType::Pred => ".pred".to_string(),
        }
    }

    /// Size in bytes.
    pub fn size_bytes(&self) -> u32 {
        match self {
            PtxType::U8 | PtxType::S8 | PtxType::B8 => 1,
            PtxType::U16 | PtxType::S16 | PtxType::B16 | PtxType::F16 | PtxType::Bf16 => 2,
            PtxType::F16x2 => 4,
            PtxType::U32 | PtxType::S32 | PtxType::B32 | PtxType::F32 => 4,
            PtxType::U64 | PtxType::S64 | PtxType::B64 | PtxType::F64 => 8,
            PtxType::Pred => 0,
        }
    }

    /// Check if integer type.
    pub fn is_integer(&self) -> bool {
        matches!(self, PtxType::U8 | PtxType::S8 | PtxType::U16 | PtxType::S16 | PtxType::U32 | PtxType::S32 | PtxType::U64 | PtxType::S64)
    }

    /// Check if float type.
    pub fn is_float(&self) -> bool {
        matches!(self, PtxType::F16 | PtxType::F16x2 | PtxType::F32 | PtxType::F64 | PtxType::Bf16)
    }

    /// Get unsigned variant of an integer type.
    pub fn to_unsigned(&self) -> PtxType {
        match self {
            PtxType::S8 | PtxType::B8 | PtxType::U8 => PtxType::U8,
            PtxType::S16 | PtxType::B16 | PtxType::U16 => PtxType::U16,
            PtxType::S32 | PtxType::B32 | PtxType::U32 => PtxType::U32,
            PtxType::S64 | PtxType::B64 | PtxType::U64 => PtxType::U64,
            other => *other,
        }
    }

    /// Get signed variant of an integer type.
    pub fn to_signed(&self) -> PtxType {
        match self {
            PtxType::U8 | PtxType::B8 | PtxType::S8 => PtxType::S8,
            PtxType::U16 | PtxType::B16 | PtxType::S16 => PtxType::S16,
            PtxType::U32 | PtxType::B32 | PtxType::S32 => PtxType::S32,
            PtxType::U64 | PtxType::B64 | PtxType::S64 => PtxType::S64,
            other => *other,
        }
    }
}

// ============================================================================
// Vector Types
// ============================================================================

/// PTX vector type specifiers.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxVector {
    V2,
    V4,
}

impl PtxVector {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxVector::V2 => ".v2".to_string(),
            PtxVector::V4 => ".v4".to_string(),
        }
    }

    /// Number of elements in the vector.
    pub fn element_count(&self) -> u32 {
        match self {
            PtxVector::V2 => 2,
            PtxVector::V4 => 4,
        }
    }
}

// ============================================================================
// Cache Operators
// ============================================================================

/// PTX cache operator modifiers for load/store instructions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxCacheOp {
    /// Cache at all levels (default).
    Ca,
    /// Cache at global level (L2 cache).
    Cg,
    /// Cache volatile (bypass cache).
    Cv,
    /// Write-back (for stores).
    Wb,
    /// Write-through (for stores).
    Wt,
    /// Evict-first (L1 evict first).
    Ef,
    /// Evict-last (L1 evict last).
    El,
    /// No-cache streaming.
    Nc,
}

impl PtxCacheOp {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxCacheOp::Ca => ".ca".to_string(),
            PtxCacheOp::Cg => ".cg".to_string(),
            PtxCacheOp::Cv => ".cv".to_string(),
            PtxCacheOp::Wb => ".wb".to_string(),
            PtxCacheOp::Wt => ".wt".to_string(),
            PtxCacheOp::Ef => ".ef".to_string(),
            PtxCacheOp::El => ".el".to_string(),
            PtxCacheOp::Nc => ".nc".to_string(),
        }
    }

    /// Whether this is a read-only cache modifier.
    pub fn is_load_cache(&self) -> bool {
        matches!(self, PtxCacheOp::Ca | PtxCacheOp::Cg | PtxCacheOp::Cv | PtxCacheOp::Nc)
    }

    /// Whether this is a write cache modifier.
    pub fn is_store_cache(&self) -> bool {
        matches!(self, PtxCacheOp::Wb | PtxCacheOp::Wt)
    }
}

// ============================================================================
// Synchronization Primitives
// ============================================================================

/// PTX synchronization scope.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxScope {
    /// Thread-block / CTA scope.
    Cta,
    /// GPU device scope.
    Gpu,
    /// System scope (GPU + CPU + peer devices).
    Sys,
}

impl PtxScope {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxScope::Cta => ".cta".to_string(),
            PtxScope::Gpu => ".gpu".to_string(),
            PtxScope::Sys => ".sys".to_string(),
        }
    }
}

/// PTX memory fence level.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxMemLevel {
    /// Thread-block / CTA level.
    Cta,
    /// GPU device level.
    Gl,
    /// System level.
    Sys,
}

impl PtxMemLevel {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxMemLevel::Cta => "cta".to_string(),
            PtxMemLevel::Gl => "gl".to_string(),
            PtxMemLevel::Sys => "sys".to_string(),
        }
    }
}

/// PTX barrier type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxBarType {
    /// Synchronize threads: bar.sync.
    Sync,
    /// Arrive at barrier without waiting: bar.arrive.
    Arrive,
    /// Reduce across threads at barrier.
    Red,
    /// Scan across threads at barrier.
    Scan,
}

impl PtxBarType {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxBarType::Sync => "bar.sync".to_string(),
            PtxBarType::Arrive => "bar.arrive".to_string(),
            PtxBarType::Red => "bar.red".to_string(),
            PtxBarType::Scan => "bar.scan".to_string(),
        }
    }
}

/// Synchronization emitter: generates bar, membar, and fence instructions.
pub struct SyncEmitter;

impl SyncEmitter {
    /// Emit a barrier synchronization instruction.
    pub fn bar(bar_type: PtxBarType, num_threads: u32) -> String {
        format!("\t{}\t{};\n", bar_type.to_ptx(), num_threads)
    }

    /// Emit a named barrier.
    pub fn bar_named(bar_type: PtxBarType, num_threads: u32, name: &str) -> String {
        format!("\t{}\t{}, {};\n", bar_type.to_ptx(), num_threads, name)
    }

    /// Emit a memory barrier (membar).
    pub fn membar(level: PtxMemLevel) -> String {
        format!("\tmembar.{};\n", level.to_ptx())
    }

    /// Emit a fence instruction with scope.
    pub fn fence(scope: PtxScope, sem: &str) -> String {
        format!("\tfence{}.{};\n", scope.to_ptx(), sem)
    }

    /// Emit a thread fence with acquire semantics.
    pub fn fence_acquire(scope: PtxScope) -> String {
        Self::fence(scope, "acquire")
    }

    /// Emit a thread fence with release semantics.
    pub fn fence_release(scope: PtxScope) -> String {
        Self::fence(scope, "release")
    }

    /// Emit a thread fence with acquire-release semantics.
    pub fn fence_acq_rel(scope: PtxScope) -> String {
        Self::fence(scope, "acq_rel")
    }

    /// Emit a warp-level synchronization (__syncwarp).
    pub fn sync_warp(mask: u32) -> String {
        format!("\tbar.warp.sync\t{};\n", mask)
    }

    /// Arrive at named barrier and wait.
    pub fn arrive_and_wait(name: &str) -> String {
        format!("\tbar.sync\t0, {};\n", name)
    }
}

// ============================================================================
// Special Register Access
// ============================================================================

/// PTX special register enumeration and formatting.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxSpecialReg {
    /// Thread ID: %tid.x, %tid.y, %tid.z
    Tid { dim: char },
    /// Block dimension: %ntid.x, %ntid.y, %ntid.z
    Ntid { dim: char },
    /// Block ID: %ctaid.x, %ctaid.y, %ctaid.z
    Ctaid { dim: char },
    /// Grid dimension: %nctaid.x, %nctaid.y, %nctaid.z
    Nctaid { dim: char },
    /// Lane ID within warp (0-31): %laneid
    Laneid,
    /// Warp ID: %warpid
    Warpid,
    /// Clock counter: %clock, %clock64
    Clock,
    /// 64-bit clock: %clock64
    Clock64,
    /// Performance monitor: %pm0..%pm7
    Pm(u8),
    /// Streaming multiprocessor ID: %smid
    Smid,
    /// Number of SMs: %nsmid
    Nsmid,
    /// Grid-wide ID: %gridid
    Gridid,
    /// Lane mask equal: %lanemask_eq
    LanemaskEq,
    /// Lane mask less or equal: %lanemask_le
    LanemaskLe,
    /// Lane mask less than: %lanemask_lt
    LanemaskLt,
    /// Lane mask greater or equal: %lanemask_ge
    LanemaskGe,
    /// Lane mask greater than: %lanemask_gt
    LanemaskGt,
    /// Total number of SMs launched: %total_smem_size
    TotalSmemSize,
    /// Aggregated block size: %aggr_smem_size
    AggrSmemSize,
    /// Dynamic shared memory size: %dynamic_smem_size
    DynamicSmemSize,
}

impl PtxSpecialReg {
    /// Format the special register as PTX text, e.g., "%tid.x".
    pub fn to_ptx(&self) -> String {
        match self {
            PtxSpecialReg::Tid { dim } => format!("%tid.{}", dim),
            PtxSpecialReg::Ntid { dim } => format!("%ntid.{}", dim),
            PtxSpecialReg::Ctaid { dim } => format!("%ctaid.{}", dim),
            PtxSpecialReg::Nctaid { dim } => format!("%nctaid.{}", dim),
            PtxSpecialReg::Laneid => "%laneid".to_string(),
            PtxSpecialReg::Warpid => "%warpid".to_string(),
            PtxSpecialReg::Clock => "%clock".to_string(),
            PtxSpecialReg::Clock64 => "%clock64".to_string(),
            PtxSpecialReg::Pm(n) => format!("%pm{}", n),
            PtxSpecialReg::Smid => "%smid".to_string(),
            PtxSpecialReg::Nsmid => "%nsmid".to_string(),
            PtxSpecialReg::Gridid => "%gridid".to_string(),
            PtxSpecialReg::LanemaskEq => "%lanemask_eq".to_string(),
            PtxSpecialReg::LanemaskLe => "%lanemask_le".to_string(),
            PtxSpecialReg::LanemaskLt => "%lanemask_lt".to_string(),
            PtxSpecialReg::LanemaskGe => "%lanemask_ge".to_string(),
            PtxSpecialReg::LanemaskGt => "%lanemask_gt".to_string(),
            PtxSpecialReg::TotalSmemSize => "%total_smem_size".to_string(),
            PtxSpecialReg::AggrSmemSize => "%aggr_smem_size".to_string(),
            PtxSpecialReg::DynamicSmemSize => "%dynamic_smem_size".to_string(),
        }
    }

    /// Returns the hardware register number used by the driver.
    pub fn hw_reg_num(&self) -> u16 {
        match self {
            PtxSpecialReg::Tid { dim } => match dim {
                'x' => 10000,
                'y' => 10001,
                'z' => 10002,
                _ => 10000,
            },
            PtxSpecialReg::Ntid { dim } => match dim {
                'x' => 10003,
                'y' => 10004,
                'z' => 10005,
                _ => 10003,
            },
            PtxSpecialReg::Ctaid { dim } => match dim {
                'x' => 10006,
                'y' => 10007,
                'z' => 10008,
                _ => 10006,
            },
            PtxSpecialReg::Nctaid { dim } => match dim {
                'x' => 10009,
                'y' => 10010,
                'z' => 10011,
                _ => 10009,
            },
            PtxSpecialReg::Laneid => 10013,
            PtxSpecialReg::Warpid => 10014,
            PtxSpecialReg::Clock => 10017,
            PtxSpecialReg::Clock64 => 10017,
            PtxSpecialReg::Pm(n) => 10018 + *n as u16,
            PtxSpecialReg::Smid => 10016,
            PtxSpecialReg::Nsmid => 10015,
            PtxSpecialReg::Gridid => 10012,
            PtxSpecialReg::LanemaskEq => 10026,
            PtxSpecialReg::LanemaskLe => 10027,
            PtxSpecialReg::LanemaskLt => 10028,
            PtxSpecialReg::LanemaskGe => 10029,
            PtxSpecialReg::LanemaskGt => 10030,
            PtxSpecialReg::TotalSmemSize => 10040,
            PtxSpecialReg::AggrSmemSize => 10041,
            PtxSpecialReg::DynamicSmemSize => 10042,
        }
    }
}

// ============================================================================
// Atomic Operations with Scope
// ============================================================================

/// PTX atomic operation type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PtxAtomicOp {
    Add,
    Sub,
    Exch,
    Cas,
    Min,
    Max,
    Inc,
    Dec,
    And,
    Or,
    Xor,
}

impl PtxAtomicOp {
    pub fn to_ptx(&self) -> String {
        match self {
            PtxAtomicOp::Add => "atom.add".to_string(),
            PtxAtomicOp::Sub => "atom.sub".to_string(),
            PtxAtomicOp::Exch => "atom.exch".to_string(),
            PtxAtomicOp::Cas => "atom.cas".to_string(),
            PtxAtomicOp::Min => "atom.min".to_string(),
            PtxAtomicOp::Max => "atom.max".to_string(),
            PtxAtomicOp::Inc => "atom.inc".to_string(),
            PtxAtomicOp::Dec => "atom.dec".to_string(),
            PtxAtomicOp::And => "atom.and".to_string(),
            PtxAtomicOp::Or => "atom.or".to_string(),
            PtxAtomicOp::Xor => "atom.xor".to_string(),
        }
    }
}

/// Full atomic instruction emitter with scope and address space qualifiers.
pub struct AtomicEmitter;

impl AtomicEmitter {
    /// Emit an atomic operation with full qualification.
    /// Format: atom.op.scope.sem.space.type addr, val [, cmp]
    pub fn emit(
        op: PtxAtomicOp,
        scope: PtxScope,
        sem: &str,
        space: &str,
        ty: PtxType,
        addr: &str,
        val: &str,
        cmp: Option<&str>,
    ) -> String {
        let mut line = format!(
            "\t{}{}{}.{}{}\t{}, {}",
            op.to_ptx(),
            scope.to_ptx(),
            if sem.is_empty() {
                String::new()
            } else {
                format!(".{}", sem)
            },
            space,
            ty.to_ptx(),
            addr,
            val
        );
        if let Some(c) = cmp {
            line.push_str(&format!(", {}", c));
        }
        line.push_str(";\n");
        line
    }

    /// Emit an atomic add at global scope.
    pub fn atomic_add_global(ty: PtxType, addr: &str, val: &str) -> String {
        Self::emit(PtxAtomicOp::Add, PtxScope::Gpu, "", "global", ty, addr, val, None)
    }

    /// Emit an atomic compare-and-swap at GPU scope.
    pub fn atomic_cas_gpu(ty: PtxType, addr: &str, cmp: &str, val: &str) -> String {
        Self::emit(
            PtxAtomicOp::Cas,
            PtxScope::Gpu,
            "",
            "global",
            ty,
            addr,
            val,
            Some(cmp),
        )
    }

    /// Emit an atomic exchange at system scope.
    pub fn atomic_exch_sys(ty: PtxType, addr: &str, val: &str) -> String {
        Self::emit(PtxAtomicOp::Exch, PtxScope::Sys, "", "global", ty, addr, val, None)
    }

    /// Emit a shared memory atomic (CTA scope).
    pub fn atomic_shared(op: PtxAtomicOp, ty: PtxType, addr: &str, val: &str) -> String {
        Self::emit(op, PtxScope::Cta, "", "shared", ty, addr, val, None)
    }

    /// Emit an atomic with acquire semantics.
    pub fn atomic_acquire(
        op: PtxAtomicOp,
        scope: PtxScope,
        space: &str,
        ty: PtxType,
        addr: &str,
        val: &str,
    ) -> String {
        Self::emit(op, scope, "acquire", space, ty, addr, val, None)
    }

    /// Emit an atomic with release semantics.
    pub fn atomic_release(
        op: PtxAtomicOp,
        scope: PtxScope,
        space: &str,
        ty: PtxType,
        addr: &str,
        val: &str,
    ) -> String {
        Self::emit(op, scope, "release", space, ty, addr, val, None)
    }
}

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

    #[test]
    fn test_encoder_new() {
        let enc = NvptxMCEncoder::new(75, true);
        assert_eq!(enc.sm_version, 75);
        assert!(enc.is_64bit);
    }

    #[test]
    fn test_encode_basic_instruction() {
        let mut enc = NvptxMCEncoder::new(75, true);
        let mut mi = MachineInstr::new(20000); // ADD
        mi.push_imm(1);
        mi.push_imm(2);
        enc.encode_instruction(&mi);
        assert!(enc.output.contains("add.s32"));
    }

    #[test]
    fn test_encode_ret() {
        let mut enc = NvptxMCEncoder::new(75, true);
        let mi = MachineInstr::new(20102); // RET
        enc.encode_instruction(&mi);
        assert!(enc.output.contains("ret;"));
    }

    #[test]
    fn test_ptx_version() {
        let enc = NvptxMCEncoder::new(75, true);
        assert_eq!(enc.ptx_version(), "7.0");
        let enc2 = NvptxMCEncoder::new(52, false);
        assert_eq!(enc2.ptx_version(), "5.0");
        let enc3 = NvptxMCEncoder::new(35, false);
        assert_eq!(enc3.ptx_version(), "4.3");
    }
}