Exemplo n.º 1
0
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
Exemplo n.º 2
0
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))
Exemplo n.º 3
0
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))
Exemplo n.º 4
0
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
Exemplo n.º 5
0
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
Exemplo n.º 6
0
 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()
Exemplo n.º 7
0
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
Exemplo n.º 8
0
 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
Exemplo n.º 9
0
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
Exemplo n.º 10
0
def decl_matrix_arg(matrixtype):
    return Matrix(
        [[decl_scalar_arg(matrixtype.dtype) for _ in range(matrixtype.m)]
         for _ in range(matrixtype.n)])
Exemplo n.º 11
0
        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
Exemplo n.º 12
0
 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)