llvm-native-core 0.1.12

//! AMDGPU Instruction Selection — selects AMDGPU instructions from
//! LLVM IR operations.
//!
//! Maps generic IR operations to AMDGPU-specific instructions based on
//! operand types, available register classes, and target features.
//!
//! Clean-room reconstruction from AMD GCN ISA manuals.
//! Zero LLVM source code consultation.

use std::collections::HashMap;

use crate::basic_block::BasicBlock;
use crate::codegen::{MachineBasicBlock, MachineFunction, MachineInstr, MachineOperand};
use crate::function::Function;
use crate::opcode::Opcode;
use crate::types::Type;
use crate::value::{Value, ValueRef};

use super::amdgpu_instr_info::{AmdgpuInstrInfo, AmdgpuOpcode};
use super::amdgpu_register_info::AmdgpuRegisterInfo;
use super::amdgpu_target_machine::AmdgpuTargetMachine;

/// Selects AMDGPU-specific instructions from LLVM IR operations.
pub struct AmdgpuInstructionSelector {
    pub instr_info: AmdgpuInstrInfo,
    pub reg_info: AmdgpuRegisterInfo,
    pub target: AmdgpuTargetMachine,
}

impl AmdgpuInstructionSelector {
    pub fn new(target: AmdgpuTargetMachine) -> Self {
        Self {
            instr_info: AmdgpuInstrInfo::new(),
            reg_info: AmdgpuRegisterInfo::new(),
            target,
        }
    }

    /// Select an AMDGPU instruction for a given IR opcode and operand types.
    pub fn select(
        &self,
        ir_opcode: Opcode,
        dst_type: &Type,
        src_types: &[Type],
    ) -> Option<AmdgpuOpcode> {
        let is_f32 = Self::is_f32(dst_type);
        let is_f64 = Self::is_f64(dst_type);
        let is_i32 = Self::is_i32(dst_type);
        let is_v2f16 = Self::is_v2f16(dst_type);

        match ir_opcode {
            // ── Integer Arithmetic ─────────────────────────────────
            Opcode::Add => {
                if is_i32 {
                    Some(AmdgpuOpcode::VAddF32)
                } else if is_f32 {
                    Some(AmdgpuOpcode::VAddF32)
                } else if is_f64 {
                    Some(AmdgpuOpcode::VAddF64)
                } else {
                    None
                }
            }
            Opcode::Sub => {
                if is_f32 {
                    Some(AmdgpuOpcode::VSubF32)
                } else if is_f64 {
                    Some(AmdgpuOpcode::VSubF64)
                } else {
                    None
                }
            }
            Opcode::Mul => {
                if is_f32 {
                    Some(AmdgpuOpcode::VMulF32Vop3)
                } else if is_f64 {
                    Some(AmdgpuOpcode::VMulF64)
                } else if is_v2f16 {
                    Some(AmdgpuOpcode::VPkMulF16)
                } else {
                    None
                }
            }
            Opcode::SDiv | Opcode::UDiv => None, // Software emulated
            Opcode::SRem | Opcode::URem => None, // Software emulated

            // ── Floating-Point ────────────────────────────────────
            Opcode::FAdd => {
                if is_f32 {
                    Some(AmdgpuOpcode::VAddF32)
                } else if is_f64 {
                    Some(AmdgpuOpcode::VAddF64)
                } else {
                    None
                }
            }
            Opcode::FSub => {
                if is_f32 {
                    Some(AmdgpuOpcode::VSubF32)
                } else {
                    None
                }
            }
            Opcode::FMul => {
                if is_f32 {
                    Some(AmdgpuOpcode::VMulF32Vop3)
                } else {
                    None
                }
            }
            Opcode::FDiv => {
                if is_f32 {
                    Some(AmdgpuOpcode::VRcpF32)
                }
                // Approximate; needs refinement
                else {
                    None
                }
            }
            Opcode::FRem => None,

            // ── Bitwise ───────────────────────────────────────────
            Opcode::And => Some(AmdgpuOpcode::VAndB32),
            Opcode::Or => Some(AmdgpuOpcode::VOrB32),
            Opcode::Xor => Some(AmdgpuOpcode::VXorB32),
            Opcode::Shl => Some(AmdgpuOpcode::VLshlrevB32),
            Opcode::LShr => Some(AmdgpuOpcode::VLshrrevB32),
            Opcode::AShr => Some(AmdgpuOpcode::VAshrrevI32),

            // ── Comparison ────────────────────────────────────────
            Opcode::ICmp => Some(AmdgpuOpcode::VCmpI32Eq),
            Opcode::FCmp => Some(AmdgpuOpcode::VCmpF32Eq),

            // ── Conversions ───────────────────────────────────────
            Opcode::SIToFP => {
                if is_f32 {
                    Some(AmdgpuOpcode::VCvtF32I32)
                } else {
                    None
                }
            }
            Opcode::FPToSI => {
                if is_f32 {
                    Some(AmdgpuOpcode::VCvtI32F32)
                } else {
                    None
                }
            }

            // ── Memory ────────────────────────────────────────────
            Opcode::Load => Some(AmdgpuOpcode::GlobalLoadDword),
            Opcode::Store => Some(AmdgpuOpcode::GlobalStoreDword),

            // ── Moves / Copies ────────────────────────────────────
            Opcode::Phi | Opcode::Select => Some(AmdgpuOpcode::VMovB32),

            _ => None,
        }
    }

