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