Passing multi-dimensional arrays to C++ functions via custom_call for custom primitive

[maurorigo]

Hello! Hopefully this is not a trivial question.

I need to define a custom primitive which has some multi-dimensional inputs (for the time being, only for CPU). For now, for simplicity, I'm just flattening the inputs on the python side to pass them to custom_call and unflattening them in C++. On the C++ side, the XLA custom call looks like this (I'm putting the whole code, but the important parts are the ones associated to pos, edgesx, edgesy, edgesz and eventually also field):

template <typename T>
void cppaint(void* out, void** in){
    // Parse inputs (flattened for simplicity)
    T* _pos = reinterpret_cast<T*>(in[0]);
    int Nparts = *reinterpret_cast<int*>(in[1]);
    T* mass = reinterpret_cast<T*>(in[2]);
    int* Nmesh = reinterpret_cast<int*>(in[3]);
    int* _edgesx = reinterpret_cast<int*>(in[4]);
    int* _edgesy = reinterpret_cast<int*>(in[5]);
    int* _edgesz = reinterpret_cast<int*>(in[6]);
    MPI_Comm comm = reinterpret_cast<MPI_Comm>(*reinterpret_cast<uintptr_t*>(in[7]));

    // Unflatten arrays
    T** pos = new T*[Nparts];
    for (int i=0; i<Nparts; i++) pos[i] = &(_pos[i * 3]); // 3 is dimensionality

    int comm_size;
    MPI_Comm_size(comm, &comm_size);
    int** edgesx = new int*[comm_size];
    int** edgesy = new int*[comm_size];
    int** edgesz = new int*[comm_size];
    for (int i=0; i<comm_size; i++){
        edgesx[i] = &(_edgesx[i * 2]); // 2 is lower lim and upper lim
        edgesy[i] = &(_edgesy[i * 2]);
        edgesz[i] = &(_edgesz[i * 2]);
    }

    // Output
    T*** field = ppaint<T>(pos, Nparts, mass, Nmesh, edgesx, edgesy, edgesz, comm);
    // TEMPORARY, return positions
    T* outp = reinterpret_cast<T*>(out);
    for (int i=0; i<3*Nparts; i++) outp[i] = _pos[i];
}

while the lowering function in python is:

def _ppaint_lowering(ctx, pos, mass, Nmesh, edgesx, edgesy, edgesz, comm):

    comm = unpack_hashable(comm)

    # Extract the numpy type of the inputs
    pos_aval = ctx.avals_in[0]
    np_dtype = np.dtype(pos_aval.dtype)

    out_type = mlir.aval_to_ir_type(ctx.avals_out[0])

    # Number of particles is length of pos / 3
    Nparts = (np.prod(pos_aval.shape) / 3).astype(np.int64)

    # Dealing with comm as in mpi4jax, see for instance barrier.py in collective ops
    comm = as_mhlo_constant(to_mpi_handle(comm), np.uintp)

    # Dispatch a different call depending on the dtype
    if np_dtype == np.float32:
        op_name = "ppaint_f32"
    elif np_dtype == np.float64:
        op_name = "ppaint_f64"
    else:
        raise NotImplementedError(f"Unsupported dtype {np_dtype}")

    return custom_call(
        op_name,
        # Output types
        result_types=[out_type],
        # The inputs:
        operands=[pos, mlir.ir_constant(Nparts), mass, Nmesh, edgesx, edgesy, edgesz, comm],
    ).results

For example, originally pos in python is (Nparts, 3), and I'm passing it to this lowering function flattened. As you can see, on the C++ side I need to pass a T** to ppaint (it could also be a T*[3] if it's easier to deal with), and for now I'm just expecting a T* (_pos) from the XLA custom call, defining pos via:

T** pos = new T*[Nparts];
for (int i=0; i<Nparts; i++) pos[i] = &(_pos[i * 3]); // 3 is dimensionality

My question is: how could I do this in an easier way? In other words, how should I pass for instance pos (and what layout should I specify, if needed) to the custom_call in python in order to have directly T** pos = reinterpret_cast<T**>(in[0]); on the C++ side?
I hope this is clear enough, thank you in advance!