NVIDIA / NVIDIA/cccl

[EXTERNAL] Numba pointer arithmetic

Open
#2,861 0 comments 1 reaction 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.