Calling into Rust from python

2

Here is how I've approached a similar situation: calling into Rust from Python for faster array-based computation. (Disclaimer -- there definitely may be some issues with this approach! But it's worked for me so far. I put this together from a number of helpful posts online. Any corrections or suggestions are welcome).

For my use case, I do a lot of I/O and pre-processing in Python and using NumPy. I then drop down into Rust for heavy lifting. I create the arrays in Python as NumPy arrays then pass these into Rust through pointers. Rust then performs in-place mutation.

In lib.rs:

#![crate_type = "dylib"]

use std::mem;

...

#[derive(Debug, Clone, Copy)]
pub struct Args {
    arg1: usize,
    arg2: f64
}

pub unsafe fn construct_array_1d<T>(input: *mut T, num_elems: usize) -> Vec<T> {
    // Create a Rust `Vec` associated with a `numpy` array via the raw pointer. This is unsafe in Rust parlance
    // because Rust can't verify that the pointer actually points to legitimate memory, so to speak.
    // At least check that the pointer is not null.
    assert!(!input.is_null());
    Vec::from_raw_parts(input, num_elems, num_elems)
}

pub fn do_computation(args: Args, v1: &Vec<f64>, v2: &Vec<f64>, mut v3: &mut Vec<f64>) {
    // Do in-place heavy lifting here.
    ...
}

#[no_mangle]
pub extern fn rust_entry_fn(arg1: usize, arg2: f64, v1_ptr: *mut f64,
    v1_num_elems: usize, v2_ptr: *mut f64, v2_num_elems: usize,
    v3_ptr: *mut f64, v3_num_elems: usize) {
    
    ...
    
    // Build ergonomic struct from input parameters.
    let args = Args { arg1: arg1, arg2 : arg2 };
    
    // Build immutable input vectors from raw parts.
    let v1 = unsafe { construct_array_1d(v1_ptr, v1_num_elems) };
    let v2 = unsafe { construct_array_1d(v2_ptr, v2_num_elems) };
    
    // Build mutable output vector from raw parts -- will mutate inplace in Rust.
    let mut v3 = unsafe { construct_array_1d(v3_ptr, v3_num_elems) };
    
    // Call into Rust function that performs computations and performs in-state
    // mutation of the output array `v3`.
    do_computation(args, &v1, &v2, &mut v3);
    
    // Don't let Rust free these arrays after leaving this function... This is
    // an FFI call from Python and the arrays are actually "owned" and
    // garbage collected by NumPy/Python.
    mem::forget(v1);
    mem::forget(v2);
    mem::forget(v3);
}

Then on the Python side, in main.py:

import numpy as np
from cffi import FFI

## Helper functions for Python-Rust FFI through `cffi`.
def _as_f64(num):
    """ Cast np.float64 for Rust."""
    return ffi.cast("double", num)

def _as_f64_array(array):
    """ Cast np.float64 array to a pointer to float 64s."""
    return ffi.cast("double*", num)

def _as_usize(num):
    """ Cast `num` to Rust `usize`."""
    return ffi.cast("unsigned long", num)

## Instantiate `cffi` object.
ffi = FFI()
ffi.cdef(
    """
    void* rust_entry_fn(unsigned long arg1, double arg2, double* v1_ptr, unsigned long v1_num_elems,
        double* v2_ptr, unsigned long v2_num_elems, double* v3_ptr, unsigned long v3_num_elems);
    """
)

## Go get the Rust library.
lib = ffi.dlopen("../target/release/mylib.dylib")

arg1 = _as_usize(calc_arg1(...))
arg2 = _as_f64(calc_arg2(...))
v1 = calc_input_array_v1(...)
v1_ptr = _as_f64_array(v1)
v1_num_elems = _as_usize(v1.size)
v2 = calc_input_array_v2(...)
v2_ptr = _as_f64_array(v2)
v2_num_elems = _as_usize(v2.size)
v3 = np.zeros_like(v1)
v3_ptr = _as_f64_array(v3)
v3_num_elems = _as_usize(v3.size)

lib.rust_entry_fn(arg1, arg2, v1_ptr, v1_num_elems, v2_ptr, v2_num_elems, v3_ptr, v3_num_elems)

... ## Do stuff with `v3`.
5 Likes