awkward What is the state of Awkward Arrays in Numba's CUDA target?

Way back in 2020, @gmarkall said that it should be possible to add custom types to Numbafied functions in the CUDA context, and then at an Nvidia meeting in 2022, demonstrated that it is possible.

Does this work for us out of the box? What would it take to get it to work? It is definitely a goal to have Awkward Arrays work in Numba's CUDA target (and a little embarrassing that it's been on the back burner this long).

Sep 01 '22 13:09 jpivarski

When we last went though the checklist: https://github.com/gmarkall/awkward-checklist - I think we concluded that all the required support in Numba was in place, and it would mainly be a case of bug fixes / small extensions as necessary to implement this for Awkward. That said, the CUDA target is now more capable than it was back then so it might be worth us revisiting the checklist to see how we'd approach meeting all the requirements today.

Sep 06 '22 13:09 gmarkall

That said, the CUDA target is now more capable than it was back then

In particular, the PR that brings support for the high-level extension API to the CUDA target is: https://github.com/numba/numba/pull/8271 - this is the one that was brought up in the meeting, and should make it quicker and easier to write extensions - we'll prioritise this as it would be good to make it available and if possible get some feedback / external use of it.

Sep 06 '22 16:09 gmarkall

Thanks, @gmarkall! I've been thinking for a while that the onus of testing lies on our side, and now @ianna is ready to do it. The links you provide help, too.

Sep 06 '22 18:09 jpivarski

Survey of Numba functions used (updated for v2) Please, see the v1 list in this demo.

decorators
---------------------------------------------------------
  108 numba.extending.lower_builtin
   51 numba.core.typing.templates.bound_function
   45 numba.extending.typeof_impl.register
   18 numba.core.typing.templates.infer_global
   12 numba.extending.register_model
   12 numba.extending.unbox
    9 numba.core.imputils.lower_constant
    9 numba.extending.box
    9 numba.core.typing.templates.infer_getattr
    6 numba.extending.lower_getattr_generic
    6 numba.extending.overload
    3 numba.core.typing.templates.infer
    3 numba.core.datamodel.registry.register_default
    3 numba.core.imputils.iternext_impl
    3 numba.extending.type_callable

classes
---------------------------------------------------------
    8 numba.core.typing.templates.AbstractTemplate
    5 numba.core.datamodel.models.StructModel
    4 numba.extending.NativeValue
    3 numba.core.typing.templates.AttributeTemplate
    2 numba.core.typing.Context
    1 numba.core.typing.templates.Signature

functions
---------------------------------------------------------
    4 numba.core.cgutils.alloca_once_value
    4 numba.core.cgutils.is_not_null
    4 numba.extending.lower_builtin
    3 numba.core.cgutils.get_null_value
    3 numba.typeof
    2 numba.core.cgutils.pointer_add
    2 numba.from_dtype
    1 numba.core.typing.ctypes_utils.make_function_type
    1 numba.core.cgutils.alloca_once
    1 numba.core.cgutils.as_bool_bit
    1 numba.core.cgutils.for_range
    1 numba.core.imputils.impl_ret_new_ref
    1 numba.core.cgutils.increment_index

builder.*
---------------------------------------------------------
   63 builder.add
   48 builder.icmp_signed
   36 builder.load
   30 builder.if_then
   24 builder.store
   15 builder.sub
   12 builder.if_else
   12 builder.inttoptr
    9 builder.mul
    6 builder.zext
    6 builder.and_
    3 builder.sext
    3 builder.trunc
    3 builder.or_
    3 builder.sdiv
    3 builder.srem
    3 builder.lshr
    3 builder.shl
    3 builder._layout
    3 builder._layout._ptr
    3 builder.bitcast
    3 builder.icmp_unsigned
    3 builder.sitofp
    3 builder.uitofp
    3 builder.fpext
    3 builder.fptrunc

context.*
---------------------------------------------------------
  192 context.make_helper
  105 context.get_constant
   54 context.get_dummy_value
   42 context.get_value_type
   12 context.add_dynamic_addr
    9 context.enable_nrt
    6 context.call_conv.return_user_exc
    6 context.unify_types
    6 context.nrt.decref
    6 context.make_tuple
    3 context.get_function_pointer_type
    3 context.call_function_pointer
    3 context.make_constant_array
    3 context.nrt.incref

