llvm-native-core 0.1.5

//! AMDGPU Code Object Format — ELF code object emission for AMDGPU.
//!
//! Generates ELF code objects containing:
//! - .text section with kernel machine code
//! - .note section with AMDGPU-specific note records
//! - Kernel descriptors
//! - ISA version metadata
//!
//! The AMDGPU code object format uses ELF with special note records
//! that describe kernel properties (register usage, LDS size, etc.).
//!
//! Clean-room reconstruction from AMD ROCm code object specification.
//! Zero LLVM source code consultation.

use std::collections::HashMap;

use super::amdgpu_kernel_lowering::AmdgpuKernelDescriptor;
use super::amdgpu_target_machine::AmdgpuIsaVersion;

// ═══════════════════════════════════════════════════════════════════════════
// ELF Note Types for AMDGPU
// ═══════════════════════════════════════════════════════════════════════════

/// AMDGPU ELF note type for kernel code properties.
pub const NT_AMDGPU_CODE_PROPERTY: u32 = 0x20000001;

/// AMDGPU ELF note type for ISA version.
pub const NT_AMDGPU_ISA_VERSION: u32 = 0x20000002;

/// AMDGPU ELF note type for metadata.
pub const NT_AMDGPU_METADATA: u32 = 0x20000003;

// ═══════════════════════════════════════════════════════════════════════════
// Code Object
// ═══════════════════════════════════════════════════════════════════════════

/// An AMDGPU code object containing compiled kernels.
#[derive(Debug)]
pub struct AmdgpuCodeObject {
    /// ISA version for this code object.
    pub isa_version: AmdgpuIsaVersion,
    /// Kernel code bytes.
    pub text_section: Vec<u8>,
    /// Kernel descriptors keyed by kernel name.
    pub kernel_descriptors: HashMap<String, AmdgpuKernelDescriptor>,
    /// Note records.
    pub notes: Vec<AmdgpuNoteRecord>,
    /// ELF header and section data.
    elf_data: Vec<u8>,
}

/// An ELF note record entry.
#[derive(Debug, Clone)]
pub struct AmdgpuNoteRecord {
    pub note_type: u32,
    pub name: String,
    pub desc: Vec<u8>,
}

impl AmdgpuCodeObject {
    pub fn new(isa: AmdgpuIsaVersion) -> Self {
        Self {
            isa_version: isa,
            text_section: Vec::new(),
            kernel_descriptors: HashMap::new(),
            notes: Vec::new(),
            elf_data: Vec::new(),
        }
    }

    /// Add a kernel to the code object.
    pub fn add_kernel(&mut self, name: &str, code: &[u8], descriptor: AmdgpuKernelDescriptor) {
        self.kernel_descriptors.insert(name.to_string(), descriptor);
        self.text_section.extend_from_slice(code);
    }

    /// Add an ISA version note.
    pub fn add_isa_note(&mut self) {
        let mut isa_data = Vec::new();
        // ISA version: major.minor.stepping encoded as 3 u16 values
        let (major, minor, stepping) = self.isa_components();
        isa_data.extend_from_slice(&major.to_le_bytes());
        isa_data.extend_from_slice(&minor.to_le_bytes());
        isa_data.extend_from_slice(&stepping.to_le_bytes());

        self.notes.push(AmdgpuNoteRecord {
            note_type: NT_AMDGPU_ISA_VERSION,
            name: "AMD".to_string(),
            desc: isa_data,
        });
    }

    /// Add kernel code property notes.
    pub fn add_kernel_notes(&mut self) {
        for (name, desc) in &self.kernel_descriptors {
            let prop_bytes = desc.to_bytes();
            self.notes.push(AmdgpuNoteRecord {
                note_type: NT_AMDGPU_CODE_PROPERTY,
                name: name.clone(),
                desc: prop_bytes,
            });
        }
    }

    /// Get ISA version components.
    fn isa_components(&self) -> (u16, u16, u16) {
        match self.isa_version {
            AmdgpuIsaVersion::GFX600 => (6, 0, 0),
            AmdgpuIsaVersion::GFX700 => (7, 0, 0),
            AmdgpuIsaVersion::GFX800 => (8, 0, 2),
            AmdgpuIsaVersion::GFX803 => (8, 0, 3),
            AmdgpuIsaVersion::GFX900 => (9, 0, 0),
            AmdgpuIsaVersion::GFX906 => (9, 0, 6),
            AmdgpuIsaVersion::GFX908 => (9, 0, 8),
            AmdgpuIsaVersion::GFX90A => (9, 0, 10),
            AmdgpuIsaVersion::GFX1010 => (10, 1, 0),
            AmdgpuIsaVersion::GFX1030 => (10, 3, 0),
            AmdgpuIsaVersion::GFX1100 => (11, 0, 0),
            AmdgpuIsaVersion::GFX1150 => (11, 5, 0),
        }
    }

    /// Generate a minimal ELF code object.
    /// Returns the complete ELF file bytes.
    pub fn emit_elf(&mut self) -> Vec<u8> {
        self.elf_data.clear();
        self.add_isa_note();
        self.add_kernel_notes();

        // Build minimal ELF
        self.build_elf_header();
        self.build_section_headers();
        self.build_note_section();

        self.elf_data.clone()
    }

