Image Engine: Performance 8 byte images.

Previously we used to cache a float image representation of the image in
rect_float. This adds some incorrect behavior as many areas only expect
one of these buffers to be used.

This patch stores float buffers inside the image engine. This is done per
instance. In the future we should consider making a global cache.

4 files changed, 187 insertions, 20 deletions

diff --git a/source/blender/draw/engines/image/image_buffer_cache.hh b/source/blender/draw/engines/image/image_buffer_cache.hh
new file mode 100644
index 00000000000..ef11551c879
--- /dev/null
+++ b/source/blender/draw/engines/image/image_buffer_cache.hh
@@ -0,0 +1,131 @@
+/*
+ * This program is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU General Public License
+ * as published by the Free Software Foundation; either version 2
+ * of the License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+ *
+ * Copyright 2022, Blender Foundation.
+ */
+
+/** \file
+ * \ingroup draw_engine
+ */
+
+#pragma once
+
+#include "BLI_vector.hh"
+
+#include "IMB_imbuf.h"
+#include "IMB_imbuf_types.h"
+
+struct FloatImageBuffer {
+  ImBuf *source_buffer = nullptr;
+  ImBuf *float_buffer = nullptr;
+  bool is_used = true;
+
+  FloatImageBuffer(ImBuf *source_buffer, ImBuf *float_buffer)
+      : source_buffer(source_buffer), float_buffer(float_buffer)
+  {
+  }
+
+  FloatImageBuffer(FloatImageBuffer &&other) noexcept
+  {
+    source_buffer = other.source_buffer;
+    float_buffer = other.float_buffer;
+    is_used = other.is_used;
+    other.source_buffer = nullptr;
+    other.float_buffer = nullptr;
+  }
+
+  virtual ~FloatImageBuffer()
+  {
+    IMB_freeImBuf(float_buffer);
+    float_buffer = nullptr;
+    source_buffer = nullptr;
+  }
+
+  FloatImageBuffer &operator=(FloatImageBuffer &&other) noexcept
+  {
+    this->source_buffer = other.source_buffer;
+    this->float_buffer = other.float_buffer;
+    is_used = other.is_used;
+    other.source_buffer = nullptr;
+    other.float_buffer = nullptr;
+    return *this;
+  }
+};
+
+struct FloatBufferCache {
+ private:
+  blender::Vector<FloatImageBuffer> cache_;
+
+ public:
+  ImBuf *ensure_float_buffer(ImBuf *image_buffer)
+  {
+    /* Check if we can use the float buffer of the given image_buffer. */
+    if (image_buffer->rect_float != nullptr) {
+      return image_buffer;
+    }
+
+    /* Do we have a cached float buffer. */
+    for (FloatImageBuffer &item : cache_) {
+      if (item.source_buffer == image_buffer) {
+        item.is_used = true;
+        return item.float_buffer;
+      }
+    }
+
+    /* Generate a new float buffer. */
+    IMB_float_from_rect(image_buffer);
+    ImBuf *new_imbuf = IMB_allocImBuf(image_buffer->x, image_buffer->y, image_buffer->planes, 0);
+    new_imbuf->rect_float = image_buffer->rect_float;
+    new_imbuf->flags |= IB_rectfloat;
+    new_imbuf->mall |= IB_rectfloat;
+    image_buffer->rect_float = nullptr;
+    image_buffer->flags &= ~IB_rectfloat;
+    image_buffer->mall &= ~IB_rectfloat;
+
+    cache_.append(FloatImageBuffer(image_buffer, new_imbuf));
+    return new_imbuf;
+  }
+
+  void reset_usage_flags()
+  {
+    for (FloatImageBuffer &buffer : cache_) {
+      buffer.is_used = false;
+    }
+  }
+
+  void mark_used(const ImBuf *image_buffer)
+  {
+    for (FloatImageBuffer &item : cache_) {
+      if (item.source_buffer == image_buffer) {
+        item.is_used = true;
+        return;
+      }
+    }
+  }
+
+  void remove_unused_buffers()
+  {
+    for (int64_t i = cache_.size() - 1; i >= 0; i--) {
+      if (!cache_[i].is_used) {
+        cache_.remove_and_reorder(i);
+      }
+    }
+  }
+
+  void clear()
+  {
+    cache_.clear();
+  }
+};
diff --git a/source/blender/draw/engines/image/image_drawing_mode.hh b/source/blender/draw/engines/image/image_drawing_mode.hh
index ccb0f3e963a..f0f7f221069 100644
--- a/source/blender/draw/engines/image/image_drawing_mode.hh
+++ b/source/blender/draw/engines/image/image_drawing_mode.hh
@@ -172,6 +172,7 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
         if (tile_buffer == nullptr) {
           continue;
         }