pyapi.*
---------------------------------------------------------
   81 c.pyapi.decref
   33 c.pyapi.object_getattr_string
   24 c.pyapi.unserialize
   24 c.pyapi.serialize_object
   12 c.pyapi.err_occurred
   12 c.pyapi.number_as_ssize_t
   12 c.pyapi.long_from_ssize_t
    9 c.pyapi.call_method
    9 c.pyapi.call_function_objargs
    6 c.pyapi.to_native_value
    3 c.pyapi.long_as_voidptr
    3 c.pyapi.incref

types
---------------------------------------------------------
   54 numba.intp
   48 numba.types.Array
   25 numba.types.none
   17 numba.int64
   17 numba.types.Integer
   16 numba.uint8
   13 numba.int8
   12 numba.types.Type
   10 numba.int32
   10 numba.types.Boolean
   10 numba.types.StringLiteral
    9 numba.none
    9 numba.uint32
    6 numba.types.float64
    6 numba.types.Float
    4 numba.types.optional
    4 numba.types.voidptr
    4 numba.types.Complex
    3 numba.types.CPointer
    3 numba.types.Optional
    2 numba.types.NPDatetime
    2 numba.types.NPTimedelta
    2 numba.types.pyobject
    2 numba.types.UnicodeType
    2 numba.types.NoneType
    2 numba.types.Number
    2 numba.types.UniTuple
    2 numba.types.intp
    1 numba.int16
    1 numba.types.uintp
    1 numba.types.IterableType
    1 numba.types.Sized
    1 numba.types.SliceType
    1 numba.types.slice2_type
    1 numba.types.common.SimpleIteratorType
    1 numba.types.EphemeralPointer
    1 numba.types.Literal
    1 numba.types.BoundFunction
    1 numba.types.boolean
    1 numba.types.int8
    1 numba.types.int16
    1 numba.types.int32
    1 numba.types.int64
    1 numba.types.uint8
    1 numba.types.uint16
    1 numba.types.uint32
    1 numba.types.uint64
    1 numba.types.float32
    1 numba.types.complex64
    1 numba.types.complex128

other
---------------------------------------------------------
  108 numba.intp (<class 'NoneType'>)
   57 numba.arrayview.wrap (<class 'NoneType'>)
   50 numba.types.none (<class 'NoneType'>)
   34 numba.types.Integer (<class 'NoneType'>)
   27 numba.arrayview.tonumbatype (<class 'NoneType'>)
   20 numba.types.Boolean (<class 'NoneType'>)
   20 numba.types.StringLiteral (<class 'NoneType'>)
   18 numba.int64 (<class 'NoneType'>)
   18 numba.arrayview.RecordViewType (<class 'NoneType'>)
   18 numba.layout.castint (<class 'NoneType'>)
   16 numba.types.Array (<class 'NoneType'>)
   16 numba.uint8 (<class 'NoneType'>)
   16 numba.core.typing.templates.AbstractTemplate (<class 'NoneType'>)
   12 numba.types.float64 (<class 'NoneType'>)
   12 numba.types.Float (<class 'NoneType'>)
   12 numba.arrayview.ArrayViewType (<class 'NoneType'>)
   10 numba.int8 (<class 'NoneType'>)
   10 numba.core.datamodel.models.StructModel (<class 'NoneType'>)
    9 numba.layout.getat (<class 'NoneType'>)
    8 numba.types.Type (<class 'NoneType'>)
    8 numba.core.cgutils.alloca_once_value (<class 'NoneType'>)
    8 numba.types.optional (<class 'NoneType'>)
    8 numba.core.cgutils.true_bit (<class 'NoneType'>)
    8 numba.types.voidptr (<class 'NoneType'>)
    8 numba.core.cgutils.is_not_null (<class 'NoneType'>)
    8 numba.extending.NativeValue (<class 'NoneType'>)
    8 numba.types.Complex (<class 'NoneType'>)
    8 numba.extending.lower_builtin (<class 'NoneType'>)
    6 numba.types.CPointer (<class 'NoneType'>)
    6 numba.core.cgutils.false_bit (<class 'NoneType'>)
    6 numba.core.cgutils.get_null_value (<class 'NoneType'>)
    6 numba.core.typing.templates.AttributeTemplate (<class 'NoneType'>)
    6 numba.types.Optional (<class 'NoneType'>)
    6 numba.typeof (<class 'NoneType'>)
    6 numba.layout.NumpyArrayType (<class 'NoneType'>)
    4 numba.int32 (<class 'NoneType'>)
    4 numba.types.NPDatetime (<class 'NoneType'>)
    4 numba.types.NPTimedelta (<class 'NoneType'>)
    4 numba.core.cgutils.pointer_add (<class 'NoneType'>)
    4 numba.from_dtype (<class 'NoneType'>)
    4 numba.core.typing.Context (<class 'NoneType'>)
    4 numba.core.typing (<class 'NoneType'>)
    4 numba.core.typing.ctypes_utils (<class 'NoneType'>)
    4 numba.types.pyobject (<class 'NoneType'>)
    4 numba.types.UnicodeType (<class 'NoneType'>)
    4 numba.types.NoneType (<class 'NoneType'>)
    4 numba.types.Number (<class 'NoneType'>)
    4 numba.types.UniTuple (<class 'NoneType'>)
    4 numba.types.intp (<class 'NoneType'>)
    4 numba.core.cgutils.true_bit (<class 'llvmlite.ir.values.Constant'>)
    3 numba.arrayview.repr_behavior (<class 'NoneType'>)
    3 numba.layout.NumpyArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.RegularArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.ListArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.IndexedArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.IndexedOptionArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.ByteMaskedArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.BitMaskedArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.UnmaskedArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.RecordArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.UnionArrayType.from_form (<class 'NoneType'>)
    3 numba.layout.RecordArrayType (<class 'NoneType'>)
    3 numba.layout.posat (<class 'NoneType'>)
    3 numba.layout.type_bitwidth (<class 'NoneType'>)
    3 numba.np.arrayobj.make_array (<class 'NoneType'>)
    3 numba.np.arrayobj.populate_array (<class 'NoneType'>)
    3 numba.core.cgutils.false_bit (<class 'llvmlite.ir.values.Constant'>)
    2 numba.none (<class 'NoneType'>)
    2 numba.uint32 (<class 'NoneType'>)
    2 numba.int16 (<class 'NoneType'>)
    2 numba.intp.bitwidth (<class 'NoneType'>)
    2 numba.types.uintp (<class 'NoneType'>)
    2 numba.core.typing.ctypes_utils.make_function_type (<class 'NoneType'>)
    2 numba.core.cgutils.alloca_once (<class 'NoneType'>)
    2 numba.core.cgutils.as_bool_bit (<class 'NoneType'>)
    2 numba.core.cgutils.for_range (<class 'NoneType'>)
    2 numba.types.IterableType (<class 'NoneType'>)
    2 numba.types.Sized (<class 'NoneType'>)
    2 numba.types.SliceType (<class 'NoneType'>)
    2 numba.types.slice2_type (<class 'NoneType'>)
    2 numba.types.common.SimpleIteratorType (<class 'NoneType'>)
    2 numba.types.EphemeralPointer (<class 'NoneType'>)
    2 numba.core.imputils.impl_ret_new_ref (<class 'NoneType'>)
    2 numba.core.imputils.RefType.BORROWED (<class 'NoneType'>)
    2 numba.core.cgutils.increment_index (<class 'NoneType'>)
    2 numba.core.typing.templates.Signature (<class 'NoneType'>)
    2 numba.types.Literal (<class 'NoneType'>)
    2 numba.types.BoundFunction (<class 'NoneType'>)
    2 numba.types.boolean (<class 'NoneType'>)
    2 numba.types.int8 (<class 'NoneType'>)
    2 numba.types.int16 (<class 'NoneType'>)
    2 numba.types.int32 (<class 'NoneType'>)
    2 numba.types.int64 (<class 'NoneType'>)
    2 numba.types.uint8 (<class 'NoneType'>)
    2 numba.types.uint16 (<class 'NoneType'>)
    2 numba.types.uint32 (<class 'NoneType'>)
    2 numba.types.uint64 (<class 'NoneType'>)
    2 numba.types.float32 (<class 'NoneType'>)
    2 numba.types.complex64 (<class 'NoneType'>)
    2 numba.types.complex128 (<class 'NoneType'>)
    1 numba.intp.bitwidth (<class 'int'>)
    1 numba.core.imputils.RefType.BORROWED (<enum 'RefType'>)

Sep 07 '22 09:09 ianna

@ianna, there may be a use-case for this feature, which bumps it up in priority and would give us an example problem to use as a test: https://numba.discourse.group/t/tuple-of-cupy-arrays-numba-cuda/1589/13

Oct 17 '22 17:10 jpivarski