melior/
execution_engine.rs

1use crate::{Error, ir::Module, logical_result::LogicalResult, string_ref::StringRef};
2use mlir_sys::{
3    MlirExecutionEngine, mlirExecutionEngineCreate, mlirExecutionEngineDestroy,
4    mlirExecutionEngineDumpToObjectFile, mlirExecutionEngineInvokePacked,
5    mlirExecutionEngineLookup, mlirExecutionEngineRegisterSymbol,
6};
7
8/// An execution engine.
9pub struct ExecutionEngine {
10    raw: MlirExecutionEngine,
11}
12
13impl ExecutionEngine {
14    /// Creates an execution engine.
15    pub fn new(
16        module: &Module,
17        optimization_level: usize,
18        shared_library_paths: &[&str],
19        enable_object_dump: bool,
20        enable_pic: bool,
21    ) -> Self {
22        Self {
23            raw: unsafe {
24                mlirExecutionEngineCreate(
25                    module.to_raw(),
26                    optimization_level as i32,
27                    shared_library_paths.len() as i32,
28                    shared_library_paths
29                        .iter()
30                        .map(|&string| StringRef::new(string).to_raw())
31                        .collect::<Vec<_>>()
32                        .as_ptr(),
33                    enable_object_dump,
34                    enable_pic,
35                )
36            },
37        }
38    }
39
40    /// Searches a symbol in a module and returns a pointer to it.
41    pub fn lookup(&self, name: &str) -> *mut () {
42        unsafe { mlirExecutionEngineLookup(self.raw, StringRef::new(name).to_raw()) as *mut () }
43    }
44
45    /// Invokes a function in a module. The `arguments` argument includes
46    /// pointers to results of the function as well as arguments.
47    ///
48    /// # Safety
49    ///
50    /// This function modifies memory locations pointed by the `arguments`
51    /// argument. If those pointers are invalid or misaligned, calling this
52    /// function might result in undefined behavior.
53    pub unsafe fn invoke_packed(&self, name: &str, arguments: &mut [*mut ()]) -> Result<(), Error> {
54        let result = LogicalResult::from_raw(unsafe {
55            mlirExecutionEngineInvokePacked(
56                self.raw,
57                StringRef::new(name).to_raw(),
58                arguments.as_mut_ptr() as _,
59            )
60        });
61
62        if result.is_success() {
63            Ok(())
64        } else {
65            Err(Error::InvokeFunction)
66        }
67    }
68
69    /// Register a symbol. This symbol will be accessible to the JIT'd codes.
70    ///
71    /// # Safety
72    ///
73    /// This function makes a pointer accessible to the execution engine. If a
74    /// given pointer is invalid or misaligned, calling this function might
75    /// result in undefined behavior.
76    pub unsafe fn register_symbol(&self, name: &str, ptr: *mut ()) {
77        unsafe {
78            mlirExecutionEngineRegisterSymbol(self.raw, StringRef::new(name).to_raw(), ptr as _);
79        }
80    }
81
82    /// Dumps a module to an object file.
83    pub fn dump_to_object_file(&self, path: &str) {
84        unsafe { mlirExecutionEngineDumpToObjectFile(self.raw, StringRef::new(path).to_raw()) }
85    }
86}
87
88impl Drop for ExecutionEngine {
89    fn drop(&mut self) {
90        unsafe { mlirExecutionEngineDestroy(self.raw) }
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use super::*;
97    use crate::{pass, test::create_test_context};
98
99    #[test]
100    fn invoke_packed() {
101        let context = create_test_context();
102
103        let mut module = Module::parse(
104            &context,
105            r#"
106            module {
107                func.func @add(%arg0 : i32) -> i32 attributes { llvm.emit_c_interface } {
108                    %res = arith.addi %arg0, %arg0 : i32
109                    return %res : i32
110                }
111            }
112            "#,
113        )
114        .unwrap();
115
116        let pass_manager = pass::PassManager::new(&context);
117        pass_manager.add_pass(pass::conversion::create_to_llvm());
118
119        assert_eq!(pass_manager.run(&mut module), Ok(()));
120
121        let engine = ExecutionEngine::new(&module, 2, &[], false, false);
122
123        let mut argument = 42;
124        let mut result = -1;
125
126        assert_eq!(
127            unsafe {
128                engine.invoke_packed(
129                    "add",
130                    &mut [
131                        &mut argument as *mut i32 as *mut (),
132                        &mut result as *mut i32 as *mut (),
133                    ],
134                )
135            },
136            Ok(())
137        );
138
139        assert_eq!(argument, 42);
140        assert_eq!(result, 84);
141    }
142
143    #[test]
144    fn dump_to_object_file() {
145        let context = create_test_context();
146
147        let mut module = Module::parse(
148            &context,
149            r#"
150            module {
151                func.func @add(%arg0 : i32) -> i32 {
152                    %res = arith.addi %arg0, %arg0 : i32
153                    return %res : i32
154                }
155            }
156            "#,
157        )
158        .unwrap();
159
160        let pass_manager = pass::PassManager::new(&context);
161        pass_manager.add_pass(pass::conversion::create_to_llvm());
162
163        assert_eq!(pass_manager.run(&mut module), Ok(()));
164
165        // TODO: use `tempfile` crate
166        ExecutionEngine::new(&module, 2, &[], true, false)
167            .dump_to_object_file("/tmp/melior/test.o");
168    }
169}
melior/execution_engine.rs

melior/
execution_engine.rs
