llvm-native-core 0.1.6

//! NVPTX Assembly Printer — formats PTX machine instructions as
//! assembly text with proper syntax, register names, and directives.
//!
//! Clean-room behavioral reconstruction from the NVIDIA PTX ISA
//! Reference Manual. Zero LLVM source code consultation.
//!
//! This module prints complete functions, basic blocks, and
//! individual instructions in PTX assembly syntax.

use super::nvptx_instr_info::NvptxInstrInfo;
use super::nvptx_register_info::{NvptxRegisterInfo, PRED_BASE, PRED_MAX, VREG_BASE};
use crate::codegen::*;
use crate::module::Module;

// ---------------------------------------------------------------------------
// NvptxAsmPrinter
// ---------------------------------------------------------------------------

/// NVPTX assembly printer: formats machine instructions into
/// PTX assembly text with proper mnemonics, register names,
/// and operand formatting.
///
/// # Output format
/// ```asm
///     .version 7.0
///     .target sm_75
///     .address_size 64
///
///     .visible .entry my_kernel(
///         .param .u64 __retval0
///     )
///     {
///         ld.param.u64    %rd1, [__retval0];
///         add.s32         %r1, %r2, %r3;
///         ret;
///     }
/// ```
pub struct NvptxAsmPrinter {
    /// Accumulated assembly output text.
    pub output: String,
    /// Whether using 64-bit addressing.
    pub is_64bit: bool,
    /// SM compute capability version.
    pub sm_version: u32,
    /// Instruction info for mnemonic lookup.
    pub instr_info: NvptxInstrInfo,
}

impl NvptxAsmPrinter {
    /// Create a new PTX assembly printer.
    pub fn new(is_64bit: bool, sm_version: u32) -> Self {
        NvptxAsmPrinter {
            output: String::with_capacity(8192),
            is_64bit,
            sm_version,
            instr_info: NvptxInstrInfo::new(),
        }
    }

    /// Print a complete function with prologue, body, and epilogue.
    pub fn print_function(&mut self, mf: &MachineFunction) {
        self.print_prologue(mf);

        for bb in &mf.blocks {
            self.print_basic_block(bb);
        }

        self.print_epilogue(mf);
    }

    /// Print a complete module: emit version, target, address_size,
    /// and call `print_function` for each provided function.
    pub fn print_module(&mut self, module: &Module) {
        // Module-level directives
        let ptx_ver = self.ptx_version();
        self.output.push_str(&format!("\t.version\t{}\n", ptx_ver));
        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");
        }

