Skip to content

Instantly share code, notes, and snippets.

@brandonwillard
Last active October 15, 2021 17:46
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save brandonwillard/5533fd777c2c5f6482d8a89a372f1bf4 to your computer and use it in GitHub Desktop.
Save brandonwillard/5533fd777c2c5f6482d8a89a372f1bf4 to your computer and use it in GitHub Desktop.
Custom NumPy C API extension testing setup (look at the raw document)

NumPy C API Testing

1 Development Setup

First, clone the NumPy project, then create a Conda env:

# Create the venv
mamba env create -n numpy-dev -f environment.yml --no-default-packages
# Activate the venv
conda activate numpy-dev
# Setup submodules
git submodule update --init

Now, we need to build the NumPy project’s extensions:

# Build with debug symbols:
# This didn't work the first time, but that might've been for a different reason.
# pip install --no-binary :all: --global-option build --global-option --debug -e ./

python setup.py build_ext --inplace -j 4 --debug
pip install -e ./

This describes how to install Python debug symbols in a Conda environment. I believe this would also include some helpful Python gdb commands (after a few more setup steps).

For src_elisp[:eval never]{lsp-clangd} to work correctly, we need a compile_commands.json. One could probably be generated using the spec (i.e. spec) and the parsed output from src_bash[:eval never]{python setup.py build_ext –inplace –debug –force >! build_ext_output.txt}.

2 TODO Investigate the fancy/advanced indexing, src_python[:eval never]{ufunc.reduce}, and src_python[:eval never]{ufunc.at} implementations

The implementations for src_python[:eval never]{ufunc.reduce} and src_python[:eval never]{ufunc.at} are:

  • PyUFunc_GenericReduction
  • ufunc_at

This is the implementation for fancy/advanced indexing: PyArray_MapIterArrayCopyIfOverlap. Here’s a high-level explanation.

The unit tests are here: TestFancyIndexing and TestMapIter.

Here’s a reference for some NumPy C types.

3 Testing A Custom Python Extension

Use src_elisp[:eval never]{org-babel-tangle} to create the files corresponding to the blocks that follow.

#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION

#include <stdio.h>
/* #include <math.h> */
/* #include <float.h> */
#include <Python.h>
#include <numpy/arrayobject.h>

static PyObject *inplace_increment(PyObject *self, PyObject *args);

static PyMethodDef module_methods[] = {
  {"inplace_increment", (PyCFunction)inplace_increment, METH_VARARGS, NULL},
  {NULL, NULL, 0, NULL}
};

static struct PyModuleDef mymod_def = {
    PyModuleDef_HEAD_INIT,
    "test_mod",
    NULL,
    -1,
    module_methods
};

PyMODINIT_FUNC PyInit_test_mod(void)
{
    PyObject *m = PyModule_Create(&mymod_def);
    if (m == NULL)
        return NULL;

    /* Load `numpy` functionality. */
    import_array();
    return m;
}

typedef void (*inplace_map_binop)(PyArrayMapIterObject *, PyArrayIterObject *);

static void npy_float64_inplace_add(PyArrayMapIterObject *mit, PyArrayIterObject *it)
{
  int index = mit->size;
  while (index--) {
    ((npy_float64*)mit->dataptr)[0] = ((npy_float64*)mit->dataptr)[0] + ((npy_float64*)it->dataptr)[0];

    PyArray_MapIterNext(mit);
    PyArray_ITER_NEXT(it);
  }
}

/*
 * The following code is from `_multiarray_tests.c.src`
 */
static int
map_increment(PyArrayMapIterObject *mit, PyObject *op, inplace_map_binop add_inplace)
{
  PyArrayObject *arr = NULL;
  PyArrayIterObject *it;
  PyArray_Descr *descr;

  if (mit->ait == NULL) {
    return -1;
  }
  descr = PyArray_DESCR(mit->ait->ao);
  Py_INCREF(descr);
  arr = (PyArrayObject *)PyArray_FromAny(op, descr,
                                         0, 0, NPY_ARRAY_FORCECAST, NULL);
  if (arr == NULL) {
    return -1;
  }

  if ((mit->subspace != NULL) && (mit->consec)) {
    PyArray_MapIterSwapAxes(mit, (PyArrayObject **)&arr, 0);
    if (arr == NULL) {
      return -1;
    }
  }

  if ((it = (PyArrayIterObject *)\
       PyArray_BroadcastToShape((PyObject *)arr, mit->dimensions,
                                mit->nd)) == NULL) {
    Py_DECREF(arr);

    return -1;
  }

  (*add_inplace)(mit, it);

  Py_DECREF(arr);
  Py_DECREF(it);
  return 0;
}

static PyObject *
inplace_increment(PyObject *self, PyObject *args)
{
    PyObject *arg_a = NULL, *index=NULL, *inc=NULL;
    PyArrayObject *a;
    inplace_map_binop add_inplace = npy_float64_inplace_add;
    int i =0;
    PyArrayMapIterObject * mit;

    if (!PyArg_ParseTuple(args, "OOO", &arg_a, &index,
            &inc)) {
        return NULL;
    }
    if (!PyArray_Check(arg_a)) {
         PyErr_SetString(PyExc_ValueError, "needs an ndarray as first argument");
         return NULL;
    }
    a = (PyArrayObject *) arg_a;

    if (PyArray_FailUnlessWriteable(a, "input/output array") < 0) {
        return NULL;
    }

    if (PyArray_NDIM(a) == 0) {
        PyErr_SetString(PyExc_IndexError, "0-d arrays can't be indexed.");
        return NULL;
    }

    if (add_inplace == NULL) {
        PyErr_SetString(PyExc_TypeError, "unsupported type for a");
        return NULL;
    }

    // Defined at [[file:numpy/core/include/numpy/ndarraytypes.h::typedef struct {]]
    mit = (PyArrayMapIterObject *) PyArray_MapIterArray(a, index);

    fprintf(stdout, "mit.subspace = ");
    PyObject_Print(mit->subspace, stdout, 0);
    /* PyObject_Print(mit->array, stdout, 0); */
    /* PyObject_Print(mit->outer, stdout, 0); */
    /* PyObject_Print(mit->extra_op, stdout, 0); */
    /* fprintf(stdout, "mit.fancy_strides = %s", mit.fancy_strides); */
    fprintf(stdout, "\n");

    if (mit == NULL) {
        goto fail;
    }

    if (map_increment(mit, inc, add_inplace) != 0) {
        goto fail;
    }

    Py_DECREF(mit);

    Py_RETURN_NONE;
fail:
    Py_XDECREF(mit);

    return NULL;
}

Create the setup.py file for this extension package:

from distutils.core import setup, Extension
import numpy.distutils.misc_util

c_ext = Extension(
    "test_mod",
    sources=["test_mod.c"],
    include_dirs=numpy.distutils.misc_util.get_numpy_include_dirs(),
    # library_dirs = [os.getcwd(),],  # path to .a or .so file(s)
    extra_compile_args=[
        "-g3",
        "-O0",
        "-DDEBUG=0",
        "-UNDEBUG",
        # "-std=c++11",
    ],
    language="c"  # "c++11"
)

setup(name="test_mod", version="0.0.1", ext_modules=[c_ext])

Compile the extension (rerun this every time the C code is changed):

conda activate numpy-dev
python setup.py build_ext --inplace --debug --force

If you want the extension package to be available from anywhere (other than the project root), run this once:

conda activate numpy-dev
pip install -e ./

Importing this module will run the extension code:

import numpy as np

import test_mod


a = np.arange(np.prod((5, 4, 3, 2, 1))).reshape((5, 4, 3, 2, 1)).astype(float)
index = ([1, 1, 2, 0], 0, [0, 0, 1, 2])
vals = 1

test_mod.inplace_increment(a, index, vals)

For more immediate debugging, the run_test module can be run from within gdb and breakpoints can be set (e.g. run gdb python, set b /tmp/test_ext/test_mod.c:125, then run with r).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment