melior/pass/
external.rs

1//! External passes
2
3use super::Pass;
4use crate::{
5    dialect::DialectHandle,
6    ir::{r#type::TypeId, OperationRef},
7    ContextRef, StringRef,
8};
9use mlir_sys::{
10    mlirCreateExternalPass, mlirExternalPassSignalFailure, MlirContext, MlirExternalPass,
11    MlirExternalPassCallbacks, MlirLogicalResult, MlirOperation,
12};
13use std::{ffi::c_void, marker::PhantomData, mem::transmute, ptr::drop_in_place};
14
15#[derive(Clone, Copy, Debug)]
16pub struct ExternalPass<'a> {
17    raw: MlirExternalPass,
18    _reference: PhantomData<&'a MlirExternalPass>,
19}
20
21impl ExternalPass<'_> {
22    /// Signals that the pass has failed.
23    pub fn signal_failure(self) {
24        unsafe { mlirExternalPassSignalFailure(self.raw) }
25    }
26
27    /// Converts an external pass to a raw object.
28    pub const fn to_raw(self) -> MlirExternalPass {
29        self.raw
30    }
31
32    /// Creates an external pass handle from a raw object.
33    ///
34    /// # Safety
35    ///
36    /// A raw object must be valid.
37    pub const unsafe fn from_raw(raw: MlirExternalPass) -> Self {
38        Self {
39            raw,
40            _reference: PhantomData,
41        }
42    }
43}
44
45unsafe extern "C" fn callback_construct<'a, T: RunExternalPass<'a>>(pass: *mut T) {
46    pass.as_mut()
47        .expect("pass should be valid when called")
48        .construct();
49}
50
51unsafe extern "C" fn callback_destruct<'a, T: RunExternalPass<'a>>(pass: *mut T) {
52    pass.as_mut()
53        .expect("pass should be valid when called")
54        .destruct();
55    drop_in_place(pass);
56}
57
58unsafe extern "C" fn callback_initialize<'a, T: RunExternalPass<'a>>(
59    context: MlirContext,
60    pass: *mut T,
61) -> MlirLogicalResult {
62    pass.as_mut()
63        .expect("pass should be valid when called")
64        .initialize(ContextRef::from_raw(context));
65
66    MlirLogicalResult { value: 1 }
67}
68
69unsafe extern "C" fn callback_run<'a, T: RunExternalPass<'a>>(
70    operation: MlirOperation,
71    mlir_pass: MlirExternalPass,
72    pass: *mut T,
73) {
74    pass.as_mut()
75        .expect("pass should be valid when called")
76        .run(
77            OperationRef::from_raw(operation),
78            ExternalPass::from_raw(mlir_pass),
79        )
80}
81
82unsafe extern "C" fn callback_clone<'a, T: RunExternalPass<'a>>(pass: *mut T) -> *mut T {
83    Box::<T>::into_raw(Box::new(
84        pass.as_mut()
85            .expect("pass should be valid when called")
86            .clone(),
87    ))
88}
89
90/// A trait for MLIR passes written in Rust.
91///
92/// This trait is implemented for any type that implements `FnMut`,
93/// but can be implemented for any struct that implements `Clone`.
94///
95/// # Examples
96///
97/// The following example pass dumps operations.
98///
99/// ```
100/// use melior::{
101///     ir::OperationRef,
102///     pass::{ExternalPass, RunExternalPass},
103///     ContextRef,
104/// };
105///
106/// #[derive(Clone, Debug)]
107/// struct ExamplePass;
108///
109/// impl<'c> RunExternalPass<'c> for ExamplePass {
110///     fn construct(&mut self) {
111///         println!("Constructed pass!");
112///     }
113///
114///     fn initialize(&mut self, context: ContextRef<'c>) {
115///         println!("Initialize called!");
116///     }
117///
118///     fn run(&mut self, operation: OperationRef<'c, '_>, _pass: ExternalPass<'_>) {
119///         operation.dump();
120///     }
121/// }
122/// ```
123pub trait RunExternalPass<'c>: Sized + Clone {
124    fn construct(&mut self) {}
125    fn destruct(&mut self) {}
126    fn initialize(&mut self, context: ContextRef<'c>);
127    fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>);
128}
129
130impl<'c, F: FnMut(OperationRef<'c, '_>, ExternalPass<'_>) + Clone> RunExternalPass<'c> for F {
131    fn initialize(&mut self, _context: ContextRef<'c>) {}
132
133    fn run(&mut self, operation: OperationRef<'c, '_>, pass: ExternalPass<'_>) {
134        self(operation, pass)
135    }
136}
137
138/// Creates a `Pass` object from an external pass
139///
140/// # Examples
141///
142/// ```
143/// use melior::{
144///     Context,
145///     ir::{r#type::TypeId, OperationRef},
146///     pass::{create_external, ExternalPass},
147/// };
148///
149/// #[repr(align(8))]
150/// struct PassId;
151///
152/// static EXAMPLE_PASS: PassId = PassId;
153///
154/// let context = Context::new();
155///
156/// create_external(
157///     |operation: OperationRef, _pass: ExternalPass| {
158///         operation.dump();
159///     },
160///     TypeId::create(&EXAMPLE_PASS),
161///     "name",
162///     "argument",
163///     "description",
164///     "",
165///     &[],
166/// );
167/// ```
168#[allow(clippy::too_many_arguments)]
169pub fn create_external<'c, T: RunExternalPass<'c>>(
170    pass: T,
171    pass_id: TypeId,
172    name: &str,
173    argument: &str,
174    description: &str,
175    op_name: &str,
176    dependent_dialects: &[DialectHandle],
177) -> Pass {
178    unsafe {
179        Pass::from_raw(mlirCreateExternalPass(
180            pass_id.to_raw(),
181            StringRef::new(name).to_raw(),
182            StringRef::new(argument).to_raw(),
183            StringRef::new(description).to_raw(),
184            StringRef::new(op_name).to_raw(),
185            dependent_dialects.len() as isize,
186            dependent_dialects.as_ptr().cast_mut() as _,
187            MlirExternalPassCallbacks {
188                construct: Some(transmute::<*const (), unsafe extern "C" fn(*mut c_void)>(
189                    callback_construct::<T> as *const (),
190                )),
191                destruct: Some(transmute::<*const (), unsafe extern "C" fn(*mut c_void)>(
192                    callback_destruct::<T> as *const (),
193                )),
194                initialize: Some(transmute::<
195                    *const (),
196                    unsafe extern "C" fn(MlirContext, *mut c_void) -> MlirLogicalResult,
197                >(callback_initialize::<T> as *const ())),
198                run: Some(transmute::<
199                    *const (),
200                    unsafe extern "C" fn(MlirOperation, MlirExternalPass, *mut c_void),
201                >(callback_run::<T> as *const ())),
202                clone: Some(transmute::<
203                    *const (),
204                    unsafe extern "C" fn(*mut c_void) -> *mut c_void,
205                >(callback_clone::<T> as *const ())),
206            },
207            Box::into_raw(Box::new(pass)) as _,
208        ))
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use crate::{
216        dialect::func,
217        ir::{
218            attribute::{StringAttribute, TypeAttribute},
219            r#type::FunctionType,
220            Block, Identifier, Location, Module, Region,
221        },
222        pass::PassManager,
223        test::create_test_context,
224        Context,
225    };
226
227    #[repr(align(8))]
228    struct PassId;
229
230    fn create_module(context: &Context) -> Module {
231        let location = Location::unknown(context);
232        let module = Module::new(location);
233
234        module.body().append_operation(func::func(
235            context,
236            StringAttribute::new(context, "foo"),
237            TypeAttribute::new(FunctionType::new(context, &[], &[]).into()),
238            {
239                let block = Block::new(&[]);
240                block.append_operation(func::r#return(&[], location));
241
242                let region = Region::new();
243                region.append_block(block);
244                region
245            },
246            &[],
247            location,
248        ));
249        module
250    }
251
252    #[test]
253    fn external_pass() {
254        static TEST_PASS: PassId = PassId;
255
256        #[derive(Clone, Debug)]
257        struct TestPass<'c> {
258            context: &'c Context,
259            value: i32,
260        }
261
262        impl<'c> RunExternalPass<'c> for TestPass<'c> {
263            fn construct(&mut self) {
264                assert_eq!(self.value, 10);
265            }
266
267            fn destruct(&mut self) {
268                assert_eq!(self.value, 30);
269            }
270
271            fn initialize(&mut self, _context: ContextRef<'c>) {
272                assert_eq!(self.value, 10);
273                self.value = 20;
274            }
275
276            fn run(&mut self, operation: OperationRef<'c, '_>, _pass: ExternalPass<'_>) {
277                assert_eq!(self.value, 20);
278                self.value = 30;
279                assert!(operation.verify());
280                assert!(
281                    operation
282                        .region(0)
283                        .expect("module has a body")
284                        .first_block()
285                        .expect("module has a body")
286                        .first_operation()
287                        .expect("body has a function")
288                        .name()
289                        == Identifier::new(self.context, "func.func")
290                );
291            }
292        }
293
294        impl TestPass<'_> {
295            fn into_pass(self) -> Pass {
296                create_external(
297                    self,
298                    TypeId::create(&TEST_PASS),
299                    "test pass",
300                    "test argument",
301                    "a test pass",
302                    "",
303                    &[DialectHandle::func()],
304                )
305            }
306        }
307
308        let context = create_test_context();
309
310        let mut module = create_module(&context);
311        let pass_manager = PassManager::new(&context);
312
313        let test_pass = TestPass {
314            context: &context,
315            value: 10,
316        };
317        pass_manager.add_pass(test_pass.into_pass());
318        pass_manager.run(&mut module).unwrap();
319    }
320
321    #[test]
322    fn external_fn_pass_failure() {
323        static TEST_FN_PASS: PassId = PassId;
324
325        let context = create_test_context();
326
327        let mut module = create_module(&context);
328        let pass_manager = PassManager::new(&context);
329
330        pass_manager.add_pass(create_external(
331            |operation: OperationRef, pass: ExternalPass| {
332                assert!(operation.verify());
333                assert!(
334                    operation
335                        .region(0)
336                        .expect("module has a body")
337                        .first_block()
338                        .expect("module has a body")
339                        .first_operation()
340                        .expect("body has a function")
341                        .name()
342                        == Identifier::new(&context, "func.func")
343                );
344                pass.signal_failure();
345            },
346            TypeId::create(&TEST_FN_PASS),
347            "test closure",
348            "test argument",
349            "test",
350            "",
351            &[DialectHandle::func()],
352        ));
353        assert!(pass_manager.run(&mut module).is_err());
354    }
355}
melior/pass/external.rs

melior/pass/
external.rs
