CODEWITHSUNDEEP
pythonc++templatescastingcython
hbristow
hbristow

Reputation: 1896

Casting Cython fused types to C++ pointers

This is a general question about casting from Cython fused types to C++ types, which I'll describe with a minimal example. Consider the superficial C++ function template:

template <typename T>
void scale_impl(const T * x, T * y, const T a, const size_t N) {
    for (size_t n = 0; n < N; ++n) {
        y[n] = a*x[n];
    }
}

I would like to be able to call this function on any numpy ndarray of any type and shape. Using Cython, we first declare the function template:

cdef extern:
    void scale_impl[T](const T * x, T * y, const T a, const size_t N)

then declare the valid scalar types we wish to operate on:

ctypedef fused Scalar:
    float
    double
    ...

And finally implement the actual Cython shim:

def scale(ndarray[Scalar] x, Scalar a):
    """Scale an array x by the value a"""
    cdef ndarray[Scalar] y = np.empty_like(x)
    scale_impl(<Scalar *>x.data, <Scalar *>y.data, a, x.size)
    return y

This doesn't work for two reasons:

One could obviously deduce the specializations explicitly:

    if Scalar is float:
        scale_impl(<float *>x.data, <float *>y.data, a, x.size)
    if Scalar is double:
        scale_impl(<double *>x.data, <double *>y.data, a, x.size)
    if Scalar is ...

But this results in a combinatorial number of code paths I have to hand-write for functions that entertain multiple fused types, and creates the very situation (I assume) fused types were introduced to avoid.

Is there any way of passing an arbitrary dimensional (within reason) array to a Cython function and have it deduce the pointer type of the scalar data? Or, what is the most reasonable compromise to approach such functionality?

Upvotes: 4

Views: 1125

Answers (1)

DavidW
DavidW

Reputation: 30898

(See also the answer given in Using Cython to wrap a c++ template to accept any numpy array which is a very similar problem.)

Using the form &x[0] instead of trying to cast x.data solves the problem of picking the right template specialization. The problem of 2D arrays is a bit more complicated because the array isn't guaranteed to be either continuous or ordered.

I'd create a function that does the actual work on a 1D array, and wrap in a simple function that flattens as needed:

def _scale_impl(Scalar[::1] x, Scalar a):
  # the "::1" syntax ensures the array is actually continuous
  cdef np.ndarray[Scalar,ndim=1] y = np.empty_like(x)
  cdef size_t N = x.shape[0] # this seems to be necessary to avoid throwing off Cython's template deduction
  scale_impl(&x[0],&y[0],a,N)
  return y


def scale(x, a):
  """Scale an array x by the value a"""

  y = _scale_impl(np.ravel(x),a)
  return y.reshape(x.shape) # reshape needs to be kept out of Cython

Upvotes: 1

Related Questions