def ballot(predicate):
return expr.Expr(
_ti_core.insert_internal_func_call("cuda_ballot_i32",
expr.make_expr_group(predicate),
False))
def asm(source, inputs=[], outputs=[]):
_ti_core.insert_external_func_call(0, source, make_expr_group(inputs),
make_expr_group(outputs))
def barrier():
return expr.Expr(
_ti_core.insert_internal_func_call("subgroupBarrier",
expr.make_expr_group(), False))
def invocation_id():
return expr.Expr(_ti_core.insert_internal_func_call(
"subgroupInvocationId", expr.make_expr_group(), False),
dtype=i32)
def external_func_call_wrapper(args=[], outputs=[]):
func_addr = ctypes.cast(self.so.__getattr__(item),
ctypes.c_void_p).value
_ti_core.insert_external_func_call(func_addr, '', '', '',
make_expr_group(args),
make_expr_group(outputs))
def elect():
return expr.Expr(
_ti_core.insert_internal_func_call("subgroupElect",
expr.make_expr_group(), False))
def make_index_expr(_var, _indices):
return Expr(_ti_core.make_index_expr(_var, make_expr_group(*_indices)))
def length(l, indices):
return Expr(_ti_core.insert_len(l.snode.ptr, make_expr_group(indices)))
def call_internal(name, *args):
return expr_init(
_ti_core.insert_internal_func_call(name, make_expr_group(args)))
def make_stride_expr(_var, _indices, shape, stride):
return Expr(
_ti_core.make_stride_expr(_var, make_expr_group(*_indices), shape,
stride))
def bitcode_func_call_wrapper(*args):
impl.get_runtime().prog.current_ast_builder(
).insert_external_func_call(0, '', self.bc, item,
make_expr_group(args),
make_expr_group([]))
def __init__(self,
n=1,
m=1,
dt=None,
shape=None,
offset=None,
empty=False,
layout=Layout.AOS,
needs_grad=False,
keep_raw=False,
disable_local_tensor=False,
rows=None,
cols=None):
self.local_tensor_proxy = None
self.any_array_access = None
self.grad = None
# construct from rows or cols (deprecated)
if rows is not None or cols is not None:
warning(
f"ti.Matrix(rows=[...]) or ti.Matrix(cols=[...]) is deprecated, use ti.Matrix.rows([...]) or ti.Matrix.cols([...]) instead.",
DeprecationWarning,
stacklevel=2)
if rows is not None and cols is not None:
raise Exception("cannot specify both rows and columns")
self.dt = dt
mat = Matrix.cols(cols) if cols is not None else Matrix.rows(rows)
self.n = mat.n
self.m = mat.m
self.entries = mat.entries
return
elif empty == True:
warning(
f"ti.Matrix(n, m, empty=True) is deprecated, use ti.Matrix.empty(n, m) instead",
DeprecationWarning,
stacklevel=2)
self.dt = dt
self.entries = [[None] * m for _ in range(n)]
return
elif isinstance(n, (list, tuple, np.ndarray)):
if len(n) == 0:
mat = []
elif isinstance(n[0], Matrix):
raise Exception(
'cols/rows required when using list of vectors')
elif not isinstance(n[0], Iterable):
if impl.inside_kernel():
# wrap potential constants with Expr
if keep_raw:
mat = [list([x]) for x in n]
else:
if in_python_scope(
) or disable_local_tensor or not ti.current_cfg(
).dynamic_index:
mat = [list([expr.Expr(x)]) for x in n]
else:
if not ti.is_extension_supported(
ti.cfg.arch, ti.extension.dynamic_index):
raise Exception(
'Backend ' + str(ti.cfg.arch) +
' doesn\'t support dynamic index')
if dt is None:
if isinstance(n[0], int):
dt = impl.get_runtime().default_ip
elif isinstance(n[0], float):
dt = impl.get_runtime().default_fp
else:
raise Exception(
'dt required when using dynamic_index for local tensor'
)
self.local_tensor_proxy = impl.expr_init_local_tensor(
[len(n)], dt,
expr.make_expr_group([expr.Expr(x)
for x in n]))
mat = []
for i in range(len(n)):
mat.append(
list([
ti.local_subscript_with_offset(
self.local_tensor_proxy, (i, ),
(len(n), ))
]))
else:
mat = [[x] for x in n]
else:
if in_python_scope(
) or disable_local_tensor or not ti.current_cfg(
).dynamic_index:
mat = [list(r) for r in n]
else:
if not ti.is_extension_supported(
ti.cfg.arch, ti.extension.dynamic_index):
raise Exception('Backend ' + str(ti.cfg.arch) +
' doesn\'t support dynamic index')
if dt is None:
if isinstance(n[0][0], int):
dt = impl.get_runtime().default_ip
elif isinstance(n[0][0], float):
dt = impl.get_runtime().default_fp
else:
raise Exception(
'dt required when using dynamic_index for local tensor'
)
self.local_tensor_proxy = impl.expr_init_local_tensor(
[len(n), len(n[0])], dt,
expr.make_expr_group(
[expr.Expr(x) for row in n for x in row]))
mat = []
for i in range(len(n)):
mat.append([])
for j in range(len(n[0])):
mat[i].append(
ti.local_subscript_with_offset(
self.local_tensor_proxy, (i, j),
(len(n), len(n[0]))))
self.n = len(mat)
if len(mat) > 0:
self.m = len(mat[0])
else:
self.m = 1
self.entries = [x for row in mat for x in row]
else:
if dt is None:
# create a local matrix with specific (n, m)
self.entries = [impl.expr_init(None) for i in range(n * m)]
self.n = n
self.m = m
else:
# construct global matrix (deprecated)
warning(
"Declaring global matrices using `ti.Matrix(n, m, dt, shape)` is deprecated, "
"use `ti.Matrix.field(n, m, dtype, shape)` instead",
DeprecationWarning,
stacklevel=2)
mat = Matrix.field(n=n,
m=m,
dtype=dt,
shape=shape,
offset=offset,
needs_grad=needs_grad,
layout=layout)
self.n = mat.n
self.m = mat.m
self.entries = mat.entries
self.grad = mat.grad
if self.n * self.m > 32:
warning(
f'Taichi matrices/vectors with {self.n}x{self.m} > 32 entries are not suggested.'
' Matrices/vectors will be automatically unrolled at compile-time for performance.'
' So the compilation time could be extremely long if the matrix size is too big.'
' You may use a field to store a large matrix like this, e.g.:\n'
f' x = ti.field(ti.f32, ({self.n}, {self.m})).\n'
' See https://taichi.readthedocs.io/en/stable/tensor_matrix.html#matrix-size'
' for more details.',
UserWarning,
stacklevel=2)
def shfl_xor_i32(mask, val, offset):
return expr.Expr(
_ti_core.insert_internal_func_call(
"cuda_shfl_xor_sync_i32",
expr.make_expr_group(mask, val, offset, 31), False))
def is_active(l, indices):
return Expr(
_ti_core.insert_is_active(l.snode.ptr, make_expr_group(indices)))
def subscript(value, *_indices, skip_reordered=False):
if isinstance(value, np.ndarray):
return value.__getitem__(_indices)
if isinstance(value, (tuple, list, dict)):
assert len(_indices) == 1
return value[_indices[0]]
has_slice = False
flattened_indices = []
for _index in _indices:
if is_taichi_class(_index):
ind = _index.entries
elif isinstance(_index, slice):
ind = [_index]
has_slice = True
else:
ind = [_index]
flattened_indices += ind
_indices = tuple(flattened_indices)
if len(_indices) == 1 and _indices[0] is None:
_indices = ()
if has_slice:
if not isinstance(value, Matrix):
raise SyntaxError(
f"The type {type(value)} do not support index of slice type")
else:
indices_expr_group = make_expr_group(*_indices)
index_dim = indices_expr_group.size()
if is_taichi_class(value):
return value._subscript(*_indices)
if isinstance(value, MeshElementFieldProxy):
return value.subscript(*_indices)
if isinstance(value, MeshRelationAccessProxy):
return value.subscript(*_indices)
if isinstance(value,
(MeshReorderedScalarFieldProxy,
MeshReorderedMatrixFieldProxy)) and not skip_reordered:
assert index_dim == 1
reordered_index = tuple([
Expr(
_ti_core.get_index_conversion(value.mesh_ptr,
value.element_type,
Expr(_indices[0]).ptr,
ConvType.g2r))
])
return subscript(value, *reordered_index, skip_reordered=True)
if isinstance(value, SparseMatrixProxy):
return value.subscript(*_indices)
if isinstance(value, Field):
_var = value._get_field_members()[0].ptr
if _var.snode() is None:
if _var.is_primal():
raise RuntimeError(
f"{_var.get_expr_name()} has not been placed.")
else:
raise RuntimeError(
f"Gradient {_var.get_expr_name()} has not been placed, check whether `needs_grad=True`"
)
field_dim = int(_var.get_attribute("dim"))
if field_dim != index_dim:
raise IndexError(
f'Field with dim {field_dim} accessed with indices of dim {index_dim}'
)
if isinstance(value, MatrixField):
return _MatrixFieldElement(value, indices_expr_group)
if isinstance(value, StructField):
return _IntermediateStruct(
{k: subscript(v, *_indices)
for k, v in value._items})
return Expr(_ti_core.subscript(_var, indices_expr_group))
if isinstance(value, AnyArray):
# TODO: deprecate using get_attribute to get dim
field_dim = int(value.ptr.get_attribute("dim"))
element_dim = len(value.element_shape)
if field_dim != index_dim + element_dim:
raise IndexError(
f'Field with dim {field_dim - element_dim} accessed with indices of dim {index_dim}'
)
if element_dim == 0:
return Expr(_ti_core.subscript(value.ptr, indices_expr_group))
n = value.element_shape[0]
m = 1 if element_dim == 1 else value.element_shape[1]
any_array_access = AnyArrayAccess(value, _indices)
ret = _IntermediateMatrix(n, m, [
any_array_access.subscript(i, j) for i in range(n)
for j in range(m)
])
ret.any_array_access = any_array_access
return ret
if isinstance(value, SNode):
# When reading bit structure we only support the 0-D case for now.
field_dim = 0
if field_dim != index_dim:
raise IndexError(
f'Field with dim {field_dim} accessed with indices of dim {index_dim}'
)
return Expr(_ti_core.subscript(value.ptr, indices_expr_group))
# Directly evaluate in Python for non-Taichi types
return value.__getitem__(*_indices)
def deactivate(l, indices):
_ti_core.insert_deactivate(l.snode.ptr, make_expr_group(indices))
def make_tensor_element_expr(_var, _indices, shape, stride):
return Expr(
_ti_core.make_tensor_element_expr(_var, make_expr_group(*_indices),
shape, stride))
def bitcode_func_call_wrapper(*args):
_ti_core.insert_external_func_call(0, '', self.bc, item,
make_expr_group(args),
make_expr_group([]))
def call_internal(name, *args, with_runtime_context=True):
return expr_init(
_ti_core.insert_internal_func_call(name, make_expr_group(args),
with_runtime_context))
def inclusive_or(value):
return expr.Expr(_ti_core.insert_internal_func_call(
"subgroupInclusiveOr", expr.make_expr_group(value), False),
dtype=value.ptr.get_ret_type())
def subscript_with_offset(var, indices, cols, is_aos):
return Expr(
_ti_core.subscript_with_offset(var.ptr, make_expr_group(*indices),
cols, is_aos))
def broadcast(value, index: i32):
return expr.Expr(
_ti_core.insert_internal_func_call("subgroupBroadcast",
expr.make_expr_group(value, index),
False))
def append(l, indices, val):
a = impl.expr_init(
_ti_core.insert_append(l.snode.ptr, make_expr_group(indices),
Expr(val).ptr))
return a
def group_size():
return expr.Expr(_ti_core.insert_internal_func_call(
"subgroupSize", expr.make_expr_group(), False),
dtype=i32)
def external_func_call(func, args=[], outputs=[]):
func_addr = ctypes.cast(func, ctypes.c_void_p).value
_ti_core.insert_external_func_call(func_addr, '', make_expr_group(args),
make_expr_group(outputs))
def reduce_xor(value):
return expr.Expr(_ti_core.insert_internal_func_call(
"subgroupXor", expr.make_expr_group(value), False),
dtype=value.ptr.get_ret_type())
def deactivate(l, indices):
impl.get_runtime().prog.current_ast_builder().insert_deactivate(
l.snode.ptr, expr.make_expr_group(indices))