From e94d0c3b935b39ed29d4dc0bdc3ad47ab861a047 Mon Sep 17 00:00:00 2001
From: Gabriel H <64807734+ghuau-innopeak@users.noreply.github.com>
Date: Mon, 28 Mar 2022 03:15:05 -0700
Subject: [PATCH] [gui] Allow to configure the texture data type (#4630)

With AOT modules depending on the device where the kernel is running, we
might support RGBA Float Texture and it's better to pass the data from
the Taichi Compute Kernel directly to the image (normalized between 0
and 1) instead of adding extra step to convert to 0;255
---
 .../backends/vulkan/renderables/set_image.cpp   | 17 +++++++++++++++--
 .../ui/backends/vulkan/renderables/set_image.h  |  1 +
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/taichi/ui/backends/vulkan/renderables/set_image.cpp b/taichi/ui/backends/vulkan/renderables/set_image.cpp
index 7c0fa239616c7..271ffbe3840a9 100644
--- a/taichi/ui/backends/vulkan/renderables/set_image.cpp
+++ b/taichi/ui/backends/vulkan/renderables/set_image.cpp
@@ -37,6 +37,16 @@ void SetImage::update_data(const SetImageInfo &info) {
 
   const FieldInfo &img = info.img;
 
+  // Support configuring the internal image based on the data type of the field
+  // info.  We assume that the internal image is 4 channels and allow the user
+  // to configure either a classic RGBA8 (u8) or RGBA32F (f32). The latter is
+  // useful for target that support this texture type as it allows to re-use the
+  // result of a kernel directly without normalizing the value from [0; 1] to
+  // [0; 255]
+  //
+  // @TODO: Make the number of channel configurable?
+  texture_dtype_ = img.dtype;
+
   int new_width = get_correct_dimension(img.shape[0]);
   int new_height = get_correct_dimension(img.shape[1]);
 
@@ -53,7 +63,7 @@ void SetImage::update_data(const SetImageInfo &info) {
   app_context_->device().image_transition(texture_, ImageLayout::shader_read,
                                           ImageLayout::transfer_dst);
 
-  uint64_t img_size = pixels * 4;
+  uint64_t img_size = pixels * data_type_size(texture_dtype_) * 4;
 
   // If there is no current program, VBO information should be provided directly
   // instead of accessing through the current SNode
@@ -122,11 +132,14 @@ void SetImage::init_set_image(AppContext *app_context,
 }
 
 void SetImage::create_texture() {
-  size_t image_size = width * height * 4;
+  size_t image_size = width * height * data_type_size(texture_dtype_) * 4;
 
   ImageParams params;
   params.dimension = ImageDimension::d2D;
   params.format = BufferFormat::rgba8;
+  if (texture_dtype_ == taichi::lang::PrimitiveType::f32) {
+    params.format = BufferFormat::rgba32f;
+  }
   params.initial_layout = ImageLayout::shader_read;
   // these are flipped because taichi is y-major and vulkan is x-major
   params.x = height;
diff --git a/taichi/ui/backends/vulkan/renderables/set_image.h b/taichi/ui/backends/vulkan/renderables/set_image.h
index 22db823c5f1aa..140093876d357 100644
--- a/taichi/ui/backends/vulkan/renderables/set_image.h
+++ b/taichi/ui/backends/vulkan/renderables/set_image.h
@@ -47,6 +47,7 @@ class SetImage final : public Renderable {
   taichi::lang::DeviceAllocation cpu_staging_buffer_;
   taichi::lang::DeviceAllocation gpu_staging_buffer_;
 
+  taichi::lang::DataType texture_dtype_{taichi::lang::PrimitiveType::u8};
   taichi::lang::DeviceAllocation texture_;
 
  private: