llvm-native-core 0.1.6

//! NVPTX Target Machine — PTX target description, compute capability
//! selection, PTX ISA version selection, address space mapping, and
//! kernel metadata emission.
//!
//! Clean-room behavioral reconstruction from the NVIDIA PTX ISA
//! Reference Manual and CUDA Toolkit documentation.
//! Zero LLVM source code consultation.

use super::nvptx_instr_info::{NvptxInstrInfo, NvptxOpcode};
use super::nvptx_mc_encoder::{ComputeCapability, ComputeFeature, NvptxMCEncoder, PtxType};
use super::nvptx_register_info::NvptxRegisterInfo;
use crate::codegen::*;
use crate::module::Module;
use std::collections::HashMap;

// ============================================================================
// Address Space Mapping
// ============================================================================

/// NVPTX address space identifiers as used in LLVM IR's addrspace(N) attribute.
/// Maps to PTX state spaces.
pub mod address_space {
    /// Global memory (.global) — visible to all threads, grid lifetime.
    pub const GLOBAL: u32 = 1;
    /// Constant memory (.const) — read-only, cached, grid lifetime.
    pub const CONSTANT: u32 = 4;
    /// Shared memory (.shared) — per-block, visible to all threads in CTA.
    pub const SHARED: u32 = 3;
    /// Local memory (.local) — per-thread stack/register spilling.
    pub const LOCAL: u32 = 5;
    /// Parameter memory (.param) — kernel parameters, read-only in kernel.
    pub const PARAM: u32 = 101;
    /// Generic address space (unified addressing, sm_60+).
    pub const GENERIC: u32 = 0;
    /// Texture memory (read-only, cached).
    pub const TEXTURE: u32 = 2;
    /// Surface memory.
    pub const SURFACE: u32 = 102;
}

/// Convert an LLVM IR address space number to a PTX state space name.
pub fn addrspace_to_ptx_space(addrspace: u32) -> &'static str {
    match addrspace {
        address_space::GLOBAL => "global",
        address_space::CONSTANT => "const",
        address_space::SHARED => "shared",
        address_space::LOCAL => "local",
        address_space::PARAM => "param",
        address_space::GENERIC => "generic",
        address_space::TEXTURE => "tex",
        address_space::SURFACE => "surf",
        _ => "global",
    }
}

/// Detailed address space properties.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct AddressSpaceProperties {
    /// The PTX state space name.
    pub name: &'static str,
    /// Whether this space is readable.
    pub is_readable: bool,
    /// Whether this space is writable.
    pub is_writable: bool,
    /// Whether this space supports atomic operations.
    pub supports_atomics: bool,
    /// Whether this space is coherent across the grid.
    pub is_grid_coherent: bool,
    /// Whether this space is cached in L1.
    pub l1_cached: bool,
    /// Whether this space is cached in L2.
    pub l2_cached: bool,
    /// Maximum access size in bytes.
    pub max_access_size: u32,
    /// Alignment requirement in bytes.
    pub alignment: u32,
}

/// Get properties for a given address space.
pub fn get_address_space_properties(addrspace: u32) -> AddressSpaceProperties {
    match addrspace {
        address_space::GLOBAL => AddressSpaceProperties {
            name: "global",
            is_readable: true,
            is_writable: true,
            supports_atomics: true,
            is_grid_coherent: false,
            l1_cached: true,
            l2_cached: true,
            max_access_size: 16,
            alignment: 1,
        },
        address_space::CONSTANT => AddressSpaceProperties {
            name: "const",
            is_readable: true,
            is_writable: false,
            supports_atomics: false,
            is_grid_coherent: true,
            l1_cached: true,
            l2_cached: true,
            max_access_size: 16,
            alignment: 1,
        },
        address_space::SHARED => AddressSpaceProperties {
            name: "shared",
            is_readable: true,
            is_writable: true,
            supports_atomics: true,
            is_grid_coherent: false,
            l1_cached: false,
            l2_cached: false,
            max_access_size: 16,
            alignment: 4,
        },
        address_space::LOCAL => AddressSpaceProperties {
            name: "local",
            is_readable: true,
            is_writable: true,
            supports_atomics: false,
            is_grid_coherent: false,
            l1_cached: true,
            l2_cached: false,
            max_access_size: 16,
            alignment: 1,
        },
        address_space::PARAM => AddressSpaceProperties {
            name: "param",
            is_readable: true,
            is_writable: false,
            supports_atomics: false,
            is_grid_coherent: true,
            l1_cached: true,
            l2_cached: true,
            max_access_size: 16,
            alignment: 1,
        },
        _ => AddressSpaceProperties {
            name: "global",
            is_readable: true,
            is_writable: true,
            supports_atomics: true,
            is_grid_coherent: false,
            l1_cached: true,
            l2_cached: true,
            max_access_size: 16,
            alignment: 1,
        },
    }
}

// ============================================================================
// NVPTX Target Machine
// ============================================================================

/// NVPTX target machine — configures code generation for NVIDIA GPUs
/// targeting the PTX virtual ISA.
pub struct NvptxTargetMachine {
    /// Target triple string (e.g., nvptx64-nvidia-cuda).
    pub target_triple: String,
    /// Compute capability (SM version).
    pub sm_version: u32,
    /// PTX ISA version string.
    pub ptx_version: String,
    /// Whether 64-bit addressing is used.
    pub is_64bit: bool,
    /// Optimization level (0-3).
    pub opt_level: String,
    /// CPU/GPU feature flags.
    pub features: Vec<String>,
    /// Data layout string.
    pub data_layout: String,
    /// Instruction info database.
    pub instr_info: NvptxInstrInfo,
    /// Register info.
    pub register_info: NvptxRegisterInfo,
    /// Maximum threads per block.
    pub max_threads_per_block: u32,
    /// Allowed register count per thread.
    pub max_registers_per_thread: u32,
    /// Allocation of shared memory per block in bytes.
    pub shared_memory_per_block: u32,
}

impl NvptxTargetMachine {
    /// Create a new NVPTX target machine with the given SM version.
    pub fn new(sm_version: u32, is_64bit: bool) -> Self {
        let cap = ComputeCapability::from_sm(sm_version);
        NvptxTargetMachine {
            target_triple: format!("nvptx{}-nvidia-cuda", if is_64bit { "64" } else { "" }),
            sm_version,
            ptx_version: cap.ptx_version().to_string(),
            is_64bit,
            opt_level: "O2".to_string(),
            features: Vec::new(),
            data_layout: if is_64bit {
                "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n32:64".to_string()
            } else {
                "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v16:16:16-v32:32:32-v64:64:64-v128:128:128-n32".to_string()
            },
            instr_info: NvptxInstrInfo::new(),
            register_info: NvptxRegisterInfo,
            max_threads_per_block: 1024,
            max_registers_per_thread: 255,
            shared_memory_per_block: cap.max_shared_memory(),
        }
    }

    /// Set optimization level.
    pub fn set_opt_level(&mut self, level: &str) {
        self.opt_level = level.to_string();
    }

    /// Enable a feature flag.
    pub fn enable_feature(&mut self, feature: &str) {
        if !self.features.contains(&feature.to_string()) {
            self.features.push(feature.to_string());
        }
    }

    /// Check if a feature is enabled.
    pub fn has_feature(&self, feature: &str) -> bool {
        self.features.contains(&feature.to_string())
    }

    /// Get the compute capability for this machine.
    pub fn compute_capability(&self) -> ComputeCapability {
        ComputeCapability::from_sm(self.sm_version)
    }

    /// Check if a compute feature is supported.
    pub fn supports(&self, feature: ComputeFeature) -> bool {
        self.compute_capability().supports(feature)
    }

    /// Get the data layout string.
    pub fn get_data_layout(&self) -> &str {
        &self.data_layout
    }

    /// Get the target triple.
    pub fn get_target_triple(&self) -> &str {
        &self.target_triple
    }

    /// Return if 64-bit addressing is used.
    pub fn is_64bit(&self) -> bool {
        self.is_64bit
    }

    /// Get the PTX version string.
    pub fn get_ptx_version(&self) -> &str {
        &self.ptx_version
    }

    /// Get a list of IR passes appropriate for this target and opt level.
    pub fn add_ir_passes(&self) -> Vec<&str> {
        match self.opt_level.as_str() {
            "O0" => vec!["mem2reg"],
            "O1" => vec!["mem2reg", "instcombine", "simplifycfg", "dce"],
            "O2" | "O3" => vec![
                "mem2reg",
                "instcombine",
                "simplifycfg",
                "gvn",
                "dce",
                "inline",
                "licm",
                "sccp",
                "loop-unroll",
                "slp-vectorize",
            ],
            _ => vec!["mem2reg", "instcombine", "simplifycfg"],
        }
    }

