llvm-native-core 0.1.12

//! AMDGPU Assembly Printer — emits GCN assembly syntax for all
//! instruction formats.
//!
//! Clean-room reconstruction from AMD GCN ISA assembly syntax docs.
//! Zero LLVM source code consultation.

use super::amdgpu_instr_info::{
    AmdgpuEncodingFormat, AmdgpuInstrDesc, AmdgpuInstrInfo, AmdgpuOpcode,
};
use super::amdgpu_register_info::AmdgpuRegisterInfo;

/// Prints AMDGPU instructions in GCN assembly syntax.
pub struct AmdgpuAsmPrinter {
    pub instr_info: AmdgpuInstrInfo,
    pub reg_info: AmdgpuRegisterInfo,
    output: String,
}

impl AmdgpuAsmPrinter {
    pub fn new() -> Self {
        Self {
            instr_info: AmdgpuInstrInfo::new(),
            reg_info: AmdgpuRegisterInfo::new(),
            output: String::new(),
        }
    }

    /// Print a single instruction with its operands.
    pub fn print_instruction(&mut self, opcode: AmdgpuOpcode, operands: &[String]) -> String {
        self.output.clear();

        let Some(desc) = self.instr_info.get(opcode) else {
            return format!("s_nop 0  ; unknown opcode");
        };

        // Print mnemonic
        self.output.push_str(desc.mnemonic);

        // Print operands based on format
        match desc.encoding_format {
            AmdgpuEncodingFormat::SOP1 => {
                if operands.len() >= 2 {
                    self.output
                        .push_str(&format!(" {}, {}", operands[0], operands[1]));
                }
            }
            AmdgpuEncodingFormat::SOP2 => {
                if operands.len() >= 3 {
                    self.output.push_str(&format!(
                        " {}, {}, {}",
                        operands[0], operands[1], operands[2]
                    ));
                }
            }
            AmdgpuEncodingFormat::SOPC => {
                if operands.len() >= 2 {
                    self.output
                        .push_str(&format!(" {}, {}", operands[0], operands[1]));
                }
            }
            AmdgpuEncodingFormat::SOPK => {
                if operands.len() >= 2 {
                    self.output
                        .push_str(&format!(" {}, {}", operands[0], operands[1]));
                }
            }
            AmdgpuEncodingFormat::SOPP => {
                if operands.len() >= 1 {
                    self.output.push_str(&format!(" {}", operands[0]));
                }
            }
            AmdgpuEncodingFormat::VOP1 => {
                if operands.len() >= 2 {
                    self.output
                        .push_str(&format!(" {}, {}", operands[0], operands[1]));
                }
            }
            AmdgpuEncodingFormat::VOP2 | AmdgpuEncodingFormat::VOPC => {
                if operands.len() >= 3 {
                    self.output.push_str(&format!(
                        " {}, {}, {}",
                        operands[0], operands[1], operands[2]
                    ));
                }
            }
            AmdgpuEncodingFormat::VOP3 | AmdgpuEncodingFormat::VOP3P => {
                if operands.len() >= 4 {
                    self.output.push_str(&format!(
                        " {}, {}, {}, {}",
                        operands[0], operands[1], operands[2], operands[3]
                    ));
                }
            }
            AmdgpuEncodingFormat::SMEM => {
                self.output.push_str(&format!(
                    " {}, s[{}], {}",
                    operands.get(0).unwrap_or(&"s0".to_string()),
                    operands.get(1).unwrap_or(&"0".to_string()),
                    operands.get(2).unwrap_or(&"0".to_string()),
                ));
            }
            AmdgpuEncodingFormat::FLAT => {
                self.output.push_str(&format!(
                    " {}, v[{}]",
                    operands.get(0).unwrap_or(&"v0".to_string()),
                    operands.get(1).unwrap_or(&"0".to_string()),
                ));
            }
            AmdgpuEncodingFormat::DS => {
                self.output.push_str(&format!(
                    " {}, {}, {}",
                    operands.get(0).unwrap_or(&"v0".to_string()),
                    operands.get(1).unwrap_or(&"v1".to_string()),
                    operands.get(2).unwrap_or(&"0".to_string()),
                ));
            }
            AmdgpuEncodingFormat::MIMG => {
                self.output.push_str(&format!(
                    " {}, v[{}], s[{}]",
                    operands.get(0).unwrap_or(&"v0".to_string()),
                    operands.get(1).unwrap_or(&"0".to_string()),
                    operands.get(2).unwrap_or(&"0:3".to_string()),
                ));
            }
            AmdgpuEncodingFormat::EXP => {
                self.output.push_str(&format!(
                    " mrt0 {}, {}, {}, {}",
                    operands.get(0).unwrap_or(&"v0".to_string()),
                    operands.get(1).unwrap_or(&"v1".to_string()),
                    operands.get(2).unwrap_or(&"v2".to_string()),
                    operands.get(3).unwrap_or(&"v3".to_string()),
                ));
            }
            _ => {}
        }

        // Add comment
        self.output.push_str(&format!(
            "  ; {} ({} bits)",
            desc.mnemonic,
            desc.encoding_format.size_bits()
        ));

        self.output.clone()
    }

