[lang] Texture support 3/n (Python changes)

python/taichi/examples/rendering/simple_texture.py

@@ -0,0 +1,50 @@
+from taichi.examples.patterns import taichi_logo
+
+import taichi as ti
+
+ti.init(arch=ti.vulkan)
+
+res = (512, 512)
+pixels = ti.Vector.field(3, dtype=float, shape=res)
+
+tex_format = ti.u8
+tex = ti.Texture(tex_format, 1, (128, 128))
+tex_ndarray = ti.ndarray(tex_format, shape=(128, 128))
+
+
+@ti.kernel
+def make_texture(arr: ti.types.ndarray()):
+    for i, j in ti.ndrange(128, 128):
+        ret = taichi_logo(ti.Vector([i, j]) / 128)
+        ret = ti.cast(ret * 255, ti.u8)
+        arr[i, j] = ret
+
+
+make_texture(tex_ndarray)
+tex.from_ndarray(tex_ndarray)
+
+
+@ti.kernel
+def paint(t: ti.f32, tex: ti.types.texture()):
+    for i, j in pixels:
+        uv = ti.Vector([i / res[0], j / res[1]])
+        warp_uv = uv + ti.Vector(
+            [ti.cos(t + uv.x * 5.0),
+             ti.sin(t + uv.y * 5.0)]) * 0.1
+        c = ti.math.vec4(0.0)
+        if uv.x > 0.5:
+            c = tex.sample_lod(warp_uv, 0.0)
+        else:
+            c = tex.fetch(ti.cast(warp_uv * 128, ti.i32), 0)
+        pixels[i, j] = [c.r, c.r, c.r]
+
+
+window = ti.ui.Window('UV', res)
+canvas = window.get_canvas()
+
+t = 0.0
+while window.running:
+    paint(t, tex)
+    canvas.set_image(pixels)
+    window.show()
+    t += 0.03

python/taichi/lang/__init__.py

@@ -1,6 +1,7 @@
 from taichi.lang import impl, simt
 from taichi.lang._ndarray import *
 from taichi.lang._ndrange import ndrange
+from taichi.lang._texture import Texture
 from taichi.lang.enums import Layout
 from taichi.lang.exception import *
 from taichi.lang.field import *

python/taichi/lang/_texture.py

@@ -0,0 +1,67 @@
+from taichi._lib import core as _ti_core
+from taichi.lang import impl
+from taichi.lang.util import taichi_scope
+
+import taichi as ti
+
+
+class TextureSampler:
+    def __init__(self, ptr_expr) -> None:
+        self.ptr_expr = ptr_expr
+
+    @taichi_scope
+    def sample_lod(self, uv, lod):
+        v = _ti_core.make_texture_op_expr(
+            _ti_core.TextureOpType.sample_lod, self.ptr_expr,
+            impl.make_expr_group(uv.x, uv.y, lod))
+        r = impl.call_internal("composite_extract_0",
+                               v,
+                               with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1",
+                               v,
+                               with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2",
+                               v,
+                               with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3",
+                               v,
+                               with_runtime_context=False)
+        return ti.Vector([r, g, b, a])
+
+    @taichi_scope
+    def fetch(self, index, lod):
+        v = _ti_core.make_texture_op_expr(
+            _ti_core.TextureOpType.fetch_texel, self.ptr_expr,
+            impl.make_expr_group(index.x, index.y, lod))
+        r = impl.call_internal("composite_extract_0",
+                               v,
+                               with_runtime_context=False)
+        g = impl.call_internal("composite_extract_1",
+                               v,
+                               with_runtime_context=False)
+        b = impl.call_internal("composite_extract_2",
+                               v,
+                               with_runtime_context=False)
+        a = impl.call_internal("composite_extract_3",
+                               v,
+                               with_runtime_context=False)
+        return ti.Vector([r, g, b, a])
+
+
+class Texture:
+    """Taichi Texture class.
+
+    Args:
+        dtype (DataType): Data type of each value.
+        num_channels (int): Number of channels in texture
+        shape (Tuple[int]): Shape of the Texture.
+    """
+    def __init__(self, dtype, num_channels, arr_shape):
+        self.tex = impl.get_runtime().prog.create_texture(
+            dtype, num_channels, arr_shape)
+
+    def from_ndarray(self, ndarray):
+        self.tex.from_ndarray(ndarray.arr)
+
+    def device_allocation_ptr(self):
+        return self.tex.device_allocation_ptr()

