def call(self, env, *args, unroll=None):
if len(args) == 0:
raise TypeError('range expected at least 1 argument, got 0')
elif len(args) == 1:
start, stop, step = Constant(0), args[0], Constant(1)
elif len(args) == 2:
start, stop, step = args[0], args[1], Constant(1)
elif len(args) == 3:
start, stop, step = args
else:
raise TypeError(
f'range expected at most 3 argument, got {len(args)}')
if unroll is not None:
if not all(isinstance(x, Constant)
for x in (start, stop, step, unroll)):
raise TypeError(
'loop unrolling requires constant start, stop, step and '
'unroll value')
unroll = unroll.obj
if not (isinstance(unroll, int) or isinstance(unroll, bool)):
raise TypeError(
'unroll value expected to be of type int, '
f'got {type(unroll).__name__}')
if unroll is False:
unroll = 1
if not (unroll is True or 0 < unroll < 1 << 31):
warnings.warn(
'loop unrolling is ignored as the unroll value is '
'non-positive or greater than INT_MAX')
if isinstance(step, Constant):
step_is_positive = step.obj >= 0
elif step.ctype.dtype.kind == 'u':
step_is_positive = True
else:
step_is_positive = None
stop = Data.init(stop, env)
start = Data.init(start, env)
step = Data.init(step, env)
if start.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if stop.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if step.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if env.mode == 'numpy':
ctype = _cuda_types.Scalar(int)
elif env.mode == 'cuda':
ctype = stop.ctype
else:
assert False
return Range(start, stop, step, ctype, step_is_positive, unroll=unroll)
def _eval_operand(op, args, env):
if is_constants(*args):
pyfunc = _cuda_typerules.get_pyfunc(type(op))
return Constant(pyfunc(*[x.obj for x in args]))
ufunc = _cuda_typerules.get_ufunc(env.mode, type(op))
return _call_ufunc(ufunc, args, None, env)
def _astype_scalar(
x: _T,
ctype: _cuda_types.Scalar,
casting: _CastingType,
env: Environment,
) -> _T:
if isinstance(x, Constant):
assert not isinstance(x, Data)
return Constant(ctype.dtype.type(x.obj))
# # TODO
# if not isinstance(x, Data):
# raise TypeError(f'{x} is not scalar type.')
if not isinstance(x.ctype, _cuda_types.Scalar):
raise TypeError(f'{x.code} is not scalar type.')
from_t = x.ctype.dtype
to_t = ctype.dtype
if from_t == to_t:
return x
# Uses casting rules for scalar values.
if not numpy.can_cast(from_t.type(0), to_t.type(0), casting):
raise TypeError(f"Cannot cast from '{from_t}' to {to_t} "
f"with casting rule {casting}.")
if from_t.kind == 'c' and to_t.kind != 'c':
if to_t.kind != 'b':
warnings.warn(
'Casting complex values to real discards the imaginary part',
numpy.ComplexWarning)
return Data(f'({ctype})({x.code}.real())', ctype)
return Data(f'({ctype})({x.code})', ctype)
def _eval_operand(
op: ast.AST,
args: Sequence[Union[Constant, Data]],
env: Environment,
) -> Union[Constant, Data]:
if is_constants(*args):
pyfunc = _cuda_typerules.get_pyfunc(type(op))
return Constant(pyfunc(*[x.obj for x in args]))
ufunc = _cuda_typerules.get_ufunc(env.mode, type(op))
return _call_ufunc(ufunc, args, None, env)
def call(self, env, mask, var, val_id, *, width=None):
name = self._name
var = Data.init(var, env)
ctype = var.ctype
if ctype.dtype.name not in self._dtypes:
raise TypeError(f'`{name}` does not support {ctype.dtype} input.')
try:
mask = mask.obj
except Exception:
raise TypeError('mask must be an integer')
if runtime.is_hip:
warnings.warn(f'mask {mask} is ignored on HIP', RuntimeWarning)
elif not (0x0 <= mask <= 0xffffffff):
raise ValueError('mask is out of range')
# val_id refers to "delta" for shfl_{up, down}, "srcLane" for shfl, and
# "laneMask" for shfl_xor
if self._op in ('up', 'down'):
val_id_t = _cuda_types.uint32
else:
val_id_t = _cuda_types.int32
val_id = _compile._astype_scalar(val_id, val_id_t, 'same_kind', env)
val_id = Data.init(val_id, env)
if width:
if isinstance(width, Constant):
if width.obj not in (2, 4, 8, 16, 32):
raise ValueError('width needs to be power of 2')
else:
width = Constant(64) if runtime.is_hip else Constant(32)
width = _compile._astype_scalar(width, _cuda_types.int32, 'same_kind',
env)
width = Data.init(width, env)
code = f'{name}({hex(mask)}, {var.code}, {val_id.code}'
code += f', {width.code})'
return Data(code, ctype)
def call(self, env, *args, **kwargs):
if len(args) == 0:
raise TypeError('range expected at least 1 argument, got 0')
elif len(args) == 1:
start, stop, step = Constant(0), args[0], Constant(1)
elif len(args) == 2:
start, stop, step = args[0], args[1], Constant(1)
elif len(args) == 3:
start, stop, step = args
else:
raise TypeError(
f'range expected at most 3 argument, got {len(args)}')
if isinstance(step, Constant):
step_is_positive = step.obj >= 0
elif step.ctype.dtype.kind == 'u':
step_is_positive = True
else:
step_is_positive = None
stop = Data.init(stop, env)
start = Data.init(start, env)
step = Data.init(step, env)
if start.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if stop.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if step.ctype.dtype.kind not in 'iu':
raise TypeError('range supports only for integer type.')
if env.mode == 'numpy':
ctype = _cuda_types.Scalar(int)
elif env.mode == 'cuda':
ctype = stop.ctype
else:
assert False
return Range(start, stop, step, ctype, step_is_positive)
def _indexing(
array: _internal_types.Expr,
index: _internal_types.Expr,
env: Environment,
) -> Union[Data, Constant]:
if isinstance(array, Constant):
if isinstance(index, Constant):
return Constant(array.obj[index.obj])
raise TypeError(
f'{type(array.obj)} is not subscriptable with non-constants.')
array = Data.init(array, env)
if isinstance(array.ctype, _cuda_types.Tuple):
if isinstance(index, Constant):
i = index.obj
t = array.ctype.types[i]
return Data(f'thrust::get<{i}>({array.code})', t)
raise TypeError('Tuple is not subscriptable with non-constants.')
if isinstance(array.ctype, _cuda_types.ArrayBase):
index = Data.init(index, env)
ndim = array.ctype.ndim
if isinstance(index.ctype, _cuda_types.Scalar):
index_dtype = index.ctype.dtype
if ndim != 1:
raise TypeError(
'Scalar indexing is supported only for 1-dim array.')
if index_dtype.kind not in 'ui':
raise TypeError('Array indices must be integers.')
return Data(f'{array.code}[{index.code}]', array.ctype.child_type)
if isinstance(index.ctype, _cuda_types.Tuple):
if ndim != len(index.ctype.types):
raise IndexError(f'The size of index must be {ndim}')
for t in index.ctype.types:
if not isinstance(t, _cuda_types.Scalar):
raise TypeError('Array indices must be scalar.')
if t.dtype.kind not in 'iu':
raise TypeError('Array indices must be integer.')
if ndim == 0:
return Data(f'{array.code}[0]', array.ctype.child_type)
if ndim == 1:
return Data(f'{array.code}[thrust::get<0>({index.code})]',
array.ctype.child_type)
return Data(f'{array.code}._indexing({index.code})',
array.ctype.child_type)
if isinstance(index.ctype, _cuda_types.CArray):
raise TypeError('Advanced indexing is not supported.')
assert False # Never reach.
raise TypeError(f'{array.code} is not subscriptable.')
def _transpile_function_internal(func, attributes, mode, consts, in_types,
ret_type):
consts = dict([(k, Constant(v)) for k, v, in consts.items()])
if not isinstance(func, ast.FunctionDef):
# TODO(asi1024): Support for `ast.ClassDef`.
raise NotImplementedError('Not supported: {}'.format(type(func)))
if len(func.decorator_list) > 0:
if sys.version_info >= (3, 9):
# Code path for Python versions that support `ast.unparse`.
for deco in func.decorator_list:
deco_code = ast.unparse(deco)
if not any(word in deco_code
for word in ['rawkernel', 'vectorize']):
warnings.warn(
f'Decorator {deco_code} may not supported in JIT.',
RuntimeWarning)
arguments = func.args
if arguments.vararg is not None:
raise NotImplementedError('`*args` is not supported currently.')
if len(arguments.kwonlyargs) > 0: # same length with `kw_defaults`.
raise NotImplementedError(
'keyword only arguments are not supported currently .')
if arguments.kwarg is not None:
raise NotImplementedError('`**kwargs` is not supported currently.')
if len(arguments.defaults) > 0:
raise NotImplementedError(
'Default values are not supported currently.')
args = [arg.arg for arg in arguments.args]
if len(args) != len(in_types):
raise TypeError(
f'{func.name}() takes {len(args)} positional arguments '
f'but {len(in_types)} were given.')
params = dict([(x, Data(x, t)) for x, t in zip(args, in_types)])
env = Environment(mode, consts, params, ret_type)
body = _transpile_stmts(func.body, True, env)
params = ', '.join([env[a].ctype.declvar(a) for a in args])
local_vars = [v.ctype.declvar(n) + ';' for n, v in env.locals.items()]
if env.ret_type is None:
env.ret_type = _cuda_types.void
head = f'{attributes} {env.ret_type} {func.name}({params})'
code = CodeBlock(head, local_vars + body)
return str(code), env
def _astype_scalar(x, ctype, casting, env):
if is_constants(x):
return Constant(ctype.dtype.type(x.obj))
from_t = x.ctype.dtype
to_t = ctype.dtype
if from_t == to_t:
return x
# Uses casting rules for scalar values.
if not numpy.can_cast(from_t.type(0), to_t.type(0), casting):
raise TypeError(f"Cannot cast from '{from_t}' to {to_t} "
f"with casting rule {casting}.")
if from_t.kind == 'c' and to_t.kind != 'c':
if to_t.kind != 'b':
warnings.warn(
'Casting complex values to real discards the imaginary part',
numpy.ComplexWarning)
return Data(f'({ctype})({x.code}.real())', ctype)
return Data(f'({ctype})({x.code})', ctype)
def _transpile_expr_internal(expr, env):
if isinstance(expr, ast.BoolOp):
values = [_transpile_expr(e, env) for e in expr.values]
value = values[0]
for rhs in values[1:]:
value = _eval_operand(expr.op, (value, rhs), env)
return value
if isinstance(expr, ast.BinOp):
left = _transpile_expr(expr.left, env)
right = _transpile_expr(expr.right, env)
return _eval_operand(expr.op, (left, right), env)
if isinstance(expr, ast.UnaryOp):
value = _transpile_expr(expr.operand, env)
return _eval_operand(expr.op, (value, ), env)
if isinstance(expr, ast.Lambda):
raise NotImplementedError('Not implemented.')
if isinstance(expr, ast.Compare):
values = [expr.left] + expr.comparators
if len(values) != 2:
raise NotImplementedError(
'Comparison of 3 or more values is not implemented.')
values = [_transpile_expr(e, env) for e in values]
return _eval_operand(expr.ops[0], values, env)
if isinstance(expr, ast.IfExp):
cond = _transpile_expr(expr.test, env)
x = _transpile_expr(expr.body, env)
y = _transpile_expr(expr.orelse, env)
if isinstance(expr, Constant):
return x if expr.obj else y
if cond.ctype.dtype.kind == 'c':
raise TypeError("Complex type value cannot be boolean condition.")
x, y = _infer_type(x, y, env), _infer_type(y, x, env)
if x.ctype.dtype != y.ctype.dtype:
raise TypeError('Type mismatch in conditional expression.: '
f'{x.ctype.dtype} != {y.ctype.dtype}')
cond = _astype_scalar(cond, _cuda_types.bool_, 'unsafe', env)
return Data(f'({cond.code} ? {x.code} : {y.code})', x.ctype)
if isinstance(expr, ast.Call):
func = _transpile_expr(expr.func, env)
args = [_transpile_expr(x, env) for x in expr.args]
kwargs = dict([(kw.arg, _transpile_expr(kw.value, env))
for kw in expr.keywords])
builtin_funcs = _builtin_funcs.builtin_functions_dict
if is_constants(func) and (func.obj in builtin_funcs):
func = builtin_funcs[func.obj]
if isinstance(func, _internal_types.BuiltinFunc):
return func.call(env, *args, **kwargs)
if not is_constants(func):
raise TypeError(f"'{func}' is not callable.")
func = func.obj
if is_constants(*args, *kwargs.values()):
# compile-time function call
args = [x.obj for x in args]
kwargs = dict([(k, v.obj) for k, v in kwargs.items()])
return Constant(func(*args, **kwargs))
if isinstance(func, _kernel.ufunc):
# ufunc call
dtype = kwargs.pop('dtype', Constant(None)).obj
if len(kwargs) > 0:
name = next(iter(kwargs))
raise TypeError(
f"'{name}' is an invalid keyword to ufunc {func.name}")
return _call_ufunc(func, args, dtype, env)
if inspect.isclass(func) and issubclass(func, _typeclasses):
# explicit typecast
if len(args) != 1:
raise TypeError(
f'function takes {func} invalid number of argument')
ctype = _cuda_types.Scalar(func)
return _astype_scalar(args[0], ctype, 'unsafe', env)
if isinstance(func, _interface._JitRawKernel) and func._device:
args = [Data.init(x, env) for x in args]
in_types = tuple([x.ctype for x in args])
fname, return_type = _transpile_func_obj(func._func,
['__device__'], env.mode,
in_types, None,
env.generated)
in_params = ', '.join([x.code for x in args])
return Data(f'{fname}({in_params})', return_type)
raise TypeError(f"Invalid function call '{fname}'.")
if isinstance(expr, ast.Constant):
return Constant(expr.value)
if isinstance(expr, ast.Num):
# Deprecated since py3.8
return Constant(expr.n)
if isinstance(expr, ast.Str):
# Deprecated since py3.8
return Constant(expr.s)
if isinstance(expr, ast.NameConstant):
# Deprecated since py3.8
return Constant(expr.value)
if isinstance(expr, ast.Subscript):
array = _transpile_expr(expr.value, env)
index = _transpile_expr(expr.slice, env)
return _indexing(array, index, env)
if isinstance(expr, ast.Name):
value = env[expr.id]
if value is None:
raise NameError(f'Unbound name: {expr.id}')
return env[expr.id]
if isinstance(expr, ast.Attribute):
value = _transpile_expr(expr.value, env)
if is_constants(value):
return Constant(getattr(value.obj, expr.attr))
if isinstance(value.ctype, _cuda_types.ArrayBase):
if 'ndim' == expr.attr:
return Constant(value.ctype.ndim)
if isinstance(value.ctype, _cuda_types.CArray):
if 'size' == expr.attr:
return Data(f'static_cast<long long>({value.code}.size())',
_cuda_types.Scalar('q'))
if isinstance(value.ctype, _interface._Dim3):
if expr.attr in ('x', 'y', 'z'):
return Data(f'{value.code}.{expr.attr}', _cuda_types.uint32)
# TODO(leofang): support arbitrary Python class methods
if isinstance(value.ctype, _ThreadGroup):
return _internal_types.BuiltinFunc.from_class_method(
value.code, getattr(value.ctype, expr.attr))
raise NotImplementedError('Not implemented: __getattr__')
if isinstance(expr, ast.Tuple):
elts = [_transpile_expr(x, env) for x in expr.elts]
# TODO: Support compile time constants.
elts = [Data.init(x, env) for x in elts]
elts_code = ', '.join([x.code for x in elts])
ctype = _cuda_types.Tuple([x.ctype for x in elts])
return Data(f'thrust::make_tuple({elts_code})', ctype)
if isinstance(expr, ast.Index):
return _transpile_expr(expr.value, env)
raise ValueError('Not supported: type {}'.format(type(expr)))
def _transpile_expr_internal(expr, env):
if isinstance(expr, ast.BoolOp):
values = [_transpile_expr(e, env) for e in expr.values]
value = values[0]
for rhs in values[1:]:
value = _eval_operand(expr.op, (value, rhs), env)
return value
if isinstance(expr, ast.BinOp):
left = _transpile_expr(expr.left, env)
right = _transpile_expr(expr.right, env)
return _eval_operand(expr.op, (left, right), env)
if isinstance(expr, ast.UnaryOp):
value = _transpile_expr(expr.operand, env)
return _eval_operand(expr.op, (value, ), env)
if isinstance(expr, ast.Lambda):
raise NotImplementedError('Not implemented.')
if isinstance(expr, ast.Compare):
values = [expr.left] + expr.comparators
if len(values) != 2:
raise NotImplementedError(
'Comparison of 3 or more values is not implemented.')
values = [_transpile_expr(e, env) for e in values]
return _eval_operand(expr.ops[0], values, env)
if isinstance(expr, ast.IfExp):
cond = _transpile_expr(expr.test, env)
x = _transpile_expr(expr.body, env)
y = _transpile_expr(expr.orelse, env)
if isinstance(expr, Constant):
return x if expr.obj else y
if cond.ctype.dtype.kind == 'c':
raise NotImplementedError('')
x = Data.init(x, env)
y = Data.init(y, env)
if x.ctype.dtype != y.ctype.dtype:
raise TypeError('Type mismatch in conditional expression.: '
f'{x.ctype.dtype} != {y.ctype.dtype}')
cond = _astype_scalar(cond, _cuda_types.bool_, 'unsafe', env)
return Data(f'({cond.code} ? {x.code} : {y.code})', x.ctype)
if isinstance(expr, ast.Call):
func = _transpile_expr(expr.func, env)
args = [_transpile_expr(x, env) for x in expr.args]
kwargs = dict([(kw.arg, _transpile_expr(kw.value, env))
for kw in expr.keywords])
builtin_funcs = _builtin_funcs.builtin_functions_dict
if is_constants(func) and (func.obj in builtin_funcs):
func = builtin_funcs[func.obj]
if isinstance(func, _internal_types.BuiltinFunc):
return func.call(env, *args, **kwargs)
if not is_constants(func):
raise NotImplementedError(
'device function call is not implemented.')
func = func.obj
if is_constants(*args, *kwargs.values()):
# compile-time function call
args = [x.obj for x in args]
kwargs = dict([(k, v.obj) for k, v in kwargs.items()])
return Constant(func(*args, **kwargs))
if isinstance(func, _kernel.ufunc):
# ufunc call
dtype = kwargs.pop('dtype', Constant(None)).obj
if len(kwargs) > 0:
name = next(iter(kwargs))
raise TypeError(
f"'{name}' is an invalid keyword to ufunc {func.name}")
return _call_ufunc(func, args, dtype, env)
if inspect.isclass(func) and issubclass(func, _typeclasses):
# explicit typecast
if len(args) != 1:
raise TypeError(
f'function takes {func} invalid number of argument')
ctype = _cuda_types.Scalar(func)
return _astype_scalar(args[0], ctype, 'unsafe', env)
raise NotImplementedError(
f'function call of `{func.__name__}` is not implemented')
if isinstance(expr, ast.Constant):
return Constant(expr.value)
if isinstance(expr, ast.Num):
# Deprecated since py3.8
return Constant(expr.n)
if isinstance(expr, ast.Str):
# Deprecated since py3.8
return Constant(expr.s)
if isinstance(expr, ast.NameConstant):
# Deprecated since py3.8
return Constant(expr.value)
if isinstance(expr, ast.Subscript):
value = _transpile_expr(expr.value, env)
index = _transpile_expr(expr.slice, env)
if is_constants(value):
if is_constants(index):
return Constant(value.obj[index.obj])
raise TypeError(
f'{type(value.obj)} is not subscriptable with non-constants.')
value = Data.init(value, env)
if isinstance(value.ctype, _cuda_types.Tuple):
raise NotImplementedError
if isinstance(value.ctype, _cuda_types.ArrayBase):
index = Data.init(index, env)
ndim = value.ctype.ndim
if isinstance(index.ctype, _cuda_types.Scalar):
index_dtype = index.ctype.dtype
if ndim != 1:
raise TypeError(
'Scalar indexing is supported only for 1-dim array.')
if index_dtype.kind not in 'ui':
raise TypeError('Array indices must be integers.')
return Data(f'{value.code}[{index.code}]',
value.ctype.child_type)
if isinstance(index.ctype, _cuda_types.Tuple):
if ndim != len(index.ctype.types):
raise IndexError(f'The size of index must be {ndim}')
for t in index.ctype.types:
if not isinstance(t, _cuda_types.Scalar):
raise TypeError('Array indices must be scalar.')
if t.dtype.kind not in 'iu':
raise TypeError('Array indices must be integer.')
if ndim == 0:
return Data(f'{value.code}[0]', value.ctype.child_type)
if ndim == 1:
return Data(f'{value.code}[thrust::get<0>({index.code})]',
value.ctype.child_type)
return Data(f'{value.code}._indexing({index.code})',
value.ctype.child_type)
if isinstance(index.ctype, _cuda_types.Array):
raise TypeError('Advanced indexing is not supported.')
assert False # Never reach.
raise TypeError(f'{value.code} is not subscriptable.')
if isinstance(expr, ast.Name):
value = env[expr.id]
if value is None:
raise NameError(f'Unbound name: {expr.id}')
return env[expr.id]
if isinstance(expr, ast.Attribute):
value = _transpile_expr(expr.value, env)
if is_constants(value):
return Constant(getattr(value.obj, expr.attr))
raise NotImplementedError('Not implemented: __getattr__')
if isinstance(expr, ast.Tuple):
elts = [_transpile_expr(x, env) for x in expr.elts]
# TODO: Support compile time constants.
elts = [Data.init(x, env) for x in elts]
elts_code = ', '.join([x.code for x in elts])
ctype = _cuda_types.Tuple([x.ctype for x in elts])
return Data(f'thrust::make_tuple({elts_code})', ctype)
if isinstance(expr, ast.Index):
return _transpile_expr(expr.value, env)
raise ValueError('Not supported: type {}'.format(type(expr)))