melior/
execution_engine.rs

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

melior/
execution_engine.rs