        if !module.source_filename.is_empty() {
            self.output
                .push_str(&format!("\t.file\t\"{}\"\n", module.source_filename));
        }
    }

    /// Emit function prologue: `.visible .entry` and opening brace.
    pub fn print_prologue(&mut self, mf: &MachineFunction) {
        let name = &mf.name;
        self.output
            .push_str(&format!("\t.visible .entry {}(\n", 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{\n");
    }

    /// Emit function epilogue: closing brace and newline.
    pub fn print_epilogue(&mut self, mf: &MachineFunction) {
        self.output.push_str("}\n");
        let name = &mf.name;
        self.output
            .push_str(&format!("\t// -- End function {}\n", name));
    }

    /// Emit a basic block: label followed by all instructions.
    pub fn print_basic_block(&mut self, bb: &MachineBasicBlock) {
        if !bb.name.is_empty() {
            self.output.push_str(&format!(".LBB_{}:\n", bb.name));
        }

        for mi in &bb.instructions {
            self.print_instruction(mi);
        }
    }

    /// Emit a single machine instruction as PTX assembly text.
    pub fn print_instruction(&mut self, mi: &MachineInstr) {
        let mnemonic = self.get_mnemonic_by_opcode(mi.opcode);
        if mnemonic.is_empty() || mnemonic == "INVALID" {
            return; // Skip unknown opcodes
        }

        let mn = mnemonic.as_str();

        // Handle zero-operand instructions
        if mi.operands.is_empty() {
            match mn {
                "ret" => {
                    self.output.push_str("\tret;\n");
                    return;
                }
                "exit" => {
                    self.output.push_str("\texit;\n");
                    return;
                }
                "brkpt" => {
                    self.output.push_str("\tbrkpt;\n");
                    return;
                }
                _ => {
                    self.output.push_str(&format!("\t{};\n", mn));
                    return;
                }
            }
        }

        // Format operands
        let mut op_strs: Vec<String> = Vec::new();
        for op in &mi.operands {
            op_strs.push(self.print_operand(op));
        }

        // Determine type suffix
        let type_sfx = self.get_type_suffix_for_opcode(mi.opcode);

        // Format the instruction line with predicate if present
        // Check if first operand is a predicate reg
        let mnemonic_with_pred = if matches!(&mi.operands[0], MachineOperand::PhysReg(r)
            if *r >= PRED_BASE as u32 && *r <= PRED_MAX as u32)
        {
            let pred_name = self.print_operand(&mi.operands[0]);
            let rest: Vec<String> = mi.operands[1..]
                .iter()
                .map(|op| self.print_operand(op))
                .collect();
            format!("{}\t@{} {}", mn, pred_name, rest.join(", "))
        } else {
            format!("{}{}\t{}", mn, type_sfx, op_strs.join(", "))
        };

        self.output
            .push_str(&format!("\t{};\n", mnemonic_with_pred));
    }

    /// Format a single machine operand as a PTX assembly string.
    pub fn print_operand(&self, op: &MachineOperand) -> String {
        match op {
            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
    // ------------------------------------------------------------------

    /// Look up the PTX mnemonic for a numeric opcode.
    fn get_mnemonic_by_opcode(&self, opcode_val: u32) -> String {
        // Map numeric opcodes to PTX mnemonics
        match opcode_val {
            // Integer
            20000 => "add".to_string(),
            20001 => "sub".to_string(),
            20002 => "mul".to_string(),
            20003 => "mad".to_string(),
            20004 => "div".to_string(),
            20005 => "rem".to_string(),
            20006 => "abs".to_string(),
            20007 => "neg".to_string(),
            20008 => "min".to_string(),
            20009 => "max".to_string(),
            20010 => "popc".to_string(),
            20011 => "clz".to_string(),
            20012 => "bfind".to_string(),
            20013 => "brev".to_string(),
            20014 => "bfe".to_string(),
            20015 => "bfi".to_string(),
            20016 => "sad".to_string(),
            20017 => "selp".to_string(),
            // Float
            20020 => "add".to_string(),
            20021 => "sub".to_string(),
            20022 => "mul".to_string(),
            20023 => "mad".to_string(),
            20024 => "fma".to_string(),
            20025 => "div".to_string(),
            20026 => "abs".to_string(),
            20027 => "neg".to_string(),
            20028 => "min".to_string(),
            20029 => "max".to_string(),
            20030 => "rcp".to_string(),
            20031 => "sqrt".to_string(),
            20032 => "rsqrt".to_string(),
            20033 => "sin".to_string(),
            20034 => "cos".to_string(),
            20035 => "lg2".to_string(),
            20036 => "ex2".to_string(),
            20037 => "tanh".to_string(),
            // Compare
            20040 => "setp.eq".to_string(),
            20041 => "setp.ne".to_string(),
            20042 => "setp.lt".to_string(),
            20043 => "setp.le".to_string(),
            20044 => "setp.gt".to_string(),
            20045 => "setp.ge".to_string(),
            20046 => "setp.lo".to_string(),
            20047 => "setp.ls".to_string(),
            20048 => "setp.hi".to_string(),
            20049 => "setp.hs".to_string(),
            20050 => "setp.equ".to_string(),
            20051 => "setp.neu".to_string(),
            20052 => "setp.ltu".to_string(),
            20053 => "setp.leu".to_string(),
            20054 => "setp.gtu".to_string(),
            20055 => "setp.geu".to_string(),
            20056 => "setp.num".to_string(),
            20057 => "setp.nan".to_string(),
            // Logic
            20060 => "and".to_string(),
            20061 => "or".to_string(),
            20062 => "xor".to_string(),
            20063 => "not".to_string(),
            20064 => "shl".to_string(),
            20065 => "shr".to_string(),
            20066 => "cnot".to_string(),
            // Data movement
            20070 => "ld".to_string(),
            20071 => "st".to_string(),
            20072 => "ld.global".to_string(),
            20073 => "st.global".to_string(),
            20074 => "ld.shared".to_string(),
            20075 => "st.shared".to_string(),
            20076 => "ld.local".to_string(),
            20077 => "st.local".to_string(),
            20078 => "ld.param".to_string(),
            20079 => "st.param".to_string(),
            20080 => "ld.const".to_string(),
            20081 => "ld.global.nc".to_string(),
            20082 => "ld.global.cg".to_string(),
            20083 => "ld.global.ca".to_string(),
            20084 => "cvta".to_string(),
            20085 => "mov".to_string(),
            // Conversion
            20090 => "cvt".to_string(),
            20091 => "cvt.rp".to_string(),
            20092 => "cvt.rn".to_string(),
            20093 => "cvt.rz".to_string(),
            20094 => "cvt.rm".to_string(),
            // Control
            20100 => "bra".to_string(),
            20101 => "call".to_string(),
            20102 => "ret".to_string(),
            20103 => "exit".to_string(),
            20104 => "brkpt".to_string(),
            20105 => "bar".to_string(),
            20106 => "bar.sync".to_string(),
            20107 => "bar.arrive".to_string(),
            20108 => "membar".to_string(),
            // Atomics
            20110 => "atom.add".to_string(),
            20111 => "atom.sub".to_string(),
            20112 => "atom.exch".to_string(),
            20113 => "atom.cas".to_string(),
            20114 => "atom.min".to_string(),
            20115 => "atom.max".to_string(),
            20116 => "atom.and".to_string(),
            20117 => "atom.or".to_string(),
            20118 => "atom.xor".to_string(),
            // Reduction
            20120 => "red.add".to_string(),
            20121 => "red.sub".to_string(),
            20122 => "red.min".to_string(),
            20123 => "red.max".to_string(),
            20124 => "red.and".to_string(),
            20125 => "red.or".to_string(),
            20126 => "red.xor".to_string(),
            // Vote
            20130 => "vote.all".to_string(),
            20131 => "vote.any".to_string(),
            20132 => "vote.uni".to_string(),
            20133 => "vote.ballot".to_string(),
            // Texture
            20140 => "tex".to_string(),
            20141 => "tld4".to_string(),
            20142 => "txq".to_string(),
            20143 => "suld".to_string(),
            20144 => "sust".to_string(),
            20145 => "sured".to_string(),
            20146 => "suq".to_string(),
            // Special
            20150 => "isspacep".to_string(),
            _ => "INVALID".to_string(),
        }
    }

    /// Get the PTX type suffix for an opcode.
    fn get_type_suffix_for_opcode(&self, opcode_val: u32) -> String {
        match opcode_val {
            // Integer ops
            20000..=20009 => ".s32".to_string(),
            20010 | 20011 => ".b32".to_string(),
            20012..=20017 => ".b32".to_string(),
            // Float ops
            20020..=20029 => ".f32".to_string(),
            20030..=20037 => ".f32".to_string(),
            // Compare
            20040..=20057 => ".s32".to_string(),
            // Logic
            20060..=20066 => ".b32".to_string(),
            // Loads/stores
            20070..=20083 => ".b32".to_string(),
            20084 => ".to".to_string(),
            20085 => ".b32".to_string(),
            // Conversion
            20090..=20094 => "".to_string(),
            // Control (no suffix)
            20100..=20108 => "".to_string(),
            // Atomics
            20110..=20118 => ".global.b32".to_string(),
            // Reductions
            20120..=20126 => ".shared.b32".to_string(),
            // Vote (no suffix)
            20130..=20133 => ".sync".to_string(),
            _ => "".to_string(),
        }
    }

    /// Determine the PTX version string for the given SM version.
    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"
        }
    }

    /// Clear the output buffer.
    pub fn clear(&mut self) {
        self.output.clear();
    }
}

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

    #[test]
    fn test_printer_new() {
        let printer = NvptxAsmPrinter::new(true, 75);
        assert!(printer.is_64bit);
        assert_eq!(printer.sm_version, 75);
    }

    #[test]
    fn test_print_operand_imm() {
        let printer = NvptxAsmPrinter::new(true, 75);
        let op = MachineOperand::Imm(42);
        assert_eq!(printer.print_operand(&op), "42");
    }

    #[test]
    fn test_print_operand_label() {
        let printer = NvptxAsmPrinter::new(true, 75);
        let op = MachineOperand::Label("target".to_string());
        assert_eq!(printer.print_operand(&op), ".LLabel_target");
    }

    #[test]
    fn test_get_mnemonic_by_opcode() {
        let printer = NvptxAsmPrinter::new(true, 75);
        assert_eq!(printer.get_mnemonic_by_opcode(20000), "add");
        assert_eq!(printer.get_mnemonic_by_opcode(20102), "ret");
        assert_eq!(printer.get_mnemonic_by_opcode(99999), "INVALID");
    }

    #[test]
    fn test_ptx_version() {
        let p = NvptxAsmPrinter::new(true, 75);
        assert_eq!(p.ptx_version(), "7.0");
        let p2 = NvptxAsmPrinter::new(false, 35);
        assert_eq!(p2.ptx_version(), "4.3");
    }
}