    /// Build a minimal 64-bit ELF header.
    fn build_elf_header(&mut self) {
        let mut header = Vec::with_capacity(64);

        // e_ident
        header.push(0x7F); // ELFMAG0
        header.push(b'E'); // ELFMAG1
        header.push(b'L'); // ELFMAG2
        header.push(b'F'); // ELFMAG3
        header.push(2); // EI_CLASS = ELFCLASS64
        header.push(1); // EI_DATA = ELFDATA2LSB
        header.push(1); // EI_VERSION
        header.push(0); // EI_OSABI = ELFOSABI_NONE
        header.extend_from_slice(&[0u8; 8]); // EI_PAD

        // e_type = ET_REL (1)
        header.extend_from_slice(&1u16.to_le_bytes());
        // e_machine = EM_AMDGPU (0xE0)
        header.extend_from_slice(&0xE0u16.to_le_bytes());
        // e_version
        header.extend_from_slice(&1u32.to_le_bytes());

        // e_entry, e_phoff = 0 (no program headers for relocatable)
        header.extend_from_slice(&[0u8; 16]);

        // e_shoff — section header offset (will be after all sections)
        let text_offset = 64u64;
        header.extend_from_slice(&text_offset.to_le_bytes());

        // e_flags — AMDGPU specific
        let e_flags = self.elf_flags();
        header.extend_from_slice(&e_flags.to_le_bytes());

        // e_ehsize
        header.extend_from_slice(&64u16.to_le_bytes());
        // e_phentsize = 0
        header.extend_from_slice(&0u16.to_le_bytes());
        // e_phnum = 0
        header.extend_from_slice(&0u16.to_le_bytes());
        // e_shentsize = 64
        header.extend_from_slice(&64u16.to_le_bytes());
        // e_shnum = 3 (.text, .note, .shstrtab)
        header.extend_from_slice(&3u16.to_le_bytes());
        // e_shstrndx = 2
        header.extend_from_slice(&2u16.to_le_bytes());

        self.elf_data = header;
    }

    /// Build section headers.
    fn build_section_headers(&mut self) {
        // Section headers come after the sections
        // For now, emit placeholder section headers
    }

    /// Build the .note section.
    fn build_note_section(&mut self) {
        for note in &self.notes {
            // Note header: name_size, desc_size, type
            let name_bytes = note.name.as_bytes();
            let name_size = name_bytes.len() as u32;
            let name_padded = (name_size + 3) & !3; // 4-byte align

            let desc_size = note.desc.len() as u32;
            let desc_padded = (desc_size + 3) & !3;

            self.elf_data.extend_from_slice(&name_size.to_le_bytes());
            self.elf_data.extend_from_slice(&desc_size.to_le_bytes());
            self.elf_data
                .extend_from_slice(&note.note_type.to_le_bytes());
            self.elf_data.extend_from_slice(name_bytes);
            if name_padded > name_size {
                self.elf_data
                    .extend_from_slice(&vec![0u8; (name_padded - name_size) as usize]);
            }
            self.elf_data.extend_from_slice(&note.desc);
            if desc_padded > desc_size {
                self.elf_data
                    .extend_from_slice(&vec![0u8; (desc_padded - desc_size) as usize]);
            }
        }
    }

    /// Get the ELF flags for this ISA version.
    fn elf_flags(&self) -> u32 {
        // AMDGPU ELF flags encode the ISA version
        match self.isa_version {
            AmdgpuIsaVersion::GFX600 => 0x0600,
            AmdgpuIsaVersion::GFX700 => 0x0700,
            AmdgpuIsaVersion::GFX800 => 0x0800,
            AmdgpuIsaVersion::GFX803 => 0x0803,
            AmdgpuIsaVersion::GFX900 => 0x0900,
            AmdgpuIsaVersion::GFX906 => 0x0906,
            AmdgpuIsaVersion::GFX908 => 0x0908,
            AmdgpuIsaVersion::GFX90A => 0x090A,
            AmdgpuIsaVersion::GFX1010 => 0x0A10,
            AmdgpuIsaVersion::GFX1030 => 0x0A30,
            AmdgpuIsaVersion::GFX1100 => 0x0B00,
            AmdgpuIsaVersion::GFX1150 => 0x0B50,
        }
    }
}

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

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

    #[test]
    fn test_code_object_create() {
        let co = AmdgpuCodeObject::new(AmdgpuIsaVersion::GFX900);
        assert!(co.text_section.is_empty());
        assert!(co.kernel_descriptors.is_empty());
    }

    #[test]
    fn test_code_object_add_kernel() {
        let mut co = AmdgpuCodeObject::new(AmdgpuIsaVersion::GFX900);
        let code = vec![0xBF, 0x80, 0x00, 0x00]; // s_nop
        let desc = AmdgpuKernelDescriptor::new();
        co.add_kernel("test_kernel", &code, desc);
        assert_eq!(co.kernel_descriptors.len(), 1);
        assert_eq!(co.text_section.len(), 4);
    }

    #[test]
    fn test_code_object_emit_elf() {
        let mut co = AmdgpuCodeObject::new(AmdgpuIsaVersion::GFX900);
        let desc = AmdgpuKernelDescriptor::new();
        co.add_kernel("kernel0", &[0xBF, 0x80, 0x00, 0x00], desc);
        let elf = co.emit_elf();
        assert!(elf.len() > 64); // at least ELF header size
        assert_eq!(elf[0], 0x7F); // ELF magic
        assert_eq!(elf[1], b'E');
        assert_eq!(elf[2], b'L');
        assert_eq!(elf[3], b'F');
    }

    #[test]
    fn test_code_object_elf_flags() {
        let co = AmdgpuCodeObject::new(AmdgpuIsaVersion::GFX1030);
        assert_eq!(co.elf_flags(), 0x0A30);
    }

    #[test]
    fn test_note_record() {
        let note = AmdgpuNoteRecord {
            note_type: NT_AMDGPU_ISA_VERSION,
            name: "AMD".into(),
            desc: vec![9, 0, 6, 0, 0, 0], // GFX906
        };
        assert_eq!(note.note_type, NT_AMDGPU_ISA_VERSION);
    }
}