Support more numeric types in Groupby.apply with engine='jit'

[brandon-b-miller]

draft

This PR adds additional numeric dtypes to GroupBy.apply with engine='jit'.

[brandon-b-miller]

I ran into some issues compiling the shim library with int16 and int8 dtypes here. It looks like our implementations of the underlying block level reduction functions rely on libcu++ atomic_ref which per the documentation may only be instantiated with a T that are either 4 or 8 bytes. This seems consistent with the error I'm getting:

atomic_cuda.h(111): error: static assertion failed with "atomic_ref does not support 1 or 2 byte types

This means that to support shorter ints we will need to do a little more involved development on the c++ side that I think should be separated out into another PR. I'm updating the PR title to reflect that this PR itself only goes as far as rounding us out with int32 and float32.

[brandon-b-miller]

raised #13736

[bdice]

I'm not a huge fan of repeating the typecasting logic twice in Python and once in C++... we'll need to find a better architecture for this before we get too deep into adding more types.

[bdice]

It feels like the typecasting rules were hard-coded twice. Once here, and once in GroupSum. Is there a way to reuse a single mapping / type promotion function?

[brandon-b-miller]

Yeah, this is a good point. There should only be one type mapping. Let me see what I can do here.

[brandon-b-miller]

In the latest commit, I've refactored towards an attribute generating function which checks the existing registry of cuda functions for a match in order to return a signature. I think that makes it so that this is the only place in python where the signatures are written by hand. Let me know what you think.

[brandon-b-miller]

@bdice thanks for your review. I've thought a bit about making it so that there's only one centralized place to keep all of the function signatures that could be read directly from c++ or python. Since the shim function library is built before the python is imported, I suppose it should be the source of truth as well. This leads my mind towards some kind of solution which inspects the shim file at import and generates signatures dynamically from it. @gmarkall can cuobjdump or nvdisasm help us here?

[bdice]

Much better to only define these once in Python. I'm okay with redefining in C++ and Python if that is needed. It is probably difficult to do this dynamically. I have one question, otherwise approving.

[brandon-b-miller]

/merge