    /// Select a fused multiply-add.
    pub fn select_fma(&self, is_f32: bool) -> AmdgpuOpcode {
        if is_f32 {
            AmdgpuOpcode::VFmaF32
        } else {
            AmdgpuOpcode::VMadF64
        }
    }

    /// Check if a comparison should write VCC.
    pub fn cmp_writes_vcc(&self, opcode: AmdgpuOpcode) -> bool {
        matches!(
            opcode,
            AmdgpuOpcode::VCmpF32Eq
                | AmdgpuOpcode::VCmpF32Lt
                | AmdgpuOpcode::VCmpF32Gt
                | AmdgpuOpcode::VCmpU32Eq
                | AmdgpuOpcode::VCmpI32Eq
        )
    }

    /// Whether to use VOP3 encoding (more operands, more flexibility)
    /// instead of VOP1/VOP2.
    pub fn use_vop3(&self, num_srcs: u32, has_imm: bool) -> bool {
        num_srcs > 2 || has_imm
    }

    // ── Type Helpers ─────────────────────────────────────────────────

    fn is_f32(ty: &Type) -> bool {
        ty.is_float()
    }
    fn is_f64(ty: &Type) -> bool {
        ty.is_float()
    }
    fn is_i32(ty: &Type) -> bool {
        ty.is_integer()
    }
    fn is_v2f16(ty: &Type) -> bool {
        false // V2F16 type not yet defined in type system
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Type, TypeKind};

    #[test]
    fn test_select_add_f32() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        let f32 = Type::new(TypeKind::Float);
        let result = isel.select(Opcode::Add, &f32, &[]);
        assert_eq!(result, Some(AmdgpuOpcode::VAddF32));
    }

    #[test]
    fn test_select_mul_f32() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        let f32 = Type::new(TypeKind::Float);
        let result = isel.select(Opcode::Mul, &f32, &[]);
        assert_eq!(result, Some(AmdgpuOpcode::VMulF32Vop3));
    }

    #[test]
    fn test_select_and() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        let i32 = Type::new(TypeKind::Integer { bits: 32 });
        let result = isel.select(Opcode::And, &i32, &[]);
        assert_eq!(result, Some(AmdgpuOpcode::VAndB32));
    }

    #[test]
    fn test_select_load() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        let i32 = Type::new(TypeKind::Integer { bits: 32 });
        let result = isel.select(Opcode::Load, &i32, &[]);
        assert_eq!(result, Some(AmdgpuOpcode::GlobalLoadDword));
    }

    #[test]
    fn test_select_store() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        let i32 = Type::new(TypeKind::Integer { bits: 32 });
        let result = isel.select(Opcode::Store, &i32, &[]);
        assert_eq!(result, Some(AmdgpuOpcode::GlobalStoreDword));
    }

    #[test]
    fn test_select_fma() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        assert_eq!(isel.select_fma(true), AmdgpuOpcode::VFmaF32);
        assert_eq!(isel.select_fma(false), AmdgpuOpcode::VMadF64);
    }

    #[test]
    fn test_cmp_writes_vcc() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        assert!(isel.cmp_writes_vcc(AmdgpuOpcode::VCmpF32Eq));
        assert!(!isel.cmp_writes_vcc(AmdgpuOpcode::VAddF32));
    }

    #[test]
    fn test_use_vop3() {
        let target = AmdgpuTargetMachine::from_cpu("gfx900").unwrap();
        let isel = AmdgpuInstructionSelector::new(target);
        assert!(isel.use_vop3(3, false));
        assert!(!isel.use_vop3(2, false));
    }
}