from numba.core import types
from numba.core.typing.npydecl import (parse_dtype, parse_shape,
register_number_classes)
from numba.core.typing.templates import (AttributeTemplate, ConcreteTemplate,
AbstractTemplate, CallableTemplate,
signature, Registry)
from numba.cuda.types import dim3, grid_group
from numba import cuda
registry = Registry()
register = registry.register
register_attr = registry.register_attr
register_global = registry.register_global
register_number_classes(register_global)
class GridFunction(CallableTemplate):
def generic(self):
def typer(ndim):
val = ndim.literal_value
if val == 1:
restype = types.int32
elif val in (2, 3):
restype = types.UniTuple(types.int32, val)
else:
raise ValueError('argument can only be 1, 2, 3')
return signature(restype, types.int32)
return typer
from numba.core.typing.templates import (
AttributeTemplate,
ConcreteTemplate,
AbstractTemplate,
MacroTemplate,
signature,
Registry,
)
from numba import cuda
registry = Registry()
intrinsic = registry.register
intrinsic_attr = registry.register_attr
intrinsic_global = registry.register_global
register_number_classes(intrinsic_global)
class Cuda_grid(MacroTemplate):
key = cuda.grid
class Cuda_gridsize(MacroTemplate):
key = cuda.gridsize
class Cuda_threadIdx_x(MacroTemplate):
key = cuda.threadIdx.x
class Cuda_threadIdx_y(MacroTemplate):