def __call__(self, ctx, node):
method = getattr(self, 'build_' + node.__class__.__name__, None)
try:
if method is None:
error_msg = f'Unsupported node "{node.__class__.__name__}"'
raise TaichiSyntaxError(error_msg)
return method(ctx, node)
except Exception as e:
if ctx.raised or not isinstance(node, (ast.stmt, ast.expr)):
raise e.with_traceback(None)
ctx.raised = True
e = handle_exception_from_cpp(e)
if not isinstance(e, TaichiCompilationError):
msg = ctx.get_pos_info(node) + traceback.format_exc()
raise TaichiCompilationError(msg) from None
msg = ctx.get_pos_info(node) + str(e)
raise type(e)(msg) from None
def func__(*args):
assert len(args) == len(
self.argument_annotations
), f'{len(self.argument_annotations)} arguments needed but {len(args)} provided'
tmps = []
callbacks = []
has_external_arrays = False
has_torch = has_pytorch()
actual_argument_slot = 0
launch_ctx = t_kernel.make_launch_context()
for i, v in enumerate(args):
needed = self.argument_annotations[i]
if isinstance(needed, template):
continue
provided = type(v)
# Note: do not use sth like "needed == f32". That would be slow.
if id(needed) in primitive_types.real_type_ids:
if not isinstance(v, (float, int)):
raise TaichiRuntimeTypeError.get(
i, needed.to_string(), provided)
launch_ctx.set_arg_float(actual_argument_slot, float(v))
elif id(needed) in primitive_types.integer_type_ids:
if not isinstance(v, int):
raise TaichiRuntimeTypeError.get(
i, needed.to_string(), provided)
launch_ctx.set_arg_int(actual_argument_slot, int(v))
elif isinstance(needed, sparse_matrix_builder):
# Pass only the base pointer of the ti.types.sparse_matrix_builder() argument
launch_ctx.set_arg_int(actual_argument_slot, v._get_addr())
elif isinstance(needed,
ndarray_type.NdarrayType) and isinstance(
v, taichi.lang._ndarray.Ndarray):
has_external_arrays = True
v = v.arr
launch_ctx.set_arg_ndarray(actual_argument_slot, v)
elif isinstance(
needed,
ndarray_type.NdarrayType) and (self.match_ext_arr(v)):
has_external_arrays = True
is_numpy = isinstance(v, np.ndarray)
if is_numpy:
tmp = np.ascontiguousarray(v)
# Purpose: DO NOT GC |tmp|!
tmps.append(tmp)
launch_ctx.set_arg_external_array_with_shape(
actual_argument_slot, int(tmp.ctypes.data),
tmp.nbytes, v.shape)
else:
is_ndarray = False
tmp, torch_callbacks = self.get_torch_callbacks(
v, has_torch, is_ndarray)
callbacks += torch_callbacks
launch_ctx.set_arg_external_array_with_shape(
actual_argument_slot, int(tmp.data_ptr()),
tmp.element_size() * tmp.nelement(), v.shape)
elif isinstance(needed, MatrixType):
if id(needed.dtype) in primitive_types.real_type_ids:
for a in range(needed.n):
for b in range(needed.m):
if not isinstance(v[a, b], (int, float)):
raise TaichiRuntimeTypeError.get(
i, needed.dtype.to_string(),
type(v[a, b]))
launch_ctx.set_arg_float(
actual_argument_slot, float(v[a, b]))
actual_argument_slot += 1
elif id(needed.dtype) in primitive_types.integer_type_ids:
for a in range(needed.n):
for b in range(needed.m):
if not isinstance(v[a, b], int):
raise TaichiRuntimeTypeError.get(
i, needed.dtype.to_string(),
type(v[a, b]))
launch_ctx.set_arg_int(actual_argument_slot,
int(v[a, b]))
actual_argument_slot += 1
else:
raise ValueError(
f'Matrix dtype {needed.dtype} is not integer type or real type.'
)
continue
else:
raise ValueError(
f'Argument type mismatch. Expecting {needed}, got {type(v)}.'
)
actual_argument_slot += 1
# Both the class kernels and the plain-function kernels are unified now.
# In both cases, |self.grad| is another Kernel instance that computes the
# gradient. For class kernels, args[0] is always the kernel owner.
if not self.is_grad and self.runtime.target_tape and not self.runtime.grad_replaced:
self.runtime.target_tape.insert(self, args)
if actual_argument_slot > 8 and (
impl.current_cfg().arch == _ti_core.opengl
or impl.current_cfg().arch == _ti_core.cc):
raise TaichiRuntimeError(
f"The number of elements in kernel arguments is too big! Do not exceed 8 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
)
if actual_argument_slot > 64 and (
(impl.current_cfg().arch != _ti_core.opengl
and impl.current_cfg().arch != _ti_core.cc)):
raise TaichiRuntimeError(
f"The number of elements in kernel arguments is too big! Do not exceed 64 on {_ti_core.arch_name(impl.current_cfg().arch)} backend."
)
try:
t_kernel(launch_ctx)
except Exception as e:
e = handle_exception_from_cpp(e)
raise e from None
ret = None
ret_dt = self.return_type
has_ret = ret_dt is not None
if has_ret or (impl.current_cfg().async_mode
and has_external_arrays):
runtime_ops.sync()
if has_ret:
if id(ret_dt) in primitive_types.integer_type_ids:
ret = t_kernel.get_ret_int(0)
elif id(ret_dt) in primitive_types.real_type_ids:
ret = t_kernel.get_ret_float(0)
elif id(ret_dt.dtype) in primitive_types.integer_type_ids:
it = iter(t_kernel.get_ret_int_tensor(0))
ret = Matrix([[next(it) for _ in range(ret_dt.m)]
for _ in range(ret_dt.n)])
else:
it = iter(t_kernel.get_ret_float_tensor(0))
ret = Matrix([[next(it) for _ in range(ret_dt.m)]
for _ in range(ret_dt.n)])
if callbacks:
for c in callbacks:
c()
return ret