ARTICLE AD BOX
I'm defining a CuPy ElementwiseKernel for computations on GPU with Python. I have several input and output arrays, all multidimensional, of differing shapes, so I'm wondering how the ElementwiseKernel's i variable is assigned (if they were all of shape (m, n), I image i would simply run from 0 to m * n - 1, but here I have many different m's and n's).
To be specific, inputs uu, vv and outputs u, v have shape (H, W). Input cc has shape (H, W, 2), and its corresponding output c has shape (H, W, 2, 2). Inputs Ix, Iy, It all have shape (H + K, W + K) (they're derived from a "padding" of the data from which uu, vv, and cc are derived). Finally, input R has shape (K + 1, K + 1). The remaining five inputs are scalars.
The nature of the computation is "elementwise" on uu, vv, cc, meaning each operation uses a single entry (or in cc's case, a 2-tuple) in its computation. Meanwhile, each operation uses a small window of entries in Ix, Iy, It, and uses every entry in R. I'm wondering, for example, how from i I might get my current element's row and column index in, say, uu, because I'm not sure i corresponds to uu's shape / size, thought that's my current assumption, since uu is the first argument, and I can't find clarification in the documentation.
I'm also not sure exactly how specifying "raw" on these arrays might affect the nature of i or the indexing in general.
Any and all guidance is greatly appreciated.
kernel = cp.ElementwiseKernel( in_params="raw float32 uu, raw float32 vv, raw float32 cc," "raw float32 Ix, raw float32 Iy, raw float32 It," "raw float32 R," "int32 h, int32 w, int32 ksize, int32 padw, int32 Rw", out_params="raw float32 u, raw float32 v, raw float32 c", operation= """ int idx = i; // how does this "i" work? // ... (rest of operation) """ )