llvm-native-core 0.1.15

//! NVPTX Instruction Selection — converts LLVM IR to PTX
//! machine instructions.
//!
//! Clean-room behavioral reconstruction from the NVIDIA PTX ISA
//! Reference Manual. Zero LLVM source code consultation.
//!
//! The instruction selector lowers LLVM IR opcodes to PTX
//! virtual instructions. PTX uses unlimited virtual registers
//! in SSA form with a 3-operand format.

use super::nvptx_instr_info::NvptxOpcode;
use super::nvptx_register_info::*;
use crate::codegen::*;
use crate::opcode::Opcode;
use crate::value::Value;
use std::collections::HashMap;

// ---------------------------------------------------------------------------
// NvptxInstructionSelector
// ---------------------------------------------------------------------------

/// NVPTX instruction selector: lowers LLVM IR instructions into
/// PTX virtual machine instructions.
///
/// # Lowering patterns
///
/// | IR op       | PTX instruction(s)                            |
/// |-------------|-----------------------------------------------|
/// | `add`       | `add.s32 rd, rs1, rs2`                        |
/// | `sub`       | `sub.s32 rd, rs1, rs2`                        |
/// | `mul`       | `mul.lo.s32 rd, rs1, rs2`                     |
/// | `sdiv`      | `div.s32 rd, rs1, rs2`                        |
/// | `udiv`      | `div.u32 rd, rs1, rs2`                        |
/// | `and`       | `and.b32 rd, rs1, rs2`                        |
/// | `or`        | `or.b32 rd, rs1, rs2`                         |
/// | `xor`       | `xor.b32 rd, rs1, rs2`                        |
/// | `shl`       | `shl.b32 rd, rs1, rs2`                        |
/// | `lshr`      | `shr.u32 rd, rs1, rs2`                        |
/// | `fadd`      | `add.f32 rd, rs1, rs2`                        |
/// | `fsub`      | `sub.f32 rd, rs1, rs2`                        |
/// | `fmul`      | `mul.f32 rd, rs1, rs2`                        |
/// | `fdiv`      | `div.f32 rd, rs1, rs2`                        |
/// | `icmp`      | `setp.* p, rs1, rs2; selp rd, 1, 0, p`       |
/// | `ret`       | `ret`                                         |
/// | `br`        | `bra label` or `setp.ne + bra`                |
/// | `call`      | `call func`                                   |
/// | `load`      | `ld.global rd, [rs]`                          |
/// | `store`     | `st.global [rd], rs`                          |
pub struct NvptxInstructionSelector {
    /// Map from IR value IDs to virtual register numbers.
    pub vreg_map: HashMap<usize, VirtReg>,
    /// The current machine basic block being built.
    pub mbb: MachineBasicBlock,
    /// Name of the function being compiled.
    pub func_name: String,
    /// Virtual predicate register counter.
    pub pred_counter: u32,
}

impl NvptxInstructionSelector {
    /// Create a new instruction selector.
    pub fn new() -> Self {
        NvptxInstructionSelector {
            vreg_map: HashMap::new(),
            mbb: MachineBasicBlock {
                name: String::new(),
                instructions: Vec::new(),
                successors: Vec::new(),
            },
            func_name: String::new(),
            pred_counter: 0,
        }
    }

    /// Get (or create) a virtual register for an IR value.
    fn get_or_create_vreg(&mut self, vid: usize, mf: &mut MachineFunction) -> VirtReg {
        if let Some(&vr) = self.vreg_map.get(&vid) {
            return vr;
        }
        let vr = mf.new_vreg();
        self.vreg_map.insert(vid, vr);
        vr
    }

    /// Get a fresh predicate register number.
    fn new_pred(&mut self) -> u32 {
        let p = self.pred_counter;
        self.pred_counter += 1;
        p
    }

    // ==================================================================
    // Top-level selection
    // ==================================================================

