Moved image tracking registration to create_image

Author: PENGUINLIONG

taichi/program/program.h

@@ -300,6 +300,9 @@ class TI_DLL_EXPORT Program {
                                            uint64 *result_buffer) {
     return program_impl_->allocate_memory_ndarray(alloc_size, result_buffer);
   }
+  DeviceAllocation allocate_texture(const ImageParams& params) {
+    return program_impl_->allocate_texture(params);
+  }
 
   Ndarray *create_ndarray(
       const DataType type,

taichi/program/program_impl.h

@@ -98,6 +98,11 @@ class ProgramImpl {
                                                    uint64 *result_buffer) {
     return kDeviceNullAllocation;
   }
+  
+  virtual DeviceAllocation allocate_texture(const ImageParams& params) {
+    return kDeviceNullAllocation;
+  }
+
   virtual ~ProgramImpl() {
   }
 

taichi/program/texture.cpp

@@ -35,7 +35,7 @@ Texture::Texture(Program *prog,
   img_params.y = height;
   img_params.z = depth;
   img_params.initial_layout = ImageLayout::undefined;
-  texture_alloc_ = device->create_image(img_params);
+  texture_alloc_ = prog_->allocate_texture(img_params);
 
   format_ = img_params.format;
 

taichi/rhi/device.h

@@ -642,6 +642,10 @@ class TI_DLL_EXPORT GraphicsDevice : public Device {
 
   virtual std::unique_ptr<Surface> create_surface(
       const SurfaceConfig &config) = 0;
+  // You are not expected to call this directly. If you want to use this image
+  // in a taichi kernel, you usually want to create the image via
+  // `GfxRuntime::create_image`. `GfxRuntime` is available in `ProgramImpl`
+  // of GPU backends. 
   virtual DeviceAllocation create_image(const ImageParams &params) = 0;
   virtual void destroy_image(DeviceAllocation handle) = 0;
 

taichi/runtime/gfx/runtime.cpp

@@ -558,8 +558,17 @@ void GfxRuntime::buffer_copy(DevicePtr dst, DevicePtr src, size_t size) {
   current_cmdlist_->buffer_copy(dst, src, size);
   submit_current_cmdlist_if_timeout();
 }
+
+DeviceAllocation GfxRuntime::create_image(const ImageParams& params) {
+  GraphicsDevice* gfx_device = dynamic_cast<GraphicsDevice*>(device_);
+  TI_ERROR_IF(gfx_device == nullptr, "Image can only be created on a graphics device");
+  DeviceAllocation image = gfx_device->create_image(params);
+  last_image_layouts_[image.alloc_id] = params.initial_layout;
+  return image;
+}
+
 void GfxRuntime::transition_image(DeviceAllocation image, ImageLayout layout) {
-  ImageLayout &last_layout = last_image_layouts_[image.alloc_id];
+  ImageLayout &last_layout = last_image_layouts_.at(image.alloc_id);
   ensure_current_cmdlist();
   current_cmdlist_->image_transition(image, last_layout, layout);
   last_layout = layout;

taichi/runtime/gfx/runtime.h

@@ -102,6 +102,7 @@ class TI_DLL_EXPORT GfxRuntime {
 
   void buffer_copy(DevicePtr dst, DevicePtr src, size_t size);
 
+  DeviceAllocation create_image(const ImageParams& params);
   void transition_image(DeviceAllocation image, ImageLayout layout);
 
   void signal_event(DeviceEvent *event);

taichi/runtime/program_impls/vulkan/vulkan_program.cpp

@@ -193,6 +193,10 @@ DeviceAllocation VulkanProgramImpl::allocate_memory_ndarray(
        /*export_sharing=*/false});
 }
 
+DeviceAllocation VulkanProgramImpl::allocate_texture(const ImageParams& params) {
+  return vulkan_runtime_->create_image(params);
+}
+
 std::unique_ptr<aot::Kernel> VulkanProgramImpl::make_aot_kernel(
     Kernel &kernel) {
   spirv::lower(&kernel);

taichi/runtime/program_impls/vulkan/vulkan_program.h

@@ -65,6 +65,7 @@ class VulkanProgramImpl : public ProgramImpl {
 
   DeviceAllocation allocate_memory_ndarray(std::size_t alloc_size,
                                            uint64 *result_buffer) override;
+  DeviceAllocation allocate_texture(const ImageParams& params) override;
 
   Device *get_compute_device() override {
     if (embedded_device_) {