def transform_as_kernel():
# Treat return type
if node.returns is not None:
kernel_arguments.decl_ret(ctx.func.return_type)
for i, arg in enumerate(args.args):
if isinstance(ctx.func.argument_annotations[i],
annotations.template):
ctx.create_variable(arg.arg, ctx.global_vars[arg.arg])
elif isinstance(ctx.func.argument_annotations[i],
annotations.sparse_matrix_builder):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_sparse_matrix(
to_taichi_type(ctx.arg_features[i])))
elif isinstance(ctx.func.argument_annotations[i],
ndarray_type.NdarrayType):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_ndarray_arg(
to_taichi_type(ctx.arg_features[i][0]),
ctx.arg_features[i][1], ctx.arg_features[i][2],
ctx.arg_features[i][3]))
elif isinstance(ctx.func.argument_annotations[i], MatrixType):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_matrix_arg(
ctx.func.argument_annotations[i]))
else:
ctx.global_vars[
arg.arg] = kernel_arguments.decl_scalar_arg(
ctx.func.argument_annotations[i])
# remove original args
node.args.args = []
def dtype(self):
"""Gets data type of each individual value.
Returns:
DataType: Data type of each individual value.
"""
return to_taichi_type(self.arr.dtype)
def transform_as_kernel():
# Treat return type
if node.returns is not None:
ti.lang.kernel_arguments.decl_scalar_ret(ctx.func.return_type)
for i, arg in enumerate(args.args):
if isinstance(ctx.func.argument_annotations[i], ti.template):
continue
elif isinstance(ctx.func.argument_annotations[i],
ti.linalg.sparse_matrix_builder):
ctx.create_variable(
arg.arg, ti.lang.kernel_arguments.decl_sparse_matrix())
elif isinstance(ctx.func.argument_annotations[i], ti.any_arr):
ctx.create_variable(
arg.arg,
ti.lang.kernel_arguments.decl_any_arr_arg(
to_taichi_type(ctx.arg_features[i][0]),
ctx.arg_features[i][1], ctx.arg_features[i][2],
ctx.arg_features[i][3]))
else:
ctx.global_vars[
arg.arg] = ti.lang.kernel_arguments.decl_scalar_arg(
ctx.func.argument_annotations[i])
# remove original args
node.args.args = []
def transform_as_kernel():
# Treat return type
if node.returns is not None:
ret_init = parse_stmt(
'ti.lang.kernel_arguments.decl_scalar_ret(0)')
ret_init.value.args[0] = node.returns
ctx.returns = node.returns
arg_decls.append(ret_init)
node.returns = None
for i, arg in enumerate(args.args):
# Directly pass in template arguments,
# such as class instances ("self"), fields, SNodes, etc.
if isinstance(ctx.func.argument_annotations[i], ti.template):
continue
if isinstance(ctx.func.argument_annotations[i],
ti.linalg.sparse_matrix_builder):
arg_init = parse_stmt(
'x = ti.lang.kernel_arguments.decl_sparse_matrix()')
arg_init.targets[0].id = arg.arg
ctx.create_variable(arg.arg)
arg_decls.append(arg_init)
elif isinstance(ctx.func.argument_annotations[i], ti.any_arr):
arg_init = parse_stmt(
'x = ti.lang.kernel_arguments.decl_any_arr_arg(0, 0, 0, 0)'
)
arg_init.targets[0].id = arg.arg
ctx.create_variable(arg.arg)
array_dt = ctx.arg_features[i][0]
array_dim = ctx.arg_features[i][1]
array_element_shape = ctx.arg_features[i][2]
array_layout = ctx.arg_features[i][3]
array_dt = to_taichi_type(array_dt)
dt_expr = 'ti.' + ti.core.data_type_name(array_dt)
dt = parse_expr(dt_expr)
arg_init.value.args[0] = dt
arg_init.value.args[1] = parse_expr("{}".format(array_dim))
arg_init.value.args[2] = parse_expr(
"{}".format(array_element_shape))
arg_init.value.args[3] = parse_expr(
"ti.{}".format(array_layout))
arg_decls.append(arg_init)
else:
arg_init = parse_stmt(
'x = ti.lang.kernel_arguments.decl_scalar_arg(0)')
arg_init.targets[0].id = arg.arg
dt = arg.annotation
arg_init.value.args[0] = dt
arg_decls.append(arg_init)
# remove original args
node.args.args = []
def extract(self, x):
shape = tuple(x.shape)
element_dim = len(self.element_shape)
if len(shape) < element_dim:
raise ValueError("Invalid argument passed to ti.any_arr()")
if element_dim > 0:
if self.layout == Layout.SOA:
if shape[:element_dim] != self.element_shape:
raise ValueError("Invalid argument passed to ti.any_arr()")
else:
if shape[-element_dim:] != self.element_shape:
raise ValueError("Invalid argument passed to ti.any_arr()")
return to_taichi_type(
x.dtype), len(shape), self.element_shape, self.layout
def extract_arg(arg, anno):
if isinstance(anno, template):
if isinstance(arg, taichi.lang.snode.SNode):
return arg.ptr
if isinstance(arg, taichi.lang.expr.Expr):
return arg.ptr.get_underlying_ptr_address()
if isinstance(arg, _ti_core.Expr):
return arg.get_underlying_ptr_address()
if isinstance(arg, tuple):
return tuple(
TaichiCallableTemplateMapper.extract_arg(item, anno)
for item in arg)
return arg
if isinstance(anno, ndarray_type.NdarrayType):
if isinstance(arg, taichi.lang._ndarray.ScalarNdarray):
anno._check_element_dim(arg, 0)
anno._check_element_shape(())
anno._check_field_dim(len(arg.shape))
return arg.dtype, len(arg.shape), (), Layout.AOS
if isinstance(arg, taichi.lang.matrix.VectorNdarray):
anno._check_element_dim(arg, 1)
anno._check_element_shape((arg.n, ))
anno._check_field_dim(len(arg.shape))
anno._check_layout(arg)
return arg.dtype, len(arg.shape) + 1, (arg.n, ), arg.layout
if isinstance(arg, taichi.lang.matrix.MatrixNdarray):
anno._check_element_dim(arg, 2)
anno._check_element_shape((arg.n, arg.m))
anno._check_field_dim(len(arg.shape))
anno._check_layout(arg)
return arg.dtype, len(arg.shape) + 2, (arg.n,
arg.m), arg.layout
# external arrays
element_dim = 0 if anno.element_dim is None else anno.element_dim
layout = Layout.AOS if anno.layout is None else anno.layout
shape = tuple(arg.shape)
if len(shape) < element_dim:
raise ValueError(
f"Invalid argument into ti.types.ndarray() - required element_dim={element_dim}, "
f"but the argument has only {len(shape)} dimensions")
element_shape = (
) if element_dim == 0 else shape[:
element_dim] if layout == Layout.SOA else shape[
-element_dim:]
return to_taichi_type(arg.dtype), len(shape), element_shape, layout
if isinstance(anno, sparse_matrix_builder):
return arg.dtype
# Use '#' as a placeholder because other kinds of arguments are not involved in template instantiation
return '#'
def extract_arg(arg, anno):
if isinstance(anno, template):
if isinstance(arg, taichi.lang.snode.SNode):
return arg.ptr
if isinstance(arg, taichi.lang.expr.Expr):
return arg.ptr.get_underlying_ptr_address()
if isinstance(arg, _ti_core.Expr):
return arg.get_underlying_ptr_address()
if isinstance(arg, tuple):
return tuple(
TaichiCallableTemplateMapper.extract_arg(item, anno)
for item in arg)
return arg
elif isinstance(anno, any_arr):
if isinstance(arg, taichi.lang._ndarray.ScalarNdarray):
anno.check_element_dim(arg, 0)
return arg.dtype, len(arg.shape), (), Layout.AOS
if isinstance(arg, taichi.lang.matrix.VectorNdarray):
anno.check_element_dim(arg, 1)
anno.check_layout(arg)
return arg.dtype, len(arg.shape) + 1, (arg.n, ), arg.layout
if isinstance(arg, taichi.lang.matrix.MatrixNdarray):
anno.check_element_dim(arg, 2)
anno.check_layout(arg)
return arg.dtype, len(arg.shape) + 2, (arg.n,
arg.m), arg.layout
# external arrays
element_dim = 0 if anno.element_dim is None else anno.element_dim
layout = Layout.AOS if anno.layout is None else anno.layout
shape = tuple(arg.shape)
if len(shape) < element_dim:
raise ValueError(
f"Invalid argument into ti.any_arr() - required element_dim={element_dim}, but the argument has only {len(shape)} dimensions"
)
element_shape = (
) if element_dim == 0 else shape[:
element_dim] if layout == Layout.SOA else shape[
-element_dim:]
return to_taichi_type(arg.dtype), len(shape), element_shape, layout
return (type(arg).__name__, )
def extract(self, x):
return to_taichi_type(x.dtype), len(x.shape)
def visit_FunctionDef(self, node):
args = node.args
assert args.vararg is None
assert args.kwonlyargs == []
assert args.kw_defaults == []
assert args.kwarg is None
import taichi as ti
if self.is_kernel: # ti.kernel
for decorator in node.decorator_list:
if (isinstance(decorator, ast.Attribute)
and isinstance(decorator.value, ast.Name)
and decorator.value.id == 'ti'
and decorator.attr == 'func'):
raise TaichiSyntaxError(
"Function definition not allowed in 'ti.kernel'.")
# Transform as kernel
arg_decls = []
# Treat return type
if node.returns is not None:
ret_init = self.parse_stmt(
'ti.lang.kernel_arguments.decl_scalar_ret(0)')
ret_init.value.args[0] = node.returns
self.returns = node.returns
arg_decls.append(ret_init)
node.returns = None
for i, arg in enumerate(args.args):
# Directly pass in template arguments,
# such as class instances ("self"), fields, SNodes, etc.
if isinstance(self.func.arguments[i], ti.template):
continue
import taichi as ti
if isinstance(self.func.arguments[i], ti.ext_arr):
arg_init = self.parse_stmt(
'x = ti.lang.kernel_arguments.decl_ext_arr_arg(0, 0)')
arg_init.targets[0].id = arg.arg
self.create_variable(arg.arg)
array_dt = self.arg_features[i][0]
array_dim = self.arg_features[i][1]
array_dt = to_taichi_type(array_dt)
dt_expr = 'ti.' + ti.core.data_type_name(array_dt)
dt = self.parse_expr(dt_expr)
arg_init.value.args[0] = dt
arg_init.value.args[1] = self.parse_expr(
"{}".format(array_dim))
arg_decls.append(arg_init)
else:
arg_init = self.parse_stmt(
'x = ti.lang.kernel_arguments.decl_scalar_arg(0)')
arg_init.targets[0].id = arg.arg
dt = arg.annotation
arg_init.value.args[0] = dt
arg_decls.append(arg_init)
# remove original args
node.args.args = []
else: # ti.func
for decorator in node.decorator_list:
if (isinstance(decorator, ast.Attribute)
and isinstance(decorator.value, ast.Name)
and decorator.value.id == 'ti'
and decorator.attr == 'func'):
raise TaichiSyntaxError(
"Function definition not allowed in 'ti.func'.")
# Transform as func (all parameters passed by value)
arg_decls = []
for i, arg in enumerate(args.args):
# Directly pass in template arguments,
# such as class instances ("self"), fields, SNodes, etc.
if isinstance(self.func.arguments[i], ti.template):
continue
# Create a copy for non-template arguments,
# so that they are passed by value.
arg_init = self.parse_stmt('x = ti.expr_init_func(0)')
arg_init.targets[0].id = arg.arg
self.create_variable(arg.arg)
arg_init.value.args[0] = self.parse_expr(arg.arg +
'_by_value__')
args.args[i].arg += '_by_value__'
arg_decls.append(arg_init)
with self.variable_scope():
self.generic_visit(node)
node.body = arg_decls + node.body
node.body = [self.parse_stmt('import taichi as ti')] + node.body
return node
