melior/
utility.rs

1//! Utility functions.
2
3use crate::{
4    context::Context, dialect::DialectRegistry, ir::Module, logical_result::LogicalResult, pass,
5    string_ref::StringRef, Error,
6};
7use mlir_sys::{
8    mlirLoadIRDLDialects, mlirParsePassPipeline, mlirRegisterAllDialects,
9    mlirRegisterAllLLVMTranslations, mlirRegisterAllPasses, MlirStringRef,
10};
11use std::{
12    ffi::c_void,
13    fmt::{self, Formatter},
14    sync::Once,
15};
16
17/// Registers all dialects to a dialect registry.
18pub fn register_all_dialects(registry: &DialectRegistry) {
19    unsafe { mlirRegisterAllDialects(registry.to_raw()) }
20}
21
22/// Register all translations from other dialects to the `llvm` dialect.
23pub fn register_all_llvm_translations(context: &Context) {
24    unsafe { mlirRegisterAllLLVMTranslations(context.to_raw()) }
25}
26
27/// Register all passes.
28pub fn register_all_passes() {
29    static ONCE: Once = Once::new();
30
31    // Multiple calls of `mlirRegisterAllPasses` seems to cause double free.
32    ONCE.call_once(|| unsafe { mlirRegisterAllPasses() });
33}
34
35/// Parses a pass pipeline.
36pub fn parse_pass_pipeline(manager: pass::OperationPassManager, source: &str) -> Result<(), Error> {
37    let mut error_message = None;
38
39    let result = LogicalResult::from_raw(unsafe {
40        mlirParsePassPipeline(
41            manager.to_raw(),
42            StringRef::new(source).to_raw(),
43            Some(handle_parse_error),
44            &mut error_message as *mut _ as *mut _,
45        )
46    });
47
48    if result.is_success() {
49        Ok(())
50    } else {
51        Err(Error::ParsePassPipeline(error_message.unwrap_or_else(
52            || "failed to parse error message in UTF-8".into(),
53        )))
54    }
55}
56
57/// Loads all IRDL dialects in the provided module, registering the dialects in the module's associated context.
58pub fn load_irdl_dialects(module: &Module) -> bool {
59    unsafe { mlirLoadIRDLDialects(module.to_raw()).value == 1 }
60}
61
62unsafe extern "C" fn handle_parse_error(raw_string: MlirStringRef, data: *mut c_void) {
63    let string = StringRef::from_raw(raw_string);
64    let data = &mut *(data as *mut Option<String>);
65
66    if let Some(message) = data {
67        message.extend(string.as_str())
68    } else {
69        *data = string.as_str().map(String::from).ok();
70    }
71}
72
73pub(crate) unsafe extern "C" fn print_callback(string: MlirStringRef, data: *mut c_void) {
74    let (formatter, result) = &mut *(data as *mut (&mut Formatter, fmt::Result));
75
76    if result.is_err() {
77        return;
78    }
79
80    *result = (|| {
81        write!(
82            formatter,
83            "{}",
84            StringRef::from_raw(string)
85                .as_str()
86                .map_err(|_| fmt::Error)?
87        )
88    })();
89}
90
91pub(crate) unsafe extern "C" fn print_string_callback(string: MlirStringRef, data: *mut c_void) {
92    let (writer, result) = &mut *(data as *mut (String, Result<(), Error>));
93
94    if result.is_err() {
95        return;
96    }
97
98    *result = (|| {
99        writer.push_str(StringRef::from_raw(string).as_str()?);
100
101        Ok(())
102    })();
103}
104
105#[cfg(test)]
106mod tests {
107    use crate::ir::Location;
108
109    use super::*;
110
111    #[test]
112    fn register_dialects() {
113        let registry = DialectRegistry::new();
114
115        register_all_dialects(&registry);
116    }
117
118    #[test]
119    fn register_dialects_twice() {
120        let registry = DialectRegistry::new();
121
122        register_all_dialects(&registry);
123        register_all_dialects(&registry);
124    }
125
126    #[test]
127    fn register_llvm_translations() {
128        let context = Context::new();
129
130        register_all_llvm_translations(&context);
131    }
132
133    #[test]
134    fn register_llvm_translations_twice() {
135        let context = Context::new();
136
137        register_all_llvm_translations(&context);
138        register_all_llvm_translations(&context);
139    }
140
141    #[test]
142    fn register_passes() {
143        register_all_passes();
144    }
145
146    #[test]
147    fn register_passes_twice() {
148        register_all_passes();
149        register_all_passes();
150    }
151
152    #[test]
153    fn register_passes_many_times() {
154        for _ in 0..1000 {
155            register_all_passes();
156        }
157    }
158
159    #[test]
160    fn test_load_irdl_dialects() {
161        let context = Context::new();
162        let module = Module::new(Location::unknown(&context));
163
164        assert!(load_irdl_dialects(&module));
165    }
166}
melior/utility.rs

melior/
utility.rs
