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
def match_ext_arr(self, v, needed):
needs_array = isinstance(
needed, np.ndarray) or needed == np.ndarray or isinstance(
needed, ext_arr)
has_array = isinstance(v, np.ndarray)
if not has_array and util.has_pytorch():
import torch
has_array = isinstance(v, torch.Tensor)
return has_array and needs_array
def __init__(self, dtype, shape):
self.host_accessor = None
if impl.current_cfg().ndarray_use_torch:
assert has_pytorch(
), "PyTorch must be available if you want to create a Taichi ndarray with PyTorch as its underlying storage."
self.arr = torch.zeros(shape,
dtype=to_pytorch_type(cook_dtype(dtype)))
if impl.current_cfg().arch == _ti_core.Arch.cuda:
self.arr = self.arr.cuda()
else:
self.arr = _ti_core.Ndarray(impl.get_runtime().prog,
cook_dtype(dtype), shape)
def __init__(self, dtype, shape):
if isinstance(shape, numbers.Number):
shape = (shape, )
assert has_pytorch(
), "PyTorch must be available if you want to create a Taichi ndarray."
import torch
if impl.current_cfg().arch == _ti_core.Arch.cuda:
device = 'cuda:0'
else:
device = 'cpu'
self.arr = torch.empty(shape,
dtype=to_pytorch_type(dtype),
device=device)
import numpy as np
import pytest
from taichi.lang import impl
from taichi.lang.util import has_pytorch
import taichi as ti
from tests import test_utils
if has_pytorch():
import torch
@pytest.mark.skipif(not has_pytorch(), reason='Pytorch not installed.')
@test_utils.test(exclude=[ti.opengl, ti.vulkan])
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 match_ext_arr(self, v):
has_array = isinstance(v, np.ndarray)
if not has_array and util.has_pytorch():
has_array = isinstance(v, torch.Tensor)
return has_array
def func__(*args):
assert len(args) == len(
self.argument_annotations
), '{} arguments needed but {} provided'.format(
len(self.argument_annotations), len(args))
tmps = []
callbacks = []
has_external_arrays = False
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 KernelArgError(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 KernelArgError(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.linalg.sparse_matrix_builder() argument
launch_ctx.set_arg_int(actual_argument_slot, v.get_addr())
elif isinstance(needed, any_arr) and (
self.match_ext_arr(v)
or isinstance(v, taichi.lang._ndarray.Ndarray)):
is_ndarray = False
if isinstance(v, taichi.lang._ndarray.Ndarray):
v = v.arr
is_ndarray = True
has_external_arrays = True
ndarray_use_torch = self.runtime.prog.config.ndarray_use_torch
has_torch = util.has_pytorch()
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(
actual_argument_slot, int(tmp.ctypes.data),
tmp.nbytes)
elif is_ndarray and not ndarray_use_torch:
# Use ndarray's own memory allocator
tmp = v
launch_ctx.set_arg_external_array(
actual_argument_slot, int(tmp.data_ptr()),
tmp.element_size() * tmp.nelement())
else:
def get_call_back(u, v):
def call_back():
u.copy_(v)
return call_back
assert util.has_pytorch()
assert isinstance(v, torch.Tensor)
tmp = v
taichi_arch = self.runtime.prog.config.arch
if str(v.device).startswith('cuda'):
# External tensor on cuda
if taichi_arch != _ti_core.Arch.cuda:
# copy data back to cpu
host_v = v.to(device='cpu', copy=True)
tmp = host_v
callbacks.append(get_call_back(v, host_v))
else:
# External tensor on cpu
if taichi_arch == _ti_core.Arch.cuda:
gpu_v = v.cuda()
tmp = gpu_v
callbacks.append(get_call_back(v, gpu_v))
launch_ctx.set_arg_external_array(
actual_argument_slot, int(tmp.data_ptr()),
tmp.element_size() * tmp.nelement())
shape = v.shape
max_num_indices = _ti_core.get_max_num_indices()
assert len(
shape
) <= max_num_indices, "External array cannot have > {} indices".format(
max_num_indices)
for ii, s in enumerate(shape):
launch_ctx.set_extra_arg_int(actual_argument_slot, ii,
s)
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)
t_kernel(launch_ctx)
ret = None
ret_dt = self.return_type
has_ret = ret_dt is not None
if has_external_arrays or has_ret:
ti.sync()
if has_ret:
if id(ret_dt) in primitive_types.integer_type_ids:
ret = t_kernel.get_ret_int(0)
else:
ret = t_kernel.get_ret_float(0)
if callbacks:
for c in callbacks:
c()
return ret
import taichi.lang
from taichi.core.util import ti_core as _ti_core
from taichi.lang import impl, util
from taichi.lang.ast.checkers import KernelSimplicityASTChecker
from taichi.lang.ast.transformer import ASTTransformerTotal
from taichi.lang.enums import Layout
from taichi.lang.exception import TaichiSyntaxError
from taichi.lang.shell import _shell_pop_print, oinspect
from taichi.lang.util import to_taichi_type
from taichi.linalg.sparse_matrix import sparse_matrix_builder
from taichi.misc.util import obsolete
from taichi.type import any_arr, primitive_types, template
import taichi as ti
if util.has_pytorch():
import torch
def func(fn):
"""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
Returns:
Callable: The decorated function
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()
ndarray_use_torch = impl.get_runtime().ndarray_use_torch
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(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(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.linalg.sparse_matrix_builder() argument
launch_ctx.set_arg_int(actual_argument_slot, v._get_addr())
elif isinstance(needed, any_arr) and isinstance(
v, taichi.lang._ndarray.Ndarray):
has_external_arrays = True
v = v.arr
if ndarray_use_torch:
is_ndarray = True
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)
else:
launch_ctx.set_arg_ndarray(actual_argument_slot, v)
elif isinstance(needed, any_arr) 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(
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(
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)
t_kernel(launch_ctx)
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)
else:
ret = t_kernel.get_ret_float(0)
if callbacks:
for c in callbacks:
c()
return ret
from taichi.lang import impl, runtime_ops
from taichi.lang.ast import (ASTTransformerContext, KernelSimplicityASTChecker,
transform_tree)
from taichi.lang.enums import Layout
from taichi.lang.exception import (TaichiCompilationError,
TaichiRuntimeTypeError, TaichiSyntaxError)
from taichi.lang.expr import Expr
from taichi.lang.matrix import MatrixType
from taichi.lang.shell import _shell_pop_print, oinspect
from taichi.lang.util import has_pytorch, to_taichi_type
from taichi.linalg.sparse_matrix import sparse_matrix_builder
from taichi.types import any_arr, primitive_types, template
from taichi import _logging
if has_pytorch():
import torch
def func(fn):
"""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
Returns:
Callable: The decorated function
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