python/taichi/lang/ast/ast_transformer.py

@@ -15,7 +15,8 @@
 from taichi.lang.exception import TaichiSyntaxError
 from taichi.lang.matrix import MatrixType
 from taichi.lang.util import is_taichi_class, to_taichi_type
-from taichi.types import annotations, ndarray_type, primitive_types
+from taichi.types import (annotations, ndarray_type, primitive_types,
+                          texture_type)
 
 if version_info < (3, 9):
     from astunparse import unparse
@@ -481,6 +482,10 @@ def transform_as_kernel():
                             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())
                 elif isinstance(ctx.func.arguments[i].annotation, MatrixType):
                     ctx.create_variable(
                         arg.arg,

python/taichi/lang/kernel_arguments.py

@@ -3,12 +3,13 @@
 import taichi.lang
 from taichi._lib import core as _ti_core
 from taichi.lang import impl, ops
+from taichi.lang._texture import TextureSampler
 from taichi.lang.any_array import AnyArray
 from taichi.lang.enums import Layout
 from taichi.lang.expr import Expr
 from taichi.lang.matrix import Matrix, MatrixType
 from taichi.lang.util import cook_dtype
-from taichi.types.primitive_types import RefType, u64
+from taichi.types.primitive_types import RefType, f32, u64
 
 
 class KernelArgument:
@@ -83,6 +84,11 @@ def decl_ndarray_arg(dtype, dim, element_shape, layout):
         layout)
 
 
+def decl_texture_arg():
+    arg_id = impl.get_runtime().prog.decl_arg(f32, True)
+    return TextureSampler(_ti_core.make_texture_ptr_expr(arg_id))
+
+
 def decl_ret(dtype):
     if isinstance(dtype, MatrixType):
         dtype = _ti_core.decl_tensor_type([dtype.n, dtype.m], dtype.dtype)

python/taichi/lang/kernel_impl.py

@@ -22,7 +22,7 @@
 from taichi.lang.shell import _shell_pop_print, oinspect
 from taichi.lang.util import has_paddle, has_pytorch, to_taichi_type
 from taichi.types import (ndarray_type, primitive_types, sparse_matrix_builder,
-                          template)
+                          template, texture_type)
 
 from taichi import _logging
 
@@ -336,6 +336,8 @@ def extract_arg(arg, anno):
                     TaichiCallableTemplateMapper.extract_arg(item, anno)
                     for item in arg)
             return arg
+        if isinstance(anno, texture_type.TextureType):
+            return '#'
         if isinstance(anno, ndarray_type.NdarrayType):
             if isinstance(arg, taichi.lang._ndarray.ScalarNdarray):
                 anno.check_matched(arg.get_type())
@@ -470,8 +472,8 @@ def extract_arguments(self):
                     raise TaichiSyntaxError(
                         'Taichi kernels parameters must be type annotated')
             else:
-                if isinstance(annotation,
-                              (template, ndarray_type.NdarrayType)):
+                if isinstance(annotation, (template, ndarray_type.NdarrayType,
+                                           texture_type.TextureType)):
                     pass
                 elif id(annotation) in primitive_types.type_ids:
                     pass
@@ -669,6 +671,12 @@ def func__(*args):
                     has_external_arrays = True
                     v = v.arr
                     launch_ctx.set_arg_ndarray(actual_argument_slot, v)
+                elif isinstance(needed,
+                                texture_type.TextureType) and isinstance(
+                                    v, taichi.lang._texture.Texture):
+                    has_external_arrays = True
+                    v = v.tex
+                    launch_ctx.set_arg_texture(actual_argument_slot, v)
                 elif isinstance(
                         needed,
                         ndarray_type.NdarrayType) and (self.match_ext_arr(v)):

python/taichi/types/__init__.py

@@ -12,4 +12,5 @@
 from taichi.types.compound_types import *
 from taichi.types.ndarray_type import *
 from taichi.types.primitive_types import *
+from taichi.types.texture_type import *
 from taichi.types.utils import *

python/taichi/types/texture_type.py

@@ -0,0 +1,12 @@
+class TextureType:
+    """Type annotation for Textures.
+    """
+    def __init__(self):
+        pass
+
+
+texture = TextureType
+"""Alias for :class:`~taichi.types.ndarray_type.TextureType`.
+"""
+
+__all__ = ['texture']

