How to use scan_expr in InclusiveScanKernel for more general operation

[dacapo1142]

I want to use a vector-valued binary operator in InclusiveScanKernel, the operation can be defined as follow,

array_length = 4  # a fix number
def op(lhs:List[int], rhs:List[int]):
    result = lhs.copy()
    for i in range(1, array_length):
        ridx = rhs[0]-lhs[i]
        if ridx <= array_length - 1:
            result[i] = rhs[ridx]
        else:
            result[i] = rhs[-1]
    return result

(It does have associative property for some carefully created vectors)

but the given example can only apply to number-valued function

knl = InclusiveScanKernel(np.int32, "a+b")

which means,

def op(a:int, b:int):
    return a+b

So, I got the following problems:

1. Can I pass a general function instead of a one-liner expression? like, using lambda expression in C/C++?
2. Can I pass a vector instead of a number?
3. Can I create a copy of a vector inside the function and return it?

I want to utilize the implementation of Work-Efficient Sum Scan Algorithm to speed up the process, that's why I want to use InclusiveScanKernel directly.

ref:
https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda
https://nvlabs.github.io/cub/structcub_1_1_device_scan.html#a16d7bc049ba8985dd89b10b3bcf0a8a3
https://www.tensorflow.org/api_docs/python/tf/scan
https://theano-pymc.readthedocs.io/en/latest/library/scan.html

[inducer]

Can I pass a general function instead of a one-liner expression? like, using lambda expression in C/C++?

Sure. The expression can be a call to a function defined in the preamble.

Can I pass a vector instead of a number?

As long as it's fixed-size, yes. Make a struct containing your fixed-size vector, and then register the type with pycuda.

Can I create a copy of a vector inside the function and return it?

Sure. Simple struct assignment would do that.

[henryliou-ephodtech]

I tried to write a POC following the code in test_gpuarray.py, but it didn't go well

import numpy as np
import pycuda.autoinit
from pycuda import gpuarray
from pycuda.scan import InclusiveScanKernel

preamble = """
struct array_4
{
    int data[4];
    // default constructor
    __device__ array_4(){}

    // copy constructor
    __device__ array_4(array_4 const &src){
        for (int idx = 0; idx <= 4; idx++){
            data[idx] = src.data[idx];
        }
    }

    // assign operator
    __device__ array_4 &operator=(array_4 const &src)
    {
        for (int idx = 0; idx <= 4; idx++){
            data[idx] = src.data[idx];
        }
        return *this;
    }

};

__device__ array_4 f(array_4 const &lhs, array_4 const &rhs){
    array_4 result;
    for (int idx = 0; idx <= 4; idx++){
        result.data[idx] = lhs.data[idx] + rhs.data[idx];
    }
    return result;
}
"""


dt = np.dtype([("data", (np.int32, (4,)))])
pycuda.tools.get_or_register_dtype(dtype=dt, c_names="array_4")
knl = InclusiveScanKernel(dtype=dt, scan_expr="f(a,b)", preamble=preamble)


a = np.array([[21, 56, 11, 98], [2, 86, 2, 2], [81, 18, 29, 92]], dtype=np.int32)
# array([[21, 56, 11, 98],
#        [ 2, 86,  2,  2],
#        [81, 18, 29, 92]], dtype=int32)

gpu_array = gpuarray.to_gpu(a.flatten())

gpu_result = knl(gpu_array).get().reshape(3, 4)
# array([[ 21,  56,  11,  98],
#        [104, 142,  13, 100],
#        [104, 160,  42, 192]], dtype=int32)


cpu_result = a.cumsum(axis=0)
# array([[ 21,  56,  11,  98],
#        [ 23, 142,  13, 100],
#        [104, 160,  42, 192]])

assert (gpu_result == cpu_result).all()
# fail :(

The input data type is a 4 elements vector, and the function f is vector summation.
Would you shed some light on what's going on here?

Beside the above problem, I also wonder if there is a way to use cupy array directly? Like, knl( wrapper( cupy.array([1,2,3]) ) ).
Or I have to copy cupy array to host and gpuarray.to_gpu()

[inducer]

Your for loops include index 4. (<=) They probably shouldn't.

[henryliou-ephodtech]

Sorry that's a careless mistake

After patching it, it works fine if the number of vectors is smaller than 3072
If the number increases, I get random results.
I've tried adding volatile to the functions, but it didn't help.
The inconsistency always occurs at 3072 (=12*256), and always at the first run.

import numpy as np
import pycuda.autoinit
from pycuda import gpuarray
from pycuda.scan import InclusiveScanKernel


preamble = """
struct array_4
{
    int data[4];
    // default constructor
    __device__ array_4(){}

    
    // copy constructor
    __device__ array_4(array_4 const &src){
        for (int idx = 0; idx < 4; idx++){
            data[idx] = src.data[idx];
        }
    }

    // assign operator
    __device__ array_4 &operator=(array_4 const &src)
    {
        for (int idx = 0; idx < 4; idx++){
            data[idx] = src.data[idx];
        }
        return *this;
    }

};

__device__ array_4 f(array_4 const &lhs, array_4 const &rhs) {
    array_4 result;
    for (int idx = 0; idx < 4; idx++){
        result.data[idx] = lhs.data[idx] + rhs.data[idx];
    }
    return result;
}
"""


dt = np.dtype([("data", (np.int32, 4))])
pycuda.tools.get_or_register_dtype(dtype=dt, c_names="array_4")
knl = InclusiveScanKernel(dtype=dt, scan_expr="f(a,b)", preamble=preamble)

from collections import defaultdict
from numpy.random import default_rng
rng = default_rng()

for round in range(100):
    length = 10000
    a = rng.choice(length*10, size=length*4, replace=False).astype(np.int32)
    a = a.reshape(length, 4)

    gpu_array = gpuarray.to_gpu(a.flatten())
    gpu_result = knl(gpu_array).get().reshape(length, 4)
    cpu_result = a.cumsum(axis=0)

    if (gpu_result != cpu_result).any():
        first_idx = np.argmin(gpu_result == cpu_result)
        print(f"inconsistent at round #{round}, at index #{first_idx}")
    
# inconsistent at round #0, at index #3072

[inducer]

You could look at the generated code and see if you can spot the bug. Just adding a print statement should suffice.