[EXTERNAL] Numba pointer arithmetic
- Dominant language
- C++
- Stars
- 2.5k
- Forks
- 486
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 295
Description
In support of cuda.parallel
@gmarkall already provided a sketch of an implementation (below).
@gevtushenko for awareness.
________
Capturing from chat messages:
@rwgk:
I'm running into this error (`numba.cuda.compile()`):
```
No implementation of function Function() found for signature:
>>> add(int32*, uint64)
```
Would it make sense to implement overloads for pointer arithmetic? How involved would that be?
A reproducer with some context:
```
import numba.cuda
from numba.types import CPointer, int32, uint64, void
def pointer_advance(this, distance):
this[0] = this[0] + distance
def compile_pointer_pointer_int32(): # This is what it really is.
thisty = CPointer(CPointer(int32))
numba.cuda.compile(pointer_advance, sig=void(thisty, uint64))
def compile_reinterpret_as_pointer_uint64(): # Workaround.
thisty = CPointer(uint64)
numba.cuda.compile(pointer_advance, sig=void(thisty, uint64))
# Disadvantage: pointer_advance() needs to be called with distance * sizeof(int32)
compile_reinterpret_as_pointer_uint64() # The workaround compiles.
compile_pointer_pointer_int32()
# numba.core.errors.TypingError: Failed in cuda mode pipeline (step: nopython frontend)
# No implementation of function Function() found for signature:
# >>> add(int32*, uint64)
```
@gmarkall:
I think pointer arithmetic would be a really helpful thing to implement, especially for people working on CUDA Python / C++ interop. A sketch of the implementation for this particular use case:
```
import numba.cuda
import operator
from numba.types import CPointer, int32, uint64, void, Integer
from numba.core.extending import overload, intrinsic
from llvmlite import ir
def sizeof_pointee(context, ptr):
size = context.get_abi_sizeof(ptr.type.pointee) # EDIT/FIX: .pointee was missing in the OP
return ir.Constant(ir.IntType(64), size)
@intrinsic
def pointer_add_intrinsic(context, ptr, offset):
def codegen(context, builder, sig, args):
ptr, index = args
base = builder.ptrtoint(ptr, ir.IntType(64))
offset = builder.mul(index, sizeof_pointee(context, ptr))
result = builder.add(base, offset)
return builder.inttoptr(result, ptr.type)
return ptr(ptr, offset), codegen
@overload(operator.add)
def pointer_add(ptr, offset):
if not isinstance(ptr, CPointer) or not isinstance(offset, Integer):
return
def impl(ptr, offset):
return pointer_add_intrinsic(ptr, offset)
return impl
def pointer_advance(this, distance):
this[0] = this[0] + distance
def compile_pointer_pointer_int32():
thisty = CPointer(CPointer(int32))
return numba.cuda.compile(pointer_advance, sig=void(thisty, uint64))
ptx, resty = compile_pointer_pointer_int32()
print(ptx)
```
It would need the other operators (sub, mul, ...) adding, and some tests, and maybe a note adding that pointer arithmetic is possible to the docs. I think that covers the necessary parts of an implementation
Contributor guide
Assessment
This issue has not been assessed yet.