    /// Create an MC encoder for this target machine.
    pub fn create_mc_encoder(&self) -> NvptxMCEncoder {
        NvptxMCEncoder::new(self.sm_version, self.is_64bit)
    }

    /// Emit a complete PTX module from a machine function.
    pub fn emit_module(&self, mf: &MachineFunction) -> String {
        let mut encoder = self.create_mc_encoder();
        encoder.encode_module(mf)
    }
}

impl Default for NvptxTargetMachine {
    fn default() -> Self {
        Self::new(70, true)
    }
}

// ============================================================================
// Kernel Metadata Emission
// ============================================================================

/// Kernel argument metadata: describes the PTX .param entries for a
/// kernel function.
#[derive(Debug, Clone)]
pub struct KernelArg {
    /// Parameter name.
    pub name: String,
    /// PTX type of the parameter.
    pub ty: PtxType,
    /// Address space of the parameter (0 = generic, 101 = .param).
    pub addrspace: u32,
    /// Size in bytes.
    pub size: u32,
    /// Alignment in bytes.
    pub alignment: u32,
}

impl KernelArg {
    pub fn new(name: &str, ty: PtxType) -> Self {
        KernelArg {
            name: name.to_string(),
            ty,
            addrspace: address_space::PARAM,
            size: ty.size_bytes(),
            alignment: 4,
        }
    }

    /// Emit as a .param directive in PTX.
    pub fn emit_ptx(&self) -> String {
        format!(
            "\t.param {}{} {}",
            self.ty.to_ptx(),
            if self.alignment > 1 {
                format!(".align {}", self.alignment)
            } else {
                String::new()
            },
            self.name
        )
    }
}

/// Kernel function metadata (emitted as PTX directives).
#[derive(Debug, Clone)]
pub struct KernelMetadata {
    /// Kernel name.
    pub name: String,
    /// Input parameters.
    pub params: Vec<KernelArg>,
    /// Return type (None if void).
    pub return_type: Option<KernelArg>,
    /// Required number of registers per thread.
    pub max_registers: u32,
    /// Required shared memory per block (dynamic, in bytes).
    pub dynamic_shared_memory: u32,
    /// Minimum threads per block.
    pub min_threads_per_block: u32,
    /// Maximum threads per block.
    pub max_threads_per_block: u32,
    /// Required compute capability.
    pub min_sm_version: u32,
}

impl KernelMetadata {
    pub fn new(name: &str) -> Self {
        KernelMetadata {
            name: name.to_string(),
            params: Vec::new(),
            return_type: None,
            max_registers: 64,
            dynamic_shared_memory: 0,
            min_threads_per_block: 1,
            max_threads_per_block: 1024,
            min_sm_version: 50,
        }
    }

    /// Add a parameter.
    pub fn add_param(&mut self, arg: KernelArg) {
        self.params.push(arg);
    }

    /// Set the return type.
    pub fn set_return_type(&mut self, arg: KernelArg) {
        self.return_type = Some(arg);
    }

    /// Set register usage metadata.
    pub fn set_register_usage(&mut self, max_regs: u32) {
        self.max_registers = max_regs;
    }

    /// Set shared memory metadata.
    pub fn set_shared_memory(&mut self, bytes: u32) {
        self.dynamic_shared_memory = bytes;
    }

    /// Emit the full kernel metadata as PTX directives.
    pub fn emit_ptx(&self) -> String {
        let mut out = String::new();

        // Kernel entry directive
        out.push_str(&format!("\t.entry {}(", self.name));
        if !self.params.is_empty() {
            out.push('\n');
            for (i, p) in self.params.iter().enumerate() {
                out.push_str(&p.emit_ptx());
                if i < self.params.len() - 1 || self.return_type.is_some() {
                    out.push_str(",\n");
                }
            }
            if let Some(ref ret) = self.return_type {
                out.push_str(&ret.emit_ptx());
            }
            out.push('\n');
        } else if let Some(ref ret) = self.return_type {
            out.push_str(&ret.emit_ptx());
            out.push('\n');
        }
        out.push(')');

        // Constraints
        if self.max_threads_per_block > 0 {
            out.push('\n');
            out.push_str(&format!(
                "\t.maxntid\t{}, {}, {}",
                self.max_threads_per_block, 1, 1
            ));
        }
        if self.max_registers > 0 {
            out.push('\n');
            out.push_str(&format!("\t.maxnreg\t{}", self.max_registers));
        }
        if self.dynamic_shared_memory > 0 {
            out.push('\n');
            out.push_str(&format!("\t.reqntid\t{}", self.dynamic_shared_memory));
        }
        if self.min_sm_version >= 20 {
            out.push('\n');
            out.push_str(&format!("\t.minnctapersm\t{}", self.min_sm_version * 10));
        }

        out.push('\n');
        out
    }
}

