xref: /libCEED/examples/rust-qfunctions/ex1-volume-rs/src/lib.rs (revision 2027fb9d13fe34211738d8539f90542a9801ae2c)

#![no_std]
#![allow(internal_features)]
#![feature(asm_experimental_arch, abi_ptx, core_intrinsics)]
use core::ffi::c_void;
use core::intrinsics::abort;
use core::panic::PanicInfo;

use ndarray::ArrayView;

// This is a dummy allocator that always returns null. Heap allocations do not work on GPUs
use core::alloc::{GlobalAlloc, Layout};
pub struct Allocator;
unsafe impl GlobalAlloc for Allocator {
    unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
        0 as *mut u8
    }
    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
        abort(); // since we never allocate
    }
}
#[global_allocator]
static GLOBAL_ALLOCATOR: Allocator = Allocator;

// This is a copy of the same data structure defined in the .h file. It can be autogenerated using bindgen/cbindgen
#[doc = " A structure used to pass additional data to f_build_mass"]
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct BuildContext {
    pub dim: i32,
    pub space_dim: i32,
}

// On no_std targets, its required to implement your own panic function.
#[panic_handler]
fn panic(_info: &PanicInfo) -> ! {
    abort()
}

/* The no_mangle is required because rust "mangles" names (changes them to prevent namespace conflicts)
Also note that this function ends in _rs, even though the C call `CEED_QFUNCTION_RUST(build_mass)` doesn't */
#[no_mangle]
pub unsafe extern "C" fn build_mass_rs(
    ctx: *mut c_void,
    q: i32,
    in_: *const *const f64,
    out: *mut *mut f64,
) -> i8 {
    let ctx: *mut BuildContext = unsafe { core::mem::transmute(ctx) };
    let ctx: &mut BuildContext = &mut *ctx;

    let in_slice = core::slice::from_raw_parts(in_, 2);

    // in_slice[0] is Jacobians with shape [dim, dim, Q]
    // in_slice[1] is quadrature weights with shape [1, Q]
    let j_ptr = in_slice[0];
    let w_ptr = in_slice[1];

    let j = ArrayView::from_shape_ptr((ctx.dim as usize, ctx.dim as usize, q as usize), j_ptr);

    let w = core::slice::from_raw_parts(w_ptr, q as usize);

    let out_slice = core::slice::from_raw_parts_mut(out, 1);
    let q_data = core::slice::from_raw_parts_mut(out_slice[0], q as usize);

    match ctx.dim * 10 + ctx.space_dim {
        11 => {
            // Quadrature Point Loop
            for i in 0..q as usize {
                q_data[i] = j[[0, 0, i]] * w[i];
            }
        }
        22 => {
            // Quadrature Point Loop
            for i in 0..q as usize {
                q_data[i] = (j[[0, 0, i]] * j[[1, 1, i]] - j[[0, 1, i]] * j[[1, 0, i]]) * w[i];
            }
        }
        33 => {
            // Quadrature Point Loop
            for i in 0..q as usize {
                q_data[i] = (j[[0, 0, i]]
                    * (j[[1, 1, i]] * j[[2, 2, i]] - j[[1, 2, i]] * j[[2, 1, i]])
                    - j[[0, 1, i]] * (j[[1, 0, i]] * j[[2, 2, i]] - j[[1, 2, i]] * j[[2, 0, i]])
                    + j[[0, 2, i]] * (j[[1, 0, i]] * j[[2, 1, i]] - j[[1, 1, i]] * j[[2, 0, i]]))
                    * w[i];
            }
        }
        _ => {
            abort();
        }
    }

    0
}

/* The no_mangle is required because rust "mangles" names (changes them to prevent namespace conflicts)
Also note that this function ends in _rs, even though the C call `CEED_QFUNCTION_RUST(apply_mass)` doesn't
For FFI reasons, it is also required to include all parameters in this exact form, even if you don't use all of them*/
#[no_mangle]
pub unsafe extern "C" fn apply_mass_rs(
    _ctx: *mut c_void,
    q: i32,
    in_: *const *const f64,
    out: *mut *mut f64,
) -> i8 {
    let in_slice = core::slice::from_raw_parts(in_, 2);

    let u_ptr = in_slice[0];
    let q_data_ptr = in_slice[1];

    let u = core::slice::from_raw_parts(u_ptr, q as usize);
    let q_data = core::slice::from_raw_parts(q_data_ptr, q as usize);

    let out_slice = core::slice::from_raw_parts_mut(out, 1);

    let v_ptr = out_slice[0];
    let v = core::slice::from_raw_parts_mut(v_ptr, q as usize);

    for i in 0..q as usize {
        v[i] = q_data[i] * u[i];
    }

    0
}