xref: /libCEED/examples/rust-qfunctions/ex1-volume-rs/src/lib.rs (revision 5a526491291e2ef13670ec99232a2cb0069702e5)
1 #![no_std]
2 #![allow(internal_features)]
3 #![feature(asm_experimental_arch, abi_ptx, core_intrinsics)]
4 use core::ffi::c_void;
5 use core::intrinsics::abort;
6 use core::panic::PanicInfo;
7 
8 use ndarray::ArrayView;
9 
10 // This is a dummy allocator that always returns null. Heap allocations do not work on GPUs
11 use core::alloc::{GlobalAlloc, Layout};
12 pub struct Allocator;
13 unsafe impl GlobalAlloc for Allocator {
14     unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
15         0 as *mut u8
16     }
17     unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
18         abort(); // since we never allocate
19     }
20 }
21 #[global_allocator]
22 static GLOBAL_ALLOCATOR: Allocator = Allocator;
23 
24 // This is a copy of the same data structure defined in the .h file. It can be autogenerated using bindgen/cbindgen
25 #[doc = " A structure used to pass additional data to f_build_mass"]
26 #[repr(C)]
27 #[derive(Debug, Copy, Clone)]
28 pub struct BuildContext {
29     pub dim: i32,
30     pub space_dim: i32,
31 }
32 
33 // On no_std targets, its required to implement your own panic function.
34 #[panic_handler]
35 fn panic(_info: &PanicInfo) -> ! {
36     abort()
37 }
38 
39 /* The no_mangle is required because rust "mangles" names (changes them to prevent namespace conflicts)
40 Also note that this function ends in _rs, even though the C call `CEED_QFUNCTION_RUST(build_mass)` doesn't */
41 #[no_mangle]
42 pub unsafe extern "C" fn build_mass_rs(
43     ctx: *mut c_void,
44     q: i32,
45     in_: *const *const f64,
46     out: *mut *mut f64,
47 ) -> i8 {
48     let ctx: *mut BuildContext = unsafe { core::mem::transmute(ctx) };
49     let ctx: &mut BuildContext = &mut *ctx;
50 
51     let in_slice = core::slice::from_raw_parts(in_, 2);
52 
53     // in_slice[0] is Jacobians with shape [dim, dim, Q]
54     // in_slice[1] is quadrature weights with shape [1, Q]
55     let j_ptr = in_slice[0];
56     let w_ptr = in_slice[1];
57 
58     let j = ArrayView::from_shape_ptr((ctx.dim as usize, ctx.dim as usize, q as usize), j_ptr);
59 
60     let w = core::slice::from_raw_parts(w_ptr, q as usize);
61 
62     let out_slice = core::slice::from_raw_parts_mut(out, 1);
63     let q_data = core::slice::from_raw_parts_mut(out_slice[0], q as usize);
64 
65     match ctx.dim * 10 + ctx.space_dim {
66         11 => {
67             // Quadrature Point Loop
68             for i in 0..q as usize {
69                 q_data[i] = j[[0, 0, i]] * w[i];
70             }
71         }
72         22 => {
73             // Quadrature Point Loop
74             for i in 0..q as usize {
75                 q_data[i] = (j[[0, 0, i]] * j[[1, 1, i]] - j[[0, 1, i]] * j[[1, 0, i]]) * w[i];
76             }
77         }
78         33 => {
79             // Quadrature Point Loop
80             for i in 0..q as usize {
81                 q_data[i] = (j[[0, 0, i]]
82                     * (j[[1, 1, i]] * j[[2, 2, i]] - j[[1, 2, i]] * j[[2, 1, i]])
83                     - j[[0, 1, i]] * (j[[1, 0, i]] * j[[2, 2, i]] - j[[1, 2, i]] * j[[2, 0, i]])
84                     + j[[0, 2, i]] * (j[[1, 0, i]] * j[[2, 1, i]] - j[[1, 1, i]] * j[[2, 0, i]]))
85                     * w[i];
86             }
87         }
88         _ => {
89             abort();
90         }
91     }
92 
93     0
94 }
95 
96 /* The no_mangle is required because rust "mangles" names (changes them to prevent namespace conflicts)
97 Also note that this function ends in _rs, even though the C call `CEED_QFUNCTION_RUST(apply_mass)` doesn't
98 For FFI reasons, it is also required to include all parameters in this exact form, even if you don't use all of them*/
99 #[no_mangle]
100 pub unsafe extern "C" fn apply_mass_rs(
101     _ctx: *mut c_void,
102     q: i32,
103     in_: *const *const f64,
104     out: *mut *mut f64,
105 ) -> i8 {
106     let in_slice = core::slice::from_raw_parts(in_, 2);
107 
108     let u_ptr = in_slice[0];
109     let q_data_ptr = in_slice[1];
110 
111     let u = core::slice::from_raw_parts(u_ptr, q as usize);
112     let q_data = core::slice::from_raw_parts(q_data_ptr, q as usize);
113 
114     let out_slice = core::slice::from_raw_parts_mut(out, 1);
115 
116     let v_ptr = out_slice[0];
117     let v = core::slice::from_raw_parts_mut(v_ptr, q as usize);
118 
119     for i in 0..q as usize {
120         v[i] = q_data[i] * u[i];
121     }
122 
123     0
124 }
125