def decl_any_arr_arg(dtype, dim, element_shape, layout):
dtype = cook_dtype(dtype)
arg_id = _ti_core.decl_arg(dtype, True)
element_dim = len(element_shape)
if layout == Layout.AOS:
element_dim = -element_dim
return AnyArray(
_ti_core.make_external_tensor_expr(dtype, dim, arg_id, element_dim),
element_shape, layout)
def decl_ext_arr_arg(dtype, dim):
dtype = cook_dtype(dtype)
arg_id = _ti_core.decl_arg(dtype, True)
return Expr(_ti_core.make_external_tensor_expr(dtype, dim, arg_id))
def decl_scalar_arg(dtype):
dtype = cook_dtype(dtype)
arg_id = _ti_core.decl_arg(dtype, False)
return Expr(_ti_core.make_arg_load_expr(arg_id, dtype))
def decl_sparse_matrix():
ptr_type = cook_dtype(u64)
# Treat the sparse matrix argument as a scalar since we only need to pass in the base pointer
arg_id = _ti_core.decl_arg(ptr_type, False)
return SparseMatrixProxy(_ti_core.make_arg_load_expr(arg_id, ptr_type))
def decl_any_arr_arg(dtype, dim, element_shape, layout):
dtype = cook_dtype(dtype)
arg_id = _ti_core.decl_arg(dtype, True)
return AnyArray(_ti_core.make_external_tensor_expr(dtype, dim, arg_id),
element_shape, layout)