From b6a03d20a37b49e5262ee38702399b456e5e94c6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=ABlle=20van=20Essen?= <joelle@v-essen.nl>
Date: Wed, 5 Feb 2025 16:24:44 +0100
Subject: [PATCH] LLVM unit tests: first attempt

---
 ptx/src/pass/emit_llvm.rs     | 20 ++++++++++++++++++++
 ptx/src/test/spirv_run/mod.rs | 24 ++++++++++++++++++++++++
 2 files changed, 44 insertions(+)

diff --git a/ptx/src/pass/emit_llvm.rs b/ptx/src/pass/emit_llvm.rs
index 2d1269d..af5d905 100644
--- a/ptx/src/pass/emit_llvm.rs
+++ b/ptx/src/pass/emit_llvm.rs
@@ -24,6 +24,7 @@
 // shows it fails inside amdgpu-isel. You can get a little bit furthr with "-mllvm -global-isel",
 // but it will too fail similarly, but with "unable to legalize instruction"
 
+use std::alloc::{alloc, dealloc, handle_alloc_error, Layout};
 use std::array::TryFromSliceError;
 use std::convert::TryInto;
 use std::ffi::{CStr, NulError};
@@ -32,6 +33,7 @@ use std::{i8, ptr};
 
 use super::*;
 use llvm_zluda::analysis::{LLVMVerifierFailureAction, LLVMVerifyModule};
+use llvm_zluda::bit_reader::LLVMParseBitcode2;
 use llvm_zluda::bit_writer::LLVMWriteBitcodeToMemoryBuffer;
 use llvm_zluda::{core::*, *};
 use llvm_zluda::{prelude::*, LLVMZludaBuildAtomicRMW};
@@ -148,6 +150,24 @@ impl std::fmt::Debug for Message {
 
 pub struct MemoryBuffer(LLVMMemoryBufferRef);
 
+impl MemoryBuffer {
+    pub fn print_as_asm(&self) -> &str {
+        unsafe {
+            let layout = Layout::new::<LLVMModuleRef>();
+            let p_module = alloc(layout);
+            if p_module.is_null() {
+                handle_alloc_error(layout);
+            }
+            LLVMParseBitcode2(self.0, p_module as *mut LLVMModuleRef);
+            let asm = LLVMPrintModuleToString(*(p_module as *mut LLVMModuleRef));
+            LLVMDisposeModule(*(p_module as *mut LLVMModuleRef));
+            dealloc(p_module, layout);
+            let asm = CStr::from_ptr(asm);
+            asm.to_str().unwrap().trim()
+        }
+    }
+}
+
 impl Drop for MemoryBuffer {
     fn drop(&mut self) {
         unsafe {
diff --git a/ptx/src/test/spirv_run/mod.rs b/ptx/src/test/spirv_run/mod.rs
index e4171cd..aca0909 100644
--- a/ptx/src/test/spirv_run/mod.rs
+++ b/ptx/src/test/spirv_run/mod.rs
@@ -28,6 +28,17 @@ macro_rules! test_ptx {
                 test_cuda_assert(stringify!($fn_name), ptx, &input, &mut output)
             }
         }
+
+        paste::item! {
+            #[test]
+            fn [<$fn_name _llvm>]() -> Result<(), Box<dyn std::error::Error>> {
+                let fn_name = stringify!($fn_name);
+                println!("{}", fn_name);
+                let ptx = include_str!(concat!(stringify!($fn_name), ".ptx"));
+                let ll = include_str!(concat!("../ll/", stringify!($fn_name), ".ll")).trim();
+                test_llvm_assert(ptx, &ll)
+            }
+        }
     };
 
     ($fn_name:ident) => {};
@@ -223,6 +234,19 @@ fn test_hip_assert<
     Ok(())
 }
 
+fn test_llvm_assert<
+    'a,
+>(
+    ptx_text: &'a str,
+    expected_ll: &str
+) -> Result<(), Box<dyn error::Error + 'a>> {
+    let ast = ptx_parser::parse_module_checked(ptx_text).unwrap();
+    let llvm_ir = pass::to_llvm_module(ast).unwrap();
+    let actual_ll = llvm_ir.llvm_ir.print_as_asm();
+    assert_eq!(actual_ll, expected_ll);
+    Ok(())
+}
+
 fn test_cuda_assert<
     'a,
     Input: From<u8> + Debug + Copy + PartialEq,