/// NVVM IR metadata keys (used in LLVM IR to annotate kernels).
pub mod nvvm_metadata {
    /// Annotation key for kernel functions: "nvvm.annotations"
    pub const ANNOTATIONS: &str = "nvvm.annotations";
    /// Maximum threads per block annotation.
    pub const MAXNTID_X: &str = "maxntid_x";
    pub const MAXNTID_Y: &str = "maxntid_y";
    pub const MAXNTID_Z: &str = "maxntid_z";
    /// Required number of registers.
    pub const MAXNREG: &str = "maxnreg";
    /// Minimum threads per block.
    pub const REQNTID: &str = "reqntid";
    /// Minimum CTA per SM.
    pub const MINNCTAPERSM: &str = "minnctapersm";
    /// Kernel name for reflection.
    pub const KERNEL_NAME: &str = "kernel";
    /// Global variable annotation.
    pub const GLOBAL_VARIABLE: &str = "global";
}

/// Emit a complete NVPTX kernel with metadata as a PTX assembly string.
pub fn emit_kernel(meta: &KernelMetadata, body: &str) -> String {
    let mut out = meta.emit_ptx();
    out.push_str("{\n");
    out.push_str(body);
    out.push_str("}\n");
    out
}

/// Build a kernel metadata object from LLVM IR-level function attributes.
pub fn kernel_metadata_from_function(
    name: &str,
    param_types: &[(String, PtxType)],
    ret_type: Option<PtxType>,
    attrs: &HashMap<String, String>,
) -> KernelMetadata {
    let mut meta = KernelMetadata::new(name);

    for (pname, pty) in param_types {
        meta.add_param(KernelArg::new(pname, *pty));
    }

    if let Some(rty) = ret_type {
        meta.set_return_type(KernelArg::new("retval", rty));
    }

    if let Some(val) = attrs.get("maxnreg") {
        if let Ok(n) = val.parse() {
            meta.max_registers = n;
        }
    }
    if let Some(val) = attrs.get("maxntid") {
        if let Ok(n) = val.parse() {
            meta.max_threads_per_block = n;
        }
    }
    if let Some(val) = attrs.get("reqntid") {
        if let Ok(n) = val.parse() {
            meta.dynamic_shared_memory = n;
        }
    }

    meta
}

// ============================================================================
// PTX Version Selection
// ============================================================================

/// PTX ISA version descriptor for feature gating.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct PtxIsaVersion {
    pub major: u32,
    pub minor: u32,
}

