Plugin

Plugins for PyQUDA

PyQUDA can build Cython warpper files from a C header file.

Example

We use the contract plugin as an example.

Consider that we want to expose some functions declaried in contract.h, and the definitions are compiled into a shared library libcontract.so.

First, we need to compile the shared library.

git clone https://github.com/CLQCD/contract
cd contract
mkdir -p build
cd build
cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=$(pwd)/install -DGPU_ARCH=60
cmake --build . -j8 && cmake --install .
cd ../..

And we want to expose functions declared in contract.h:

#pragma once

#include "contract_define.h"

#ifdef __cplusplus
extern "C" {
#endif

typedef enum {
  IK_JL_MN,
  IK_JN_ML,
  IL_JK_MN,
  IL_JN_MK,
  IN_JK_ML,
  IN_JL_MK,
} BaryonContractType;

void init(int device);
void baryon_two_point(void *correl, void *propag_i, void *propag_j, void *propag_m, BaryonContractType contract_type,
                      unsigned long volume, int gamma_ij, int gamma_kl, int gamma_mn);
void proton(void *correl, void *propag_i, void *propag_j, void *propag_m, int contract_type, unsigned long volume,
            int gamma_ij, int gamma_kl, int gamma_mn);

#ifdef __cplusplus
}
#endif

After installing PyQUDA-Utils by pip install pyquda-utils, you can use the following command to build the plugin to wrap libcontract.so:

cd contract/build
python -m pyquda_plugins -i contract.h -l contract -I $(pwd)/install/include -L $(pwd)/install/lib

pyquda_plugins will build Cython source files contract.pxd, _pycontract.pyx along with the stub file _pycontract.pyi and then compile them into pyquda_plugins.pycontract._pycontract module. Arguments of C basic type will be transferred from corresponding Python objects as input, and pointers will be converted from numpy.ndarray, cupy.ndarray or torch.Tensor.

contract.pxd:

cdef extern from "contract.h":
    ctypedef enum BaryonContractType:
        pass

    void init(int device)
    void baryon_two_point(void *correl, void *propag_i, void *propag_j, void *propag_m, BaryonContractType contract_type, unsigned long volume, int gamma_ij, int gamma_kl, int gamma_mn)
    void proton(void *correl, void *propag_i, void *propag_j, void *propag_m, int contract_type, unsigned long volume, int gamma_ij, int gamma_kl, int gamma_mn)

_pycontract.pyx:

from enum import IntEnum
from libcpp cimport bool
from numpy cimport ndarray
from pyquda_comm.pointer cimport Pointer, _NDArray
cimport contract

class BaryonContractType(IntEnum):
    IK_JL_MN = 0
    IK_JN_ML = 1
    IL_JK_MN = 2
    IL_JN_MK = 3
    IN_JK_ML = 4
    IN_JL_MK = 5


def init(int device):
    contract.init(device)

def baryon_two_point(correl, propag_i, propag_j, propag_m, contract.BaryonContractType contract_type, unsigned long volume, int gamma_ij, int gamma_kl, int gamma_mn):
    _correl = _NDArray(correl, 1)
    _propag_i = _NDArray(propag_i, 1)
    _propag_j = _NDArray(propag_j, 1)
    _propag_m = _NDArray(propag_m, 1)
    contract.baryon_two_point(<void *>_correl.ptr, <void *>_propag_i.ptr, <void *>_propag_j.ptr, <void *>_propag_m.ptr, contract_type, volume, gamma_ij, gamma_kl, gamma_mn)

def proton(correl, propag_i, propag_j, propag_m, int contract_type, unsigned long volume, int gamma_ij, int gamma_kl, int gamma_mn):
    _correl = _NDArray(correl, 1)
    _propag_i = _NDArray(propag_i, 1)
    _propag_j = _NDArray(propag_j, 1)
    _propag_m = _NDArray(propag_m, 1)
    contract.proton(<void *>_correl.ptr, <void *>_propag_i.ptr, <void *>_propag_j.ptr, <void *>_propag_m.ptr, contract_type, volume, gamma_ij, gamma_kl, gamma_mn)

_pycontract.pyi:

from numpy import int32, float64, complex128
from numpy.typing import NDArray

def gwu_init_machine(latt_size: NDArray[int32]) -> None: ...
def gwu_shutdown_machine() -> None: ...
def gwu_build_hw(links_in: NDArray, kappa: float) -> None: ...
def gwu_load_hw_eigen(hw_eignum: int, hw_eigprec: float, hw_eigvals: NDArray[complex128], hw_eigvecs: NDArray) -> None: ...
def gwu_build_ov(ov_poly_prec: float, ov_use_fp32: int) -> None: ...
def gwu_load_ov_eigen(ov_eignum: int, ov_eigprec: float, ov_eigvals: NDArray[complex128], ov_eigvecs: NDArray) -> None: ...
def gwu_build_hw_eigen(hw_eignum: int, hw_eigprec: float, hw_extra_krylov: int, maxiter: int, chebyshev_order: int, chebyshev_cut: float, iseed: int) -> None: ...
def gwu_invert_overlap(propag_in: NDArray, source_in: NDArray, num_mass: int, masses: NDArray[float64], tol: float, maxiter: int, one_minus_half_d: int, mode: int) -> None: ...

Then you can use functions from libcontract.so in Python. We recommend to build a pure Python file for better user interface. You can check pycontract as an example.

Limitations

ABI

Because we compile the header in C mode, ensure that all the declarations are wrapped with

#ifdef __cplusplus
extern "C" {
#endif

// function declarations

#ifdef __cplusplus
}
#endif

Argument type

Because of Cython's automatic type conversions, we have some limitations on the argument type:

C type	Python type
[unsigned] int, [unsigned] long, [unsigned] long long	int
float, double	float
float _Complex, double _Complex	complex
int *, int **, int ***	numpy.ndarray[int32][^1]
double *, double **, double ***	numpy.ndarray[float64][^1]
double _Complex *, double _Complex **, double _Complex ***	numpy.ndarray[complex128][^1]
const char *, const char []	bytes
void *, void **, void ***	numpy.ndarray, cupy.ndarray, torch.Tensor[^2]

[^1]: Typed pointers can be converted only from numpy.ndarray with corresponding ndim [^2]: Untyped pointers can be converted from numpy.ndarray, cupy.ndarray or torch.Tensor with corresponding ndim

• struct is unacceptable as the arguments for now. We are working on it.
• union cannot be the argument type

Return type

Only basic arithmetic type are allowed besides void pointer. Returned void pointers are usually instance in C/C++.

C type	Python type
[unsigned] int, [unsigned] long, [unsigned] long long	int
float, double	float
float _Complex, double _Complex	complex
void *	pyquda_comm.pointer.Pointer[^3]

[^3]: Pointer handles the pointer returned from C/C++, remember to free it in C/C++

Include

We use pycparser to parse the header, however it uses fake_libc_include to parse standard library headers, which cannot be easily installed by pip install pycparser. So you cannot include any standard library header such as stdlib.h in the header.