+        instance_data.float_buffers.mark_used(tile_buffer);
         BKE_image_release_ibuf(image, tile_buffer, lock);
 
         DRWShadingGroup *shsub = DRW_shgroup_create_sub(shgrp);
@@ -199,12 +200,14 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
     switch (changes.get_result_code()) {
       case ePartialUpdateCollectResult::FullUpdateNeeded:
         instance_data.mark_all_texture_slots_dirty();
+        instance_data.float_buffers.clear();
         break;
       case ePartialUpdateCollectResult::NoChangesDetected:
         break;
       case ePartialUpdateCollectResult::PartialChangesDetected:
         /* Partial update when wrap repeat is enabled is not supported. */
         if (instance_data.flags.do_tile_drawing) {
+          instance_data.float_buffers.clear();
           instance_data.mark_all_texture_slots_dirty();
         }
         else {
@@ -215,6 +218,33 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
     do_full_update_for_dirty_textures(instance_data, image_user);
   }
 
+  /**
+   * Update the float buffer.
+   *
+   * TODO(jbakker): This is a very expensive operation and should be optimized to perform the
+   * color space conversion + alpha premultiplication on a part of the buffer.
+   * Basically perform a float_from_rect on a given rectangle.
+   */
+  void do_partial_update_float_buffer(
+      ImBuf *float_buffer, PartialUpdateChecker<ImageTileData>::CollectResult &iterator) const
+  {
+#if 0
+    ImBuf *src = iterator.tile_data.tile_buffer;
+    src->rect_float = float_buffer->rect_float;
+    IMB_float_from_rect(src);
+
+    src->rect_float = nullptr;
+#else
+    ImBuf *src = iterator.tile_data.tile_buffer;
+    BLI_assert(float_buffer->float_rect != nullptr);
+    BLI_assert(float_buffer->rect == nullptr);
+    BLI_assert(src->float_rect == nullptr);
+    BLI_assert(src->rect != nullptr);
+    IMB_float_from_rect_ex(float_buffer, src, &iterator.changed_region.region);
+
+#endif
+  }
+
   void do_partial_update(PartialUpdateChecker<ImageTileData>::CollectResult &iterator,
                          IMAGE_InstanceData &instance_data) const
   {
@@ -223,7 +253,11 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
       if (iterator.tile_data.tile_buffer == nullptr) {
         continue;
       }
-      const bool do_free_float_buffer = ensure_float_buffer(*iterator.tile_data.tile_buffer);
+      ImBuf *tile_buffer = ensure_float_buffer(instance_data, iterator.tile_data.tile_buffer);
+      if (tile_buffer != iterator.tile_data.tile_buffer) {
+        do_partial_update_float_buffer(tile_buffer, iterator);
+      }
+
       const float tile_width = static_cast<float>(iterator.tile_data.tile_buffer->x);
       const float tile_height = static_cast<float>(iterator.tile_data.tile_buffer->y);
 
@@ -299,7 +333,6 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
             &extracted_buffer, texture_region_width, texture_region_height, 32, IB_rectfloat);
 
         int offset = 0;
-        ImBuf *tile_buffer = iterator.tile_data.tile_buffer;
         for (int y = gpu_texture_region_to_update.ymin; y < gpu_texture_region_to_update.ymax;
              y++) {
           float yf = y / (float)texture_height;
@@ -330,10 +363,6 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
                                0);
         imb_freerectImbuf_all(&extracted_buffer);
       }
-
-      if (do_free_float_buffer) {
-        imb_freerectfloatImBuf(iterator.tile_data.tile_buffer);
-      }
     }
   }
 
@@ -392,16 +421,12 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
    * rect_float as the refcounter isn't 0. To work around this we destruct any created local
    * buffers ourself.
    */
-  bool ensure_float_buffer(ImBuf &image_buffer) const
+  ImBuf *ensure_float_buffer(IMAGE_InstanceData &instance_data, ImBuf *image_buffer) const
   {
-    if (image_buffer.rect_float == nullptr) {
-      IMB_float_from_rect(&image_buffer);
-      return true;
-    }
-    return false;
+    return instance_data.float_buffers.ensure_float_buffer(image_buffer);
   }
 
