def match_ext_arr(v):
has_array = isinstance(v, np.ndarray)
if not has_array and has_pytorch():
has_array = isinstance(v, torch.Tensor)
if not has_array and has_paddle():
has_array = isinstance(v, paddle.Tensor)
return has_array
import numpy as np
import pytest
from taichi.lang import impl
from taichi.lang.util import has_paddle
import taichi as ti
from tests import test_utils
if has_paddle():
import paddle
@pytest.mark.skipif(not has_paddle(), reason='Paddle not installed.')
@test_utils.test(arch=[ti.cpu, ti.cuda])
def test_io_devices():
n = 32
x = ti.field(dtype=ti.i32, shape=n)
@ti.kernel
def load(y: ti.types.ndarray()):
for i in x:
x[i] = y[i] + 10
@ti.kernel
def inc():
for i in x:
x[i] += i
@ti.kernel
def store(y: ti.types.ndarray()):
for i in x:
def func__(*args):
assert len(args) == len(
self.arguments
), f'{len(self.arguments)} arguments needed but {len(args)} provided'
tmps = []
callbacks = []
has_external_arrays = False
has_torch = has_pytorch()
has_pp = has_paddle()
actual_argument_slot = 0
launch_ctx = t_kernel.make_launch_context()
for i, v in enumerate(args):
needed = self.arguments[i].annotation
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,
texture_type.TextureType) and isinstance(
v, taichi.lang._texture.Texture):
has_external_arrays = True
v = v.tex
launch_ctx.set_arg_texture(actual_argument_slot, v)
elif isinstance(needed,
texture_type.RWTextureType) and isinstance(
v, taichi.lang._texture.Texture):
has_external_arrays = True
v = v.tex
launch_ctx.set_arg_rw_texture(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)
is_torch = isinstance(v,
torch.Tensor) if has_torch else False
# Element shapes are already spcialized in Taichi codegen.
# The shape information for element dims are no longer needed.
# Therefore we strip the element shapes from the shape vector,
# so that it only holds "real" array shapes.
is_soa = needed.layout == Layout.SOA
array_shape = v.shape
element_dim = needed.element_dim
if element_dim:
array_shape = v.shape[
element_dim:] if is_soa else v.shape[:-element_dim]
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, array_shape)
elif is_torch:
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(), array_shape)
else:
# For now, paddle.fluid.core.Tensor._ptr() is only available on develop branch
tmp, paddle_callbacks = self.get_paddle_callbacks(
v, has_pp)
callbacks += paddle_callbacks
launch_ctx.set_arg_external_array_with_shape(
actual_argument_slot, int(tmp._ptr()),
v.element_size() * v.size, array_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 self.autodiff_mode == AutodiffMode.NONE 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:
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
TaichiRuntimeTypeError, TaichiSyntaxError,
handle_exception_from_cpp)
from taichi.lang.expr import Expr
from taichi.lang.kernel_arguments import KernelArgument
from taichi.lang.matrix import Matrix, MatrixType
from taichi.lang.shell import _shell_pop_print, oinspect
from taichi.lang.util import has_paddle, has_pytorch, to_taichi_type
from taichi.types import (ndarray_type, primitive_types, sparse_matrix_builder,
template, texture_type)
from taichi import _logging
if has_pytorch():
import torch
if has_paddle():
import paddle
def func(fn, is_real_function=False):
"""Marks a function as callable in Taichi-scope.
This decorator transforms a Python function into a Taichi one. Taichi
will JIT compile it into native instructions.
Args:
fn (Callable): The Python function to be decorated
is_real_function (bool): Whether the function is a real function
Returns:
Callable: The decorated function