    /// Convert an entire LLVM function into PTX machine instructions.
    pub fn select(&mut self, mf: &mut MachineFunction, func: &Value) {
        self.func_name = func.name.clone();
        if self.func_name.is_empty() {
            self.func_name = format!(".Lfunc{}", func.vid);
        }
        self.vreg_map.clear();
        self.pred_counter = 0;

        for bb_ref in &func.successors {
            let bb = bb_ref.borrow();
            self.mbb = MachineBasicBlock {
                name: bb.name.clone(),
                instructions: Vec::new(),
                successors: Vec::new(),
            };

            for inst_ref in &bb.operands {
                let inst = inst_ref.borrow();
                if inst.is_instruction() {
                    let instrs = self.select_instruction(mf, &inst);
                    self.mbb.instructions.extend(instrs);
                }
            }

            mf.push_block(self.mbb.clone());
        }
    }

    /// Select machine instructions for a single IR instruction.
    pub fn select_instruction(
        &mut self,
        mf: &mut MachineFunction,
        inst: &Value,
    ) -> Vec<MachineInstr> {
        let opcode = match inst.get_opcode() {
            Some(op) => op,
            None => return Vec::new(),
        };

        match opcode {
            Opcode::Add => self.lower_int_binop(mf, inst, NvptxOpcode::Add),
            Opcode::Sub => self.lower_int_binop(mf, inst, NvptxOpcode::Sub),
            Opcode::Mul => self.lower_int_binop(mf, inst, NvptxOpcode::Mul),
            Opcode::SDiv => self.lower_int_binop(mf, inst, NvptxOpcode::Div),
            Opcode::UDiv => self.lower_int_binop(mf, inst, NvptxOpcode::Div),
            Opcode::SRem => self.lower_int_binop(mf, inst, NvptxOpcode::Rem),
            Opcode::URem => self.lower_int_binop(mf, inst, NvptxOpcode::Rem),
            Opcode::And => self.lower_int_binop(mf, inst, NvptxOpcode::And),
            Opcode::Or => self.lower_int_binop(mf, inst, NvptxOpcode::Or),
            Opcode::Xor => self.lower_int_binop(mf, inst, NvptxOpcode::Xor),
            Opcode::Shl => self.lower_int_binop(mf, inst, NvptxOpcode::Shl),
            Opcode::LShr => self.lower_int_binop(mf, inst, NvptxOpcode::Shr),
            Opcode::FAdd => self.lower_fp_binop(mf, inst, NvptxOpcode::FAdd),
            Opcode::FSub => self.lower_fp_binop(mf, inst, NvptxOpcode::FSub),
            Opcode::FMul => self.lower_fp_binop(mf, inst, NvptxOpcode::FMul),
            Opcode::FDiv => self.lower_fp_binop(mf, inst, NvptxOpcode::FDiv),
            Opcode::ICmp => self.lower_icmp(mf, inst),
            Opcode::FCmp => self.lower_fcmp(mf, inst),
            Opcode::Br => self.lower_br(mf, inst),
            Opcode::Ret => self.lower_ret(mf, inst),
            Opcode::Call => self.lower_call(mf, inst),
            Opcode::Alloca => self.lower_alloca(mf, inst),
            Opcode::Load => self.lower_load(mf, inst),
            Opcode::Store => self.lower_store(mf, inst),
            Opcode::ZExt | Opcode::SExt | Opcode::Trunc | Opcode::BitCast => {
                self.lower_cast(mf, inst)
            }
            Opcode::Select => self.lower_select(mf, inst),
            Opcode::Phi => Vec::new(),
            _ => Vec::new(),
        }
    }

    // ==================================================================
    // Binary arithmetic lowering
    // ==================================================================

    fn lower_int_binop(
        &mut self,
        mf: &mut MachineFunction,
        inst: &Value,
        nvop: NvptxOpcode,
    ) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut mi = MachineInstr::new(nvop as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_reg(src2_vr);
        vec![mi]
    }

