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],
linalg.sparse_matrix_builder):
ctx.create_variable(arg.arg,
kernel_arguments.decl_sparse_matrix())
elif isinstance(ctx.func.argument_annotations[i],
annotations.any_arr):
ctx.create_variable(
arg.arg,
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]))
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 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.arguments[i].annotation,
annotations.template):
ctx.create_variable(arg.arg, ctx.global_vars[arg.arg])
elif isinstance(ctx.func.arguments[i].annotation,
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.arguments[i].annotation,
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.arguments[i].annotation,
texture_type.TextureType):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_texture_arg(
ctx.func.arguments[i].annotation.num_dimensions))
elif isinstance(ctx.func.arguments[i].annotation, MatrixType):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_matrix_arg(
ctx.func.arguments[i].annotation))
elif isinstance(ctx.func.arguments[i].annotation,
primitive_types.RefType):
ctx.create_variable(
arg.arg,
kernel_arguments.decl_scalar_arg(
ctx.func.arguments[i].annotation))
else:
ctx.global_vars[
arg.arg] = kernel_arguments.decl_scalar_arg(
ctx.func.arguments[i].annotation)
# remove original args
node.args.args = []