def polar_decompose2d(A, dt):
"""Perform polar decomposition (A=UP) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U` and `P`. `U` is a 2x2 orthogonal matrix
and `P` is a 2x2 positive or semi-positive definite matrix.
"""
U = Matrix.identity(dt, 2)
P = ops.cast(A, dt)
zero = ops.cast(0.0, dt)
# if A is a zero matrix we simply return the pair (I, A)
if (A[0, 0] == zero and A[0, 1] == zero and A[1, 0] == zero
and A[1, 1] == zero):
pass
else:
detA = A[0, 0] * A[1, 1] - A[1, 0] * A[0, 1]
adetA = abs(detA)
B = Matrix([[A[0, 0] + A[1, 1], A[0, 1] - A[1, 0]],
[A[1, 0] - A[0, 1], A[1, 1] + A[0, 0]]], dt)
if detA < zero:
B = Matrix([[A[0, 0] - A[1, 1], A[0, 1] + A[1, 0]],
[A[1, 0] + A[0, 1], A[1, 1] - A[0, 0]]], dt)
# here det(B) != 0 if A is not the zero matrix
adetB = abs(B[0, 0] * B[1, 1] - B[1, 0] * B[0, 1])
k = ops.cast(1.0, dt) / ops.sqrt(adetB)
U = B * k
P = (A.transpose() @ A + adetA * Matrix.identity(dt, 2)) * k
return U, P
def expr_init(rhs):
if rhs is None:
return Expr(_ti_core.expr_alloca())
if isinstance(rhs, Matrix):
if rhs.in_python_scope or isinstance(rhs, _IntermediateMatrix):
return Matrix(rhs.to_list())
return rhs
if isinstance(rhs, Struct):
if rhs.in_python_scope or isinstance(rhs, _IntermediateStruct):
return Struct(rhs.to_dict())
return rhs
if isinstance(rhs, list):
return [expr_init(e) for e in rhs]
if isinstance(rhs, tuple):
return tuple(expr_init(e) for e in rhs)
if isinstance(rhs, dict):
return dict((key, expr_init(val)) for key, val in rhs.items())
if isinstance(rhs, _ti_core.DataType):
return rhs
if isinstance(rhs, _ti_core.Arch):
return rhs
if isinstance(rhs, ti.ndrange):
return rhs
if isinstance(rhs, MeshElementFieldProxy):
return rhs
if isinstance(rhs, MeshRelationAccessProxy):
return rhs
if hasattr(rhs, '_data_oriented'):
return rhs
return Expr(_ti_core.expr_var(Expr(rhs).ptr))
def expr_init(rhs):
if rhs is None:
return Expr(get_runtime().prog.current_ast_builder().expr_alloca())
if isinstance(rhs, Matrix) and (hasattr(rhs, "_DIM")):
return type(rhs)(*rhs.to_list())
if isinstance(rhs, Matrix):
return Matrix(rhs.to_list())
if isinstance(rhs, Struct):
return Struct(rhs.to_dict())
if isinstance(rhs, list):
return [expr_init(e) for e in rhs]
if isinstance(rhs, tuple):
return tuple(expr_init(e) for e in rhs)
if isinstance(rhs, dict):
return dict((key, expr_init(val)) for key, val in rhs.items())
if isinstance(rhs, _ti_core.DataType):
return rhs
if isinstance(rhs, _ti_core.Arch):
return rhs
if isinstance(rhs, _Ndrange):
return rhs
if isinstance(rhs, MeshElementFieldProxy):
return rhs
if isinstance(rhs, MeshRelationAccessProxy):
return rhs
if hasattr(rhs, '_data_oriented'):
return rhs
return Expr(get_runtime().prog.current_ast_builder().expr_var(
Expr(rhs).ptr))
def produce_injected_args(kernel, symbolic_args=None):
injected_args = []
for i, arg in enumerate(kernel.arguments):
anno = arg.annotation
if isinstance(anno, template_types):
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if symbolic_args is not None:
element_shape = tuple(symbolic_args[i].element_shape)
element_dim = len(element_shape)
dtype = symbolic_args[i].dtype()
else:
element_shape = anno.element_shape
element_dim = anno.element_dim
dtype = anno.dtype
if element_shape is None or anno.field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={arg.name}')
if element_dim is None or element_dim == 0:
injected_args.append(
ScalarNdarray(dtype, (2, ) * anno.field_dim))
elif element_dim == 1:
injected_args.append(
VectorNdarray(element_shape[0],
dtype=dtype,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
elif element_dim == 2:
injected_args.append(
MatrixNdarray(element_shape[0],
element_shape[1],
dtype=dtype,
shape=(2, ) * anno.field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
elif isinstance(anno, MatrixType):
if not isinstance(symbolic_args[i], list):
raise RuntimeError('Expected a symbolic arg with Matrix type.')
symbolic_mat_n = len(symbolic_args[i])
symbolic_mat_m = len(symbolic_args[i][0])
if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
raise RuntimeError(
f'Matrix dimension mismatch, expected ({anno.n}, {anno.m}) '
f'but dispathed shape ({symbolic_mat_n}, {symbolic_mat_m}).'
)
injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
else:
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
return injected_args
def svd2d(A, dt):
"""Perform singular value decomposition (A=USV^T) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Singular_value_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U`, 'S' and `V`.
"""
R, S = polar_decompose2d(A, dt)
c, s = ops.cast(0.0, dt), ops.cast(0.0, dt)
s1, s2 = ops.cast(0.0, dt), ops.cast(0.0, dt)
if abs(S[0, 1]) < 1e-5:
c, s = 1, 0
s1, s2 = S[0, 0], S[1, 1]
else:
tao = ops.cast(0.5, dt) * (S[0, 0] - S[1, 1])
w = ops.sqrt(tao**2 + S[0, 1]**2)
t = ops.cast(0.0, dt)
if tao > 0:
t = S[0, 1] / (tao + w)
else:
t = S[0, 1] / (tao - w)
c = 1 / ops.sqrt(t**2 + 1)
s = -t * c
s1 = c**2 * S[0, 0] - 2 * c * s * S[0, 1] + s**2 * S[1, 1]
s2 = s**2 * S[0, 0] + 2 * c * s * S[0, 1] + c**2 * S[1, 1]
V = Matrix.zero(dt, 2, 2)
if s1 < s2:
tmp = s1
s1 = s2
s2 = tmp
V = Matrix([[-s, c], [-c, -s]], dt=dt)
else:
V = Matrix([[c, s], [-s, c]], dt=dt)
U = R @ V
return U, Matrix([[s1, ops.cast(0, dt)], [ops.cast(0, dt), s2]], dt=dt), V
def __init__(self, *args, **kwargs):
# converts lists to matrices and dicts to structs
if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
self.entries = args[0]
elif len(args) == 0:
self.entries = kwargs
else:
raise TaichiSyntaxError(
"Custom structs need to be initialized using either dictionary or keyword arguments"
)
for k, v in self.entries.items():
if isinstance(v, (list, tuple)):
v = Matrix(v)
if isinstance(v, dict):
v = Struct(v)
self.entries[k] = v if in_python_scope() else impl.expr_init(v)
self._register_members()
def polar_decompose2d(A, dt):
"""Perform polar decomposition (A=UP) for 2x2 matrix.
Mathematical concept refers to https://en.wikipedia.org/wiki/Polar_decomposition.
Args:
A (ti.Matrix(2, 2)): input 2x2 matrix `A`.
dt (DataType): date type of elements in matrix `A`, typically accepts ti.f32 or ti.f64.
Returns:
Decomposed 2x2 matrices `U` and `P`.
"""
x, y = A(0, 0) + A(1, 1), A(1, 0) - A(0, 1)
scale = (1.0 / ops.sqrt(x * x + y * y))
c = x * scale
s = y * scale
r = Matrix([[c, -s], [s, c]], dt=dt)
return r, r.transpose() @ A
def __init__(self, *args, **kwargs):
# converts lists to matrices and dicts to structs
if len(args) == 1 and kwargs == {} and isinstance(args[0], dict):
self.entries = args[0]
elif len(args) == 0:
self.entries = kwargs
else:
raise TaichiSyntaxError(
"Custom structs need to be initialized using either dictionary or keyword arguments"
)
for k, v in self.entries.items():
if isinstance(v, (list, tuple)):
v = Matrix(v)
if isinstance(v, dict):
v = Struct(v)
self.entries[k] = v
self.register_members()
self.local_tensor_proxy = None
self.any_array_access = None
def produce_injected_args(kernel, symbolic_args=None):
injected_args = []
for i, arg in enumerate(kernel.arguments):
anno = arg.annotation
if isinstance(anno, template_types):
if not isinstance(anno, NdarrayType):
raise TaichiCompilationError(
f'Expected Ndaray type, got {anno}')
if symbolic_args is not None:
element_shape = tuple(symbolic_args[i].element_shape)
element_dim = len(element_shape)
field_dim = symbolic_args[i].field_dim
dtype = symbolic_args[i].dtype()
else:
element_shape = anno.element_shape
element_dim = anno.element_dim
field_dim = anno.field_dim
dtype = anno.dtype
if element_shape is None or field_dim is None:
raise TaichiCompilationError(
'Please either specify both `element_shape` and `field_dim` '
'in the param annotation, or provide an example '
f'ndarray for param={arg.name}')
if anno.field_dim is not None and field_dim != anno.field_dim:
raise TaichiCompilationError(
f'{field_dim} from Arg {arg.name} doesn\'t match kernel\'s annotated field_dim={anno.field_dim}'
)
if anno.dtype is not None and not check_type_match(
dtype, anno.dtype):
raise TaichiCompilationError(
f' Arg {arg.name}\'s dtype {dtype.to_string()} doesn\'t match kernel\'s annotated dtype={anno.dtype.to_string()}'
)
if element_dim is None or element_dim == 0:
injected_args.append(ScalarNdarray(dtype, (2, ) * field_dim))
elif element_dim == 1:
injected_args.append(
VectorNdarray(element_shape[0],
dtype=dtype,
shape=(2, ) * field_dim,
layout=Layout.AOS))
elif element_dim == 2:
injected_args.append(
MatrixNdarray(element_shape[0],
element_shape[1],
dtype=dtype,
shape=(2, ) * field_dim,
layout=Layout.AOS))
else:
raise RuntimeError('')
elif isinstance(anno, (TextureType, RWTextureType)):
if symbolic_args is None:
raise RuntimeError(
'Texture type annotation doesn\'t have enough information for aot. Please either specify the channel_format, shape and num_channels in the graph arg declaration.'
)
texture_shape = tuple(symbolic_args[i].texture_shape)
channel_format = symbolic_args[i].channel_format()
num_channels = symbolic_args[i].num_channels
injected_args.append(
Texture(channel_format, num_channels, texture_shape))
elif isinstance(anno, MatrixType):
if not isinstance(symbolic_args[i], list):
raise RuntimeError('Expected a symbolic arg with Matrix type.')
symbolic_mat_n = len(symbolic_args[i])
symbolic_mat_m = len(symbolic_args[i][0])
if symbolic_mat_m != anno.m or symbolic_mat_n != anno.n:
raise RuntimeError(
f'Matrix dimension mismatch, expected ({anno.n}, {anno.m}) '
f'but dispatched shape ({symbolic_mat_n}, {symbolic_mat_m}).'
)
injected_args.append(Matrix([0] * anno.n * anno.m, dt=anno.dtype))
else:
if symbolic_args is not None:
dtype = symbolic_args[i].dtype()
else:
dtype = anno
if not check_type_match(dtype, anno):
raise TaichiCompilationError(
f' Arg {arg.name}\'s dtype {dtype.to_string()} doesn\'t match kernel\'s annotated dtype={anno.to_string()}'
)
# For primitive types, we can just inject a dummy value.
injected_args.append(0)
return injected_args
def decl_matrix_arg(matrixtype):
return Matrix(
[[decl_scalar_arg(matrixtype.dtype) for _ in range(matrixtype.m)]
for _ in range(matrixtype.n)])
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
def __getitem__(self, key):
self._initialize_host_accessors()
key = self.g2r_field[key]
key = self._pad_key(key)
return Matrix(self._host_access(key), is_ref=True)