impl PtxIsaVersion {
    pub const V4_0: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 0 };
    pub const V4_1: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 1 };
    pub const V4_2: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 2 };
    pub const V4_3: PtxIsaVersion = PtxIsaVersion { major: 4, minor: 3 };
    pub const V5_0: PtxIsaVersion = PtxIsaVersion { major: 5, minor: 0 };
    pub const V6_0: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 0 };
    pub const V6_1: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 1 };
    pub const V6_2: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 2 };
    pub const V6_3: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 3 };
    pub const V6_4: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 4 };
    pub const V6_5: PtxIsaVersion = PtxIsaVersion { major: 6, minor: 5 };
    pub const V7_0: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 0 };
    pub const V7_1: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 1 };
    pub const V7_2: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 2 };
    pub const V7_3: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 3 };
    pub const V7_4: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 4 };
    pub const V7_5: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 5 };
    pub const V7_6: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 6 };
    pub const V7_7: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 7 };
    pub const V7_8: PtxIsaVersion = PtxIsaVersion { major: 7, minor: 8 };
    pub const V8_0: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 0 };
    pub const V8_1: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 1 };
    pub const V8_2: PtxIsaVersion = PtxIsaVersion { major: 8, minor: 2 };

    /// Determine PTX ISA version from compute capability.
    pub fn from_sm(sm: u32) -> Self {
        match sm {
            10..=13 => PtxIsaVersion::V4_0,
            20..=21 => PtxIsaVersion::V4_1,
            30..=35 => PtxIsaVersion::V4_2,
            37 => PtxIsaVersion::V4_3,
            50..=53 => PtxIsaVersion::V5_0,
            60..=62 => PtxIsaVersion::V6_0,
            70..=72 => PtxIsaVersion::V7_0,
            75 => PtxIsaVersion::V7_5,
            80 => PtxIsaVersion::V7_8,
            86 => PtxIsaVersion::V8_0,
            89 => PtxIsaVersion::V8_1,
            90..=99 => PtxIsaVersion::V8_2,
            _ => PtxIsaVersion::V5_0,
        }
    }

    /// Format as a PTX .version directive string.
    pub fn to_version_string(&self) -> String {
        format!("{}.{}", self.major, self.minor)
    }

    /// Check if this version supports warp matrix operations (wmma).
    pub fn supports_wmma(&self) -> bool {
        *self >= PtxIsaVersion::V6_0
    }

    /// Check if this version supports tensor core MMA (mma).
    pub fn supports_mma(&self) -> bool {
        *self >= PtxIsaVersion::V7_0
    }

    /// Check if this version supports asynchronous copy (cp.async).
    pub fn supports_cp_async(&self) -> bool {
        *self >= PtxIsaVersion::V7_8
    }

    /// Check if this version supports stmatrix.
    pub fn supports_stmatrix(&self) -> bool {
        *self >= PtxIsaVersion::V7_8
    }

    /// Check if this version supports DPX instructions.
    pub fn supports_dpx(&self) -> bool {
        *self >= PtxIsaVersion::V8_0
    }
}

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

    #[test]
    fn test_address_space_to_ptx() {
        assert_eq!(addrspace_to_ptx_space(address_space::GLOBAL), "global");
        assert_eq!(addrspace_to_ptx_space(address_space::SHARED), "shared");
        assert_eq!(addrspace_to_ptx_space(address_space::CONSTANT), "const");
        assert_eq!(addrspace_to_ptx_space(address_space::LOCAL), "local");
        assert_eq!(addrspace_to_ptx_space(address_space::PARAM), "param");
        assert_eq!(addrspace_to_ptx_space(999), "global");
    }

    #[test]
    fn test_address_space_properties() {
        let props = get_address_space_properties(address_space::GLOBAL);
        assert!(props.is_readable);
        assert!(props.is_writable);
        assert!(props.supports_atomics);

        let shared = get_address_space_properties(address_space::SHARED);
        assert!(shared.is_readable);
        assert!(!shared.l2_cached);

        let cnst = get_address_space_properties(address_space::CONSTANT);
        assert!(cnst.is_readable);
        assert!(!cnst.is_writable);
    }

    #[test]
    fn test_target_machine_new() {
        let tm = NvptxTargetMachine::new(70, true);
        assert!(tm.is_64bit());
        assert_eq!(tm.sm_version, 70);
        assert!(tm.get_target_triple().contains("nvptx64"));
    }

    #[test]
    fn test_target_machine_32bit() {
        let tm = NvptxTargetMachine::new(50, false);
        assert!(!tm.is_64bit());
        assert!(!tm.get_target_triple().contains("64"));
    }

    #[test]
    fn test_target_machine_features() {
        let mut tm = NvptxTargetMachine::new(80, true);
        tm.enable_feature("sm_80");
        assert!(tm.has_feature("sm_80"));
        assert!(!tm.has_feature("sm_70"));
    }

    #[test]
    fn test_target_machine_opt_level() {
        let mut tm = NvptxTargetMachine::new(60, true);
        tm.set_opt_level("O3");
        assert!(!tm.add_ir_passes().is_empty());
        assert!(tm.add_ir_passes().len() > 6);
    }

    #[test]
    fn test_compute_capability_support() {
        let tm = NvptxTargetMachine::new(80, true);
        assert!(tm.supports(ComputeFeature::TensorCores));
        assert!(tm.supports(ComputeFeature::FP16));
        assert!(tm.supports(ComputeFeature::BF16));

        let tm2 = NvptxTargetMachine::new(50, true);
        assert!(!tm2.supports(ComputeFeature::TensorCores));
        assert!(!tm2.supports(ComputeFeature::BF16));
    }

    #[test]
    fn test_default_target_machine() {
        let tm = NvptxTargetMachine::default();
        assert_eq!(tm.sm_version, 70);
        assert!(tm.is_64bit());
    }

    #[test]
    fn test_kernel_metadata_empty() {
        let meta = KernelMetadata::new("my_kernel");
        let ptx = meta.emit_ptx();
        assert!(ptx.contains("my_kernel"));
        assert!(ptx.contains(".entry"));
    }

    #[test]
    fn test_kernel_metadata_with_params() {
        let mut meta = KernelMetadata::new("add_kernel");
        meta.add_param(KernelArg::new("a", PtxType::F32));
        meta.add_param(KernelArg::new("b", PtxType::F32));
        let ptx = meta.emit_ptx();
        assert!(ptx.contains(".param .f32 a"));
        assert!(ptx.contains(".param .f32 b"));
    }

    #[test]
    fn test_kernel_metadata_with_return() {
        let mut meta = KernelMetadata::new("compute");
        meta.set_return_type(KernelArg::new("retval", PtxType::U64));
        let ptx = meta.emit_ptx();
        assert!(ptx.contains(".param .u64 retval"));
    }

    #[test]
    fn test_emit_kernel_full() {
        let mut meta = KernelMetadata::new("testkern");
        meta.add_param(KernelArg::new("in", PtxType::S32));
        meta.set_register_usage(32);
        let body = "\tld.param.u32\t%r1, [in];\n\tret;\n";
        let output = emit_kernel(&meta, body);
        assert!(output.contains("testkern"));
        assert!(output.contains("ld.param.u32"));
        assert!(output.contains(".maxnreg\t32"));
    }

    #[test]
    fn test_ptx_isa_version_from_sm() {
        assert_eq!(PtxIsaVersion::from_sm(50), PtxIsaVersion::V5_0);
        assert_eq!(PtxIsaVersion::from_sm(70), PtxIsaVersion::V7_0);
        assert_eq!(PtxIsaVersion::from_sm(75), PtxIsaVersion::V7_5);
        assert_eq!(PtxIsaVersion::from_sm(80), PtxIsaVersion::V7_8);
        assert_eq!(PtxIsaVersion::from_sm(90), PtxIsaVersion::V8_2);
    }

    #[test]
    fn test_ptx_isa_version_features() {
        let v5 = PtxIsaVersion::V5_0;
        assert!(!v5.supports_wmma());
        assert!(!v5.supports_mma());

        let v7 = PtxIsaVersion::V7_0;
        assert!(v7.supports_wmma());
        assert!(v7.supports_mma());

        let v8 = PtxIsaVersion::V8_0;
        assert!(v8.supports_cp_async());
        assert!(v8.supports_dpx());
    }

    #[test]
    fn test_kernel_metadata_from_function() {
        let mut attrs = HashMap::new();
        attrs.insert("maxnreg".to_string(), "64".to_string());
        attrs.insert("maxntid".to_string(), "256".to_string());
        let params = vec![
            ("x".to_string(), PtxType::F32),
            ("y".to_string(), PtxType::F32),
        ];
        let meta = kernel_metadata_from_function("vecadd", &params, None, &attrs);
        assert_eq!(meta.name, "vecadd");
        assert_eq!(meta.params.len(), 2);
        assert_eq!(meta.max_registers, 64);
        assert_eq!(meta.max_threads_per_block, 256);
    }

    #[test]
    fn test_kernel_arg_emit_ptx() {
        let arg = KernelArg::new("my_param", PtxType::U64);
        let ptx = arg.emit_ptx();
        assert!(ptx.contains(".param .u64"));
        assert!(ptx.contains("my_param"));
    }

    #[test]
    fn test_nvvm_metadata_constants() {
        assert_eq!(nvvm_metadata::ANNOTATIONS, "nvvm.annotations");
        assert_eq!(nvvm_metadata::KERNEL_NAME, "kernel");
    }

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

    #[test]
    fn test_emit_module_basic() {
        let tm = NvptxTargetMachine::new(70, true);
        // Create a minimal machine function
        let mut mf = MachineFunction::new("test_func");
        let block = MachineBlock {
            name: "entry".to_string(),
            instructions: Vec::new(),
        };
        mf.blocks.push(block);
        let output = tm.emit_module(&mf);
        assert!(output.contains(".version"));
        assert!(output.contains(".target"));
        assert!(output.contains("test_func"));
    }
}