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):