    /// Print a kernel prologue.
    pub fn print_kernel_prologue(&mut self, kernel_name: &str) -> String {
        format!(
            r#"{name}:
  ; Kernel: {name}
  ; Wavefront size: 64
  ; Implicit args in SGPRs:
  ;   s0-s2: workgroup_id (x, y, z)
  ;   s4-s5: workgroup_size
  ;   s8-s9: kernarg_segment_ptr
  s_load_dwordx2 s[4:5], s[8:9], 0x0
  s_waitcnt lgkmcnt(0)
"#,
            name = kernel_name
        )
    }

    /// Print a kernel epilogue.
    pub fn print_kernel_epilogue(&self) -> &'static str {
        "  s_endpgm\n"
    }

    /// Print a comment line.
    pub fn print_comment(&mut self, comment: &str) -> String {
        format!("  ; {}", comment)
    }

    /// Print a label.
    pub fn print_label(&mut self, label: &str) -> String {
        format!("{}:", label)
    }
}

impl Default for AmdgpuAsmPrinter {
    fn default() -> Self {
        Self::new()
    }
}

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

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

    #[test]
    fn test_print_s_nop() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(AmdgpuOpcode::SNop, &["0".into()]);
        assert!(asm.contains("s_nop"));
    }

    #[test]
    fn test_print_v_add_f32() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(
            AmdgpuOpcode::VAddF32,
            &["v0".into(), "v1".into(), "v2".into()],
        );
        assert!(asm.contains("v_add_f32"));
        assert!(asm.contains("v0"));
        assert!(asm.contains("v1"));
    }

    #[test]
    fn test_print_s_load_dword() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(
            AmdgpuOpcode::SLoadDword,
            &["s5".into(), "2".into(), "0".into()],
        );
        assert!(asm.contains("s_load_dword"));
    }

    #[test]
    fn test_print_global_load() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(
            AmdgpuOpcode::GlobalLoadDword,
            &["v5".into(), "v[2]".into(), "0".into()],
        );
        assert!(asm.contains("global_load_dword"));
    }

    #[test]
    fn test_print_ds_write() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(
            AmdgpuOpcode::DsWriteB32,
            &["v0".into(), "42".into(), "0".into()],
        );
        assert!(asm.contains("ds_write_b32"));
    }

    #[test]
    fn test_kernel_prologue() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_kernel_prologue("my_kernel");
        assert!(asm.contains("my_kernel:"));
        assert!(asm.contains("s_load_dwordx2"));
    }

    #[test]
    fn test_print_image_sample() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_instruction(
            AmdgpuOpcode::ImageSample,
            &["v[0:3]".into(), "v[4:7]".into(), "s[0:3]".into()],
        );
        assert!(asm.contains("image_sample"));
    }

    #[test]
    fn test_comment() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_comment("this is a comment");
        assert!(asm.starts_with("  ;"));
        assert!(asm.contains("this is a comment"));
    }

    #[test]
    fn test_label() {
        let mut printer = AmdgpuAsmPrinter::new();
        let asm = printer.print_label("loop_start");
        assert_eq!(asm, "loop_start:");
    }
}