nvptx 0.2.3

Rust to PTX compiler using LLVM/NVPTX target
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use llvm_sys::bit_reader::*;
use llvm_sys::core::*;
use llvm_sys::prelude::*;

use failure::err_msg;
use std::ffi::*;
use std::os::raw::c_char;
use std::path::*;
use std::ptr::null_mut;

use error::*;

struct MemoryBuffer(LLVMMemoryBufferRef);

impl Drop for MemoryBuffer {
    fn drop(&mut self) {
        unsafe { LLVMDisposeMemoryBuffer(self.0) }
    }
}

impl MemoryBuffer {
    fn new(filename: &str) -> ResultAny<Self> {
        let input = CString::new(filename)?;
        let mut membuf: LLVMMemoryBufferRef = null_mut();
        let mut msg: *mut c_char = null_mut();
        let result = unsafe {
            LLVMCreateMemoryBufferWithContentsOfFile(
                input.into_raw(),
                &mut membuf as *mut LLVMMemoryBufferRef,
                &mut msg as *mut *mut c_char,
            )
        };
        if result != 0 {
            let msg = unsafe { CString::from_raw(msg) };
            return Err(err_msg(format!("Canont read input: {:?}", msg)));
        }
        Ok(MemoryBuffer(membuf))
    }
}

#[derive(Debug)]
struct Module(LLVMModuleRef);

#[derive(Debug)]
struct Function(LLVMValueRef);

impl Module {
    fn parse_bitcode(buf: &MemoryBuffer) -> ResultAny<Self> {
        let mut md: LLVMModuleRef = null_mut();
        let res = unsafe { LLVMParseBitcode2(buf.0, &mut md as *mut _) };
        if res != 0 {
            return Err(err_msg("Cannot parse LLVM Bitcode"));
        }
        Ok(Module(md))
    }

    fn read_bitcode(filename: &str) -> ResultAny<Self> {
        let membuf = MemoryBuffer::new(filename)?;
        Self::parse_bitcode(&membuf)
    }

    fn functions(&self) -> Vec<Function> {
        let mut funcs = Vec::new();
        let mut f = unsafe { LLVMGetFirstFunction(self.0) };
        while f != null_mut() {
            funcs.push(Function(f));
            f = unsafe { LLVMGetNextFunction(f) };
        }
        funcs
    }
}

impl Function {
    fn name(&self) -> String {
        let name = unsafe { CString::from_raw(LLVMGetValueName(self.0) as *mut _) };
        name.into_string().expect("Fail to parse function name")
    }

    // See the LLVM call convention list
    //
    // - PTX_Kernel = 71
    // - PTX_Device = 72
    //
    // http://llvm.org/doxygen/CallingConv_8h_source.html
    fn call_conv(&self) -> u32 {
        unsafe { LLVMGetFunctionCallConv(self.0) }
    }

    fn is_ptx_kernel(&self) -> bool {
        self.call_conv() == 71
    }

    fn is_ptx_device_func(&self) -> bool {
        self.call_conv() == 72
    }
}

pub fn get_ptx_functions<P: AsRef<Path>>(filename: P) -> ResultAny<Vec<String>> {
    let path = filename.as_ref().to_str().unwrap();
    let md = Module::read_bitcode(path)?;
    let ptx: Vec<_> = md
        .functions()
        .iter()
        .filter(|f| f.is_ptx_kernel() || f.is_ptx_device_func())
        .map(|f| f.name())
        .collect();
    if ptx.len() == 0 {
        return Err(err_msg("No PTX found"));
    }
    Ok(ptx)
}