[gui] Move all ggui kernels to python by using taichi fields as stagi…

…ng buffers (taichi-dev#2957)

* wip

* wip

* wip

* wip

* wip

* wip

* remove legacy code

* wip

* wip

* wip

* remove ui_kernels

* wip

* Auto Format

* remove unused file

Co-authored-by: Taichi Gardener <[email protected]>

CMakeLists.txt

@@ -103,13 +103,6 @@ foreach(arch IN LISTS HOST_ARCH CUDA_ARCH)
   add_dependencies(${CORE_LIBRARY_NAME} "generate_llvm_runtime_${arch}")
 endforeach()
 
-add_custom_target(
-      "generate_ui_interop_kernels_cuda"
-      COMMAND ${CLANG_EXECUTABLE} -S ui_kernels.cpp -o "ui_kernels_cuda.ll" -fno-exceptions -emit-llvm -std=c++17 -D "ARCH_${arch}" -I ${PROJECT_SOURCE_DIR};
-      COMMAND ${LLVM_AS_EXECUTABLE} "ui_kernels_cuda.ll" -o "ui_kernels_cuda.bc"
-      WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}/taichi/ui/backends/vulkan"
-)
-add_dependencies(${CORE_LIBRARY_NAME} "generate_ui_interop_kernels_cuda")
 
 FILE(WRITE ${CMAKE_CURRENT_LIST_DIR}/taichi/common/version.h
         "#pragma once\n"

cmake/TaichiCore.cmake

@@ -77,7 +77,7 @@ if(TI_WITH_GGUI)
 endif()
 
 # These files are compiled into .bc and loaded as LLVM module dynamically. They should not be compiled into libtaichi. So they're removed here
-file(GLOB BYTECODE_SOURCE "taichi/runtime/llvm/runtime.cpp" "taichi/runtime/llvm/ui_kernels.cpp")
+file(GLOB BYTECODE_SOURCE "taichi/runtime/llvm/runtime.cpp")
 list(REMOVE_ITEM TAICHI_CORE_SOURCE ${BYTECODE_SOURCE})
 
 

python/taichi/shaders/SetImage_vk.frag

@@ -6,7 +6,14 @@ layout(location = 0) in vec2 frag_texcoord;
 
 layout(location = 0) out vec4 out_color;
 
+layout(binding = 1) uniform UBO {
+  float x_factor;
+  float y_factor;
+}
+ubo;
+
 void main() {
-  out_color = texture(texSampler, frag_texcoord);
+  vec2 coord = frag_texcoord.yx * vec2(ubo.y_factor,ubo.x_factor);
+  out_color = texture(texSampler, coord);
   // out_color = vec4(frag_texcoord.xy,0,1);
 }

python/taichi/ui/canvas.py

@@ -4,6 +4,8 @@
 from taichi.lang.kernel_impl import kernel
 from taichi.lang.ops import get_addr
 
+from .staging_buffer import (copy_colors_to_vbo, copy_vertices_to_vbo,
+                             get_vbo_field, to_u8_rgba)
 from .utils import *
 
 
@@ -15,39 +17,53 @@ def set_background_color(self, color):
         self.canvas.set_background_color(color)
 
     def set_image(self, img):
-        info = get_field_info(img)
+        staging_img = to_u8_rgba(img)
+        info = get_field_info(staging_img)
         self.canvas.set_image(info)
 
     def triangles(self,
                   vertices,
                   color=(0.5, 0.5, 0.5),
                   indices=None,
                   per_vertex_color=None):
-        vertices_info = get_field_info(vertices)
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
+        vbo_info = get_field_info(vbo)
         indices_info = get_field_info(indices)
-        colors_info = get_field_info(per_vertex_color)
-        self.canvas.triangles(vertices_info, indices_info, colors_info, color)
+        self.canvas.triangles(vbo_info, indices_info, has_per_vertex_color,
+                              color)
 
     def lines(self,
               vertices,
               width,
               indices=None,
               color=(0.5, 0.5, 0.5),
               per_vertex_color=None):
-        vertices_info = get_field_info(vertices)
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
+        vbo_info = get_field_info(vbo)
         indices_info = get_field_info(indices)
-        colors_info = get_field_info(per_vertex_color)
-        self.canvas.lines(vertices_info, indices_info, colors_info, color,
+        self.canvas.lines(vbo_info, indices_info, has_per_vertex_color, color,
                           width)
 
     def circles(self,
                 vertices,
                 radius,
                 color=(0.5, 0.5, 0.5),
                 per_vertex_color=None):
-        vertices_info = get_field_info(vertices)
-        colors_info = get_field_info(per_vertex_color)
-        self.canvas.circles(vertices_info, colors_info, color, radius)
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
+        vbo_info = get_field_info(vbo)
+        self.canvas.circles(vbo_info, has_per_vertex_color, color, radius)
 
     def scene(self, scene):
         self.canvas.scene(scene)

python/taichi/ui/scene.py

@@ -9,6 +9,8 @@
 from taichi.lang.ops import atomic_add, get_addr
 
 from .camera import Camera
+from .staging_buffer import (copy_colors_to_vbo, copy_normals_to_vbo,
+                             copy_vertices_to_vbo, get_vbo_field)
 from .utils import get_field_info
 
 normals_field_cache = {}
@@ -88,25 +90,32 @@ def mesh(self,
              color=(0.5, 0.5, 0.5),
              per_vertex_color=None,
              two_sided=False):
-        vertices_info = get_field_info(vertices)
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
         if normals is None:
             normals = gen_normals(vertices, indices)
-        normals_info = get_field_info(normals)
+        copy_normals_to_vbo(vbo, normals)
+        vbo_info = get_field_info(vbo)
         indices_info = get_field_info(indices)
-        colors_info = get_field_info(per_vertex_color)
-        super().mesh(vertices_info, normals_info, colors_info, indices_info,
-                     color, two_sided)
-
-    def particles(
-            self,
-            vertices,
-            radius,
-            color=(0.5, 0.5, 0.5),
-            per_vertex_color=None,
-    ):
-        vertices_info = get_field_info(vertices)
-        colors_info = get_field_info(per_vertex_color)
-        super().particles(vertices_info, colors_info, color, radius)
+
+        super().mesh(vbo_info, has_per_vertex_color, indices_info, color,
+                     two_sided)
+
+    def particles(self,
+                  vertices,
+                  radius,
+                  color=(0.5, 0.5, 0.5),
+                  per_vertex_color=None):
+        vbo = get_vbo_field(vertices)
+        copy_vertices_to_vbo(vbo, vertices)
+        has_per_vertex_color = per_vertex_color is not None
+        if has_per_vertex_color:
+            copy_colors_to_vbo(vbo, per_vertex_color)
+        vbo_info = get_field_info(vbo)
+        super().particles(vbo_info, has_per_vertex_color, color, radius)
 
     def point_light(self, pos, color):
         super().point_light(pos, color)

python/taichi/ui/staging_buffer.py

@@ -0,0 +1,130 @@
+from taichi.core import ti_core as _ti_core
+from taichi.core.primitive_types import f32, u8
+from taichi.lang.impl import default_cfg, field, static
+from taichi.lang.kernel_arguments import ext_arr, template
+from taichi.lang.kernel_impl import kernel
+from taichi.lang.matrix import Vector
+from taichi.lang.ndrange import ndrange
+from taichi.lang.ops import atomic_add, get_addr
+
+import taichi as ti
+
+from .utils import get_field_info
+
+vbo_field_cache = {}
+
+
+def get_vbo_field(vertices):
+    if vertices not in vbo_field_cache:
+        N = vertices.shape[0]
+        pos = 3
+        normal = 3
+        tex_coord = 2
+        color = 3
+        vertex_stride = pos + normal + tex_coord + color
+        vbo = Vector.field(vertex_stride, f32, shape=(N, ))
+        vbo_field_cache[vertices] = vbo
+        return vbo
+    else:
+        return vbo_field_cache[vertices]
+
+
+@kernel
+def copy_to_vbo(vbo: template(), src: template(), offset: template(),
+                num_components: template()):
+    for i in src:
+        for c in ti.static(range(num_components)):
+            vbo[i][offset + c] = src[i][c]
+
+
+def validate_input_field(f, name):
+    if f.dtype != f32:
+        raise Exception(f"{name} needs to have dtype f32")
+    if hasattr(f, 'n'):
+        if f.m != 1:
+            raise Exception(
+                f'{name} needs to be a Vector field (matrix with 1 column)')
+    else:
+        raise Exception(f'{name} needs to be a Vector field')
+    if len(f.shape) != 1:
+        raise Exception(f"the shape of {name} needs to be 1-dimensional")
+
+
+def copy_vertices_to_vbo(vbo, vertices):
+    validate_input_field(vertices, "vertices")
+    if not 2 <= vertices.n <= 3:
+        raise Exception(f'vertices can only be 2D or 3D vector fields')
+    copy_to_vbo(vbo, vertices, 0, vertices.n)
+
+
+def copy_normals_to_vbo(vbo, normals):
+    validate_input_field(normals, "normals")
+    if normals.n != 3:
+        raise Exception(f'normals can only be 3D vector fields')
+    copy_to_vbo(vbo, normals, 3, normals.n)
+
+
+def copy_texcoords_to_vbo(vbo, texcoords):
+    validate_input_field(texcoords, "texcoords")
+    if texcoords.n != 2:
+        raise Exception(f'texcoords can only be 3D vector fields')
+    copy_to_vbo(vbo, texcoords, 6, texcoords.n)
+
+
+def copy_colors_to_vbo(vbo, colors):
+    validate_input_field(colors, "colors")
+    if colors.n != 3:
+        raise Exception(f'colors can only be 3D vector fields')
+    copy_to_vbo(vbo, colors, 8, colors.n)
+
+
+@ti.kernel
+def copy_image_f32_to_u8(src: ti.template(), dst: ti.template(),
+                         num_components: ti.template()):
+    for i, j in src:
+        for k in ti.static(range(num_components)):
+            c = src[i, j][k]
+            c = max(0.0, min(1.0, c))
+            c = c * 255
+            dst[i, j][k] = int(c)
+
+
+@ti.kernel
+def copy_image_u8_to_u8(src: ti.template(), dst: ti.template(),
+                        num_components: ti.template()):
+    for i, j in src:
+        for k in ti.static(range(num_components)):
+            dst[i, j][k] = src[i, j][k]
+
+
+# ggui renderer always assumes the input image to be u8 RGBA
+# if the user input is not in this format, a staging ti field is needed
+image_field_cache = {}
+
+
+def to_u8_rgba(image):
+    if not hasattr(image, 'n') or image.m != 1:
+        raise Exception(
+            f'the input image needs to be a Vector field (matrix with 1 column)'
+        )
+    if len(image.shape) != 2:
+        raise Exception(
+            f"the shape of the image must be of the form (width,height)")
+
+    if image.dtype == u8 and image.n == 4:
+        # already in the desired format
+        return image
+
+    if image not in image_field_cache:
+        staging_img = Vector.field(4, u8, image.shape)
+        image_field_cache[image] = staging_img
+    else:
+        staging_img = image_field_cache[image]
+
+    if image.dtype == u8:
+        copy_image_u8_to_u8(image, staging_img, image.n)
+    elif image.dtype == f32:
+        copy_image_f32_to_u8(image, staging_img, image.n)
+    else:
+        raise Exception("dtype of input image must either be u8 or f32")
+    return staging_img

python/taichi/ui/window.py

@@ -17,7 +17,8 @@ def __init__(self, name, res, vsync=False):
         package_path = str(pathlib.Path(__file__).parent.parent)
 
         ti_arch = default_cfg().arch
-        super().__init__(name, res, vsync, package_path, ti_arch)
+        is_packed = default_cfg().packed
+        super().__init__(name, res, vsync, package_path, ti_arch, is_packed)
 
     @property
     def running(self):

setup.py

@@ -191,12 +191,6 @@ def prepare_package(self):
                     print(f"Fetching runtime file {f} to {target} folder")
                     shutil.copy(os.path.join(llvm_runtime_dir, f), target)
 
-            ui_kernel_dir = 'taichi/ui/backends/vulkan'
-            for f in os.listdir(ui_kernel_dir):
-                if f.endswith('.bc'):
-                    print(f"Fetching ui kernel file {f} to {target} folder")
-                    shutil.copy(os.path.join(ui_kernel_dir, f), target)
-
 
 setup(name=project_name,
       packages=packages,