-  void do_full_update_texture_slot(const IMAGE_InstanceData &instance_data,
+  void do_full_update_texture_slot(IMAGE_InstanceData &instance_data,
                                    const TextureInfo &texture_info,
                                    ImBuf &texture_buffer,
                                    ImBuf &tile_buffer,
@@ -409,7 +434,7 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
   {
     const int texture_width = texture_buffer.x;
     const int texture_height = texture_buffer.y;
-    const bool do_free_float_buffer = ensure_float_buffer(tile_buffer);
+    ImBuf *float_tile_buffer = ensure_float_buffer(instance_data, &tile_buffer);
 
     /* IMB_transform works in a non-consistent space. This should be documented or fixed!.
      * Construct a variant of the info_uv_to_texture that adds the texel space
@@ -440,16 +465,12 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
       transform_mode = IMB_TRANSFORM_MODE_CROP_SRC;
     }
 
-    IMB_transform(&tile_buffer,
+    IMB_transform(float_tile_buffer,
                   &texture_buffer,
                   transform_mode,
                   IMB_FILTER_NEAREST,
                   uv_to_texel,
                   crop_rect_ptr);
-
-    if (do_free_float_buffer) {
-      imb_freerectfloatImBuf(&tile_buffer);
-    }
   }
 
  public:
@@ -468,6 +489,7 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
 
     instance_data->partial_update.ensure_image(image);
     instance_data->clear_dirty_flag();
+    instance_data->float_buffers.reset_usage_flags();
 
     /* Step: Find out which screen space textures are needed to draw on the screen. Remove the
      * screen space textures that aren't needed. */
@@ -488,8 +510,10 @@ template<typename TextureMethod> class ScreenSpaceDrawingMode : public AbstractD
     add_shgroups(instance_data);
   }
 
-  void draw_finish(IMAGE_Data *UNUSED(vedata)) const override
+  void draw_finish(IMAGE_Data *vedata) const override
   {
+    IMAGE_InstanceData *instance_data = vedata->instance_data;
+    instance_data->float_buffers.remove_unused_buffers();
   }
 
   void draw_scene(IMAGE_Data *vedata) const override
diff --git a/source/blender/draw/engines/image/image_instance_data.hh b/source/blender/draw/engines/image/image_instance_data.hh
index 420f9396f3f..f8f20fbe178 100644
--- a/source/blender/draw/engines/image/image_instance_data.hh
+++ b/source/blender/draw/engines/image/image_instance_data.hh
@@ -23,6 +23,7 @@
 #pragma once
 
 #include "image_batches.hh"
+#include "image_buffer_cache.hh"
 #include "image_partial_updater.hh"
 #include "image_private.hh"
 #include "image_shader_params.hh"
@@ -63,11 +64,18 @@ struct IMAGE_InstanceData {
     DRWPass *depth_pass;
   } passes;
 
+  /**
+   * Cache containing the float buffers when drawing byte images.
+   */
+  FloatBufferCache float_buffers;
+
   /** \brief Transform matrix to convert a normalized screen space coordinates to texture space. */
   float ss_to_texture[4][4];
   TextureInfo texture_infos[SCREEN_SPACE_DRAWING_MODE_TEXTURE_LEN];
 
  public:
+  virtual ~IMAGE_InstanceData() = default;
+
   void clear_dirty_flag()
   {
     reset_dirty_flag(false);
@@ -117,6 +125,7 @@ struct IMAGE_InstanceData {
     if (last_usage != usage) {
       last_usage = usage;
       reset_dirty_flag(true);
+      float_buffers.clear();
     }
   }
 
diff --git a/source/blender/draw/engines/image/image_usage.hh b/source/blender/draw/engines/image/image_usage.hh
index 1eadee4481d..bbd56e3db7a 100644
--- a/source/blender/draw/engines/image/image_usage.hh
+++ b/source/blender/draw/engines/image/image_usage.hh
@@ -38,6 +38,8 @@ struct ImageUsage {
   /** IMA_ALPHA_* */
   char alpha_mode;
 
+  const void *last_image = nullptr;
+
   ImageUsage() = default;
   ImageUsage(const struct Image *image, const struct ImageUser *image_user)
   {
@@ -46,6 +48,7 @@ struct ImageUsage {
     view = image_user ? image_user->multi_index : 0;
     colorspace_settings = image->colorspace_settings;
     alpha_mode = image->alpha_mode;
+    last_image = static_cast<const void *>(image);
   }
 
   bool operator==(const ImageUsage &other) const