taichi/runtime/gfx/runtime.cpp

@@ -72,7 +72,11 @@ class HostDeviceContextBlitter {
             device_->unmap(buffer);
           }
           // Substitue in the device address if supported
-          if (device_->get_cap(
+          if ((host_ctx_->device_allocation_type[i] ==
+                   RuntimeContext::DevAllocType::kNone ||
+               host_ctx_->device_allocation_type[i] ==
+                   RuntimeContext::DevAllocType::kNdarray) &&
+              device_->get_cap(
                   DeviceCapability::spirv_has_physical_storage_buffer)) {
             uint64_t addr =
                 device_->get_memory_physical_pointer(ext_arrays.at(i));
@@ -435,6 +439,7 @@ void GfxRuntime::launch_kernel(KernelHandle handle, RuntimeContext *host_ctx) {
       host_result_buffer_, args_buffer.get(), ret_buffer.get());
 
   // `any_arrays` contain both external arrays and NDArrays
+  std::vector<std::unique_ptr<DeviceAllocationGuard>> allocated_buffers;
   std::unordered_map<int, DeviceAllocation> any_arrays;
   // `ext_array_size` only holds the size of external arrays (host arrays)
   // As buffer size information is only needed when it needs to be allocated
@@ -473,10 +478,11 @@ void GfxRuntime::launch_kernel(KernelHandle handle, RuntimeContext *host_ctx) {
           ext_array_size[i] = host_ctx->array_runtime_sizes[i];
           // Alloc ext arr
           if (ext_array_size[i]) {
-            DeviceAllocation extarr_buf = device_->allocate_memory(
+            auto allocated = device_->allocate_memory_unique(
                 {ext_array_size[i], /*host_write=*/true, /*host_read=*/true,
                  /*export_sharing=*/false, AllocUsage::Storage});
-            any_arrays[i] = extarr_buf;
+            any_arrays[i] = *allocated.get();
+            allocated_buffers.push_back(std::move(allocated));
           } else {
             any_arrays[i] = kDeviceNullAllocation;
           }
@@ -529,16 +535,6 @@ void GfxRuntime::launch_kernel(KernelHandle handle, RuntimeContext *host_ctx) {
       flush();
     }
   }
-
-  // Dealloc external arrays
-  for (auto pair : any_arrays) {
-    if (pair.second != kDeviceNullAllocation) {
-      if (host_ctx->device_allocation_type[pair.first] ==
-          RuntimeContext::DevAllocType::kNone) {
-        device_->dealloc_memory(pair.second);
-      }
-    }
-  }
 }
 
 void GfxRuntime::synchronize() {

tests/python/test_api.py

@@ -65,28 +65,28 @@ def _get_expected_matrix_apis():
     'SNode', 'ScalarField', 'ScalarNdarray', 'Struct', 'StructField', 'TRACE',
     'TaichiAssertionError', 'TaichiCompilationError', 'TaichiNameError',
     'TaichiRuntimeError', 'TaichiRuntimeTypeError', 'TaichiSyntaxError',
-    'TaichiTypeError', 'TetMesh', 'TriMesh', 'Vector', 'VectorNdarray', 'WARN',
-    'abs', 'acos', 'activate', 'ad', 'aot', 'append', 'arm64', 'asin',
-    'assume_in_range', 'atan2', 'atomic_add', 'atomic_and', 'atomic_max',
-    'atomic_min', 'atomic_or', 'atomic_sub', 'atomic_xor', 'axes', 'bit_cast',
-    'bit_shr', 'block_local', 'cache_read_only', 'cast', 'cc', 'ceil', 'cos',
-    'cpu', 'cuda', 'data_oriented', 'deactivate', 'deactivate_all_snodes',
-    'dx11', 'eig', 'exp', 'experimental', 'extension', 'f16', 'f32', 'f64',
-    'field', 'float16', 'float32', 'float64', 'floor', 'func', 'get_addr',
-    'global_thread_idx', 'gpu', 'graph', 'grouped', 'hex_to_rgb', 'i', 'i16',
-    'i32', 'i64', 'i8', 'ij', 'ijk', 'ijkl', 'ijl', 'ik', 'ikl', 'il', 'init',
-    'int16', 'int32', 'int64', 'int8', 'is_active', 'is_logging_effective',
-    'j', 'jk', 'jkl', 'jl', 'k', 'kernel', 'kl', 'l', 'lang', 'length',
-    'linalg', 'log', 'loop_config', 'math', 'max', 'mesh_local',
-    'mesh_patch_idx', 'metal', 'min', 'ndarray', 'ndrange', 'no_activate',
-    'one', 'opengl', 'polar_decompose', 'pow', 'profiler', 'randn', 'random',
-    'raw_div', 'raw_mod', 'ref', 'rescale_index', 'reset', 'rgb_to_hex',
-    'root', 'round', 'rsqrt', 'select', 'set_logging_level', 'simt', 'sin',
-    'solve', 'sparse_matrix_builder', 'sqrt', 'static', 'static_assert',
-    'static_print', 'stop_grad', 'struct_class', 'svd', 'swizzle_generator',
-    'sym_eig', 'sync', 'tan', 'tanh', 'template', 'tools', 'types', 'u16',
-    'u32', 'u64', 'u8', 'ui', 'uint16', 'uint32', 'uint64', 'uint8', 'vulkan',
-    'wasm', 'x64', 'x86_64', 'zero'
+    'TaichiTypeError', 'TetMesh', 'Texture', 'TriMesh', 'Vector',
+    'VectorNdarray', 'WARN', 'abs', 'acos', 'activate', 'ad', 'aot', 'append',
+    'arm64', 'asin', 'assume_in_range', 'atan2', 'atomic_add', 'atomic_and',
+    'atomic_max', 'atomic_min', 'atomic_or', 'atomic_sub', 'atomic_xor',
+    'axes', 'bit_cast', 'bit_shr', 'block_local', 'cache_read_only', 'cast',
+    'cc', 'ceil', 'cos', 'cpu', 'cuda', 'data_oriented', 'deactivate',
+    'deactivate_all_snodes', 'dx11', 'eig', 'exp', 'experimental', 'extension',
+    'f16', 'f32', 'f64', 'field', 'float16', 'float32', 'float64', 'floor',
+    'func', 'get_addr', 'global_thread_idx', 'gpu', 'graph', 'grouped',
+    'hex_to_rgb', 'i', 'i16', 'i32', 'i64', 'i8', 'ij', 'ijk', 'ijkl', 'ijl',
+    'ik', 'ikl', 'il', 'init', 'int16', 'int32', 'int64', 'int8', 'is_active',
+    'is_logging_effective', 'j', 'jk', 'jkl', 'jl', 'k', 'kernel', 'kl', 'l',
+    'lang', 'length', 'linalg', 'log', 'loop_config', 'math', 'max',
+    'mesh_local', 'mesh_patch_idx', 'metal', 'min', 'ndarray', 'ndrange',
+    'no_activate', 'one', 'opengl', 'polar_decompose', 'pow', 'profiler',
+    'randn', 'random', 'raw_div', 'raw_mod', 'ref', 'rescale_index', 'reset',
+    'rgb_to_hex', 'root', 'round', 'rsqrt', 'select', 'set_logging_level',
+    'simt', 'sin', 'solve', 'sparse_matrix_builder', 'sqrt', 'static',
+    'static_assert', 'static_print', 'stop_grad', 'struct_class', 'svd',
+    'swizzle_generator', 'sym_eig', 'sync', 'tan', 'tanh', 'template', 'tools',
+    'types', 'u16', 'u32', 'u64', 'u8', 'ui', 'uint16', 'uint32', 'uint64',
+    'uint8', 'vulkan', 'wasm', 'x64', 'x86_64', 'zero'
 ]
 user_api[ti.ad] = [
     'FwdMode', 'Tape', 'clear_all_gradients', 'grad_for', 'grad_replaced',
@@ -120,6 +120,7 @@ def _get_expected_matrix_apis():
     'copy_from', 'element_shape', 'fill', 'from_numpy', 'get_type', 'to_numpy'
 ]
 user_api[ti.Ndarray] = ['copy_from', 'element_shape', 'fill', 'get_type']
+user_api[ti.Texture] = ['device_allocation_ptr', 'from_ndarray']
 user_api[ti.SNode] = [
     'bit_array', 'bit_struct', 'bitmasked', 'deactivate_all', 'dense',
     'dynamic', 'lazy_grad', 'parent', 'place', 'pointer', 'shape'