    fn lower_fp_binop(
        &mut self,
        mf: &mut MachineFunction,
        inst: &Value,
        nvop: NvptxOpcode,
    ) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut mi = MachineInstr::new(nvop as u32).with_def(def_vr);
        mi.push_reg(src1_vr);
        mi.push_reg(src2_vr);
        vec![mi]
    }

    // ==================================================================
    // ICmp lowering
    // ==================================================================

    /// Lower ICmp to setp + selp.
    ///
    /// ```text
    /// setp.eq.s32  %pN, rs1, rs2
    /// selp.b32     rd, 1, 0, %pN
    /// ```
    fn lower_icmp(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let pred_idx = self.new_pred();

        // setp.eq.s32 %pN, rs1, rs2
        let mut setp = MachineInstr::new(NvptxOpcode::SetpEq as u32);
        setp.push_reg(pred_idx + PRED_BASE as u32);
        setp.push_reg(src1_vr);
        setp.push_reg(src2_vr);
        instrs.push(setp);

        // selp.b32 rd, 1, 0, %pN
        let mut selp = MachineInstr::new(NvptxOpcode::Selp as u32).with_def(def_vr);
        selp.push_imm(1);
        selp.push_imm(0);
        selp.push_reg(pred_idx + PRED_BASE as u32);
        instrs.push(selp);

        instrs
    }

    /// Lower FCmp (defaults to ordered-eq).
    fn lower_fcmp(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src1_vr = self.resolve_operand(mf, &inst.operands[0]);
        let src2_vr = self.resolve_operand(mf, &inst.operands[1]);
        let pred_idx = self.new_pred();

        let mut setp = MachineInstr::new(NvptxOpcode::SetpEq as u32);
        setp.push_reg(pred_idx + PRED_BASE as u32);
        setp.push_reg(src1_vr);
        setp.push_reg(src2_vr);
        instrs.push(setp);

        let mut selp = MachineInstr::new(NvptxOpcode::Selp as u32).with_def(def_vr);
        selp.push_imm(1);
        selp.push_imm(0);
        selp.push_reg(pred_idx + PRED_BASE as u32);
        instrs.push(selp);

        instrs
    }

    // ==================================================================
    // Terminator lowering
    // ==================================================================

    fn lower_br(&mut self, _mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let num_ops = inst.operands.len();

        if num_ops == 1 {
            // Unconditional branch
            let dest = inst.operands[0].borrow().name.clone();
            self.mbb.successors.push(dest.clone());
            let mut bra = MachineInstr::new(NvptxOpcode::Bra as u32);
            bra.push_label(&dest);
            instrs.push(bra);
        } else if num_ops == 3 {
            // Conditional branch: br cond, true_label, false_label
            let t_label = inst.operands[1].borrow().name.clone();
            let f_label = inst.operands[2].borrow().name.clone();
            self.mbb.successors.push(t_label.clone());
            self.mbb.successors.push(f_label.clone());

            let cond_vid = inst.operands[0].borrow().vid as usize;
            if let Some(&cond_vr) = self.vreg_map.get(&cond_vid) {
                let pred_idx = self.new_pred();
                // setp.ne.s32 %pN, cond, 0
                let mut setp = MachineInstr::new(NvptxOpcode::SetpNe as u32);
                setp.push_reg(pred_idx + PRED_BASE as u32);
                setp.push_reg(cond_vr);
                setp.push_imm(0);
                instrs.push(setp);

                // @%pN bra t_label
                let mut bra = MachineInstr::new(NvptxOpcode::Bra as u32);
                bra.push_label(&t_label);
                instrs.push(bra);

                // bra f_label
                let mut bra2 = MachineInstr::new(NvptxOpcode::Bra as u32);
                bra2.push_label(&f_label);
                instrs.push(bra2);
            }
        }

        instrs
    }

    fn lower_ret(&mut self, _mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();

        if !inst.operands.is_empty() {
            // PTX: st.param for return value handled by st.param
            let val_vid = inst.operands[0].borrow().vid as usize;
            if let Some(&_vr) = self.vreg_map.get(&val_vid) {
                // mov to return param handled by st.param
                let mut mov = MachineInstr::new(NvptxOpcode::StParam as u32);
                mov.push_reg(_vr);
                mov.push_imm(0);
                instrs.push(mov);
            }
        }

        instrs.push(MachineInstr::new(NvptxOpcode::Ret as u32));
        instrs
    }

    fn lower_call(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        if inst.operands.is_empty() {
            return instrs;
        }

        let callee_name = inst.operands[0].borrow().name.clone();

        // Move arguments via st.param
        for i in 1..inst.operands.len() {
            let arg_vid = inst.operands[i].borrow().vid as usize;
            if let Some(&arg_vr) = self.vreg_map.get(&arg_vid) {
                let mut st_param = MachineInstr::new(NvptxOpcode::StParam as u32);
                st_param.push_reg(arg_vr);
                st_param.push_imm(i as i64 - 1);
                instrs.push(st_param);
            }
        }

        let mut call = MachineInstr::new(NvptxOpcode::Call as u32);
        call.operands.push(MachineOperand::Global(callee_name));
        instrs.push(call);

        // Load return value if not void
        if !inst.ty.is_void() {
            let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
            let mut ld_param = MachineInstr::new(NvptxOpcode::LdParam as u32).with_def(def_vr);
            ld_param.push_imm(0);
            instrs.push(ld_param);
        }

        instrs
    }

    // ==================================================================
    // Memory lowering
    // ==================================================================

    fn lower_alloca(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        // PTX allocates local memory via .local directive at function scope.
        // At instruction level, we emit cvta.local to get the address.
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let mut mov = MachineInstr::new(NvptxOpcode::Mov as u32).with_def(def_vr);
        mov.push_imm(0); // Local address handled by register alloc
        vec![mov]
    }

    fn lower_load(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src_vr = self.resolve_operand(mf, &inst.operands[0]);
        let mut ld = MachineInstr::new(NvptxOpcode::LdGlobal as u32).with_def(def_vr);
        ld.push_reg(src_vr);
        vec![ld]
    }

    fn lower_store(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let val_vr = self.resolve_operand(mf, &inst.operands[0]);
        let ptr_vr = self.resolve_operand(mf, &inst.operands[1]);
        let mut st = MachineInstr::new(NvptxOpcode::StGlobal as u32);
        st.push_reg(ptr_vr);
        st.push_reg(val_vr);
        vec![st]
    }

    // ==================================================================
    // Cast lowering
    // ==================================================================

    fn lower_cast(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let src_vr = self.resolve_operand(mf, &inst.operands[0]);
        let mut cvt = MachineInstr::new(NvptxOpcode::Cvt as u32).with_def(def_vr);
        cvt.push_reg(src_vr);
        vec![cvt]
    }

    // ==================================================================
    // Select lowering
    // ==================================================================

    fn lower_select(&mut self, mf: &mut MachineFunction, inst: &Value) -> Vec<MachineInstr> {
        let mut instrs = Vec::new();
        let def_vr = self.get_or_create_vreg(inst.vid as usize, mf);
        let true_vr = self.resolve_operand(mf, &inst.operands[1]);
        let false_vr = self.resolve_operand(mf, &inst.operands[2]);
        let cond_vid = inst.operands[0].borrow().vid as usize;

        if let Some(&cond_vr) = self.vreg_map.get(&cond_vid) {
            let pred_idx = self.new_pred();
            // setp.ne.s32 %pN, cond, 0
            let mut setp = MachineInstr::new(NvptxOpcode::SetpNe as u32);
            setp.push_reg(pred_idx + PRED_BASE as u32);
            setp.push_reg(cond_vr);
            setp.push_imm(0);
            instrs.push(setp);

            // selp.b32 rd, true, false, %pN
            let mut selp = MachineInstr::new(NvptxOpcode::Selp as u32).with_def(def_vr);
            selp.push_reg(true_vr);
            selp.push_reg(false_vr);
            selp.push_reg(pred_idx + PRED_BASE as u32);
            instrs.push(selp);
        }

        instrs
    }

    // ==================================================================
    // Helpers
    // ==================================================================

    /// Resolve an operand IR value to a virtual register, creating one
    /// if it's a constant and we need a register.
    fn resolve_operand(
        &mut self,
        mf: &mut MachineFunction,
        val_ref: &crate::value::ValueRef,
    ) -> VirtReg {
        let val = val_ref.borrow();
        if let Some(&vr) = self.vreg_map.get(&(val.vid as usize)) {
            vr
        } else if val.is_constant() {
            let vr = mf.new_vreg();
            self.vreg_map.insert(val.vid as usize, vr);
            vr
        } else {
            let vr = mf.new_vreg();
            self.vreg_map.insert(val.vid as usize, vr);
            vr
        }
    }
}

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

    #[test]
    fn test_selector_new() {
        let sel = NvptxInstructionSelector::new();
        assert!(sel.vreg_map.is_empty());
        assert_eq!(sel.pred_counter, 0);
    }

    #[test]
    fn test_new_pred() {
        let mut sel = NvptxInstructionSelector::new();
        assert_eq!(sel.new_pred(), 0);
        assert_eq!(sel.new_pred(), 1);
        assert_eq!(sel.new_pred(), 2);
    }
}