Sow MemoryManager (#3561)

Pointer to MemoryManager instance is conducted from API frontend to Plane.
It is checked that instance is specified, when necessary, but memory is
allocated as usual.

In next PR Memory manager will be used for actual memory allocation.

(cherry picked from commit 98c96657baeaf47035dd4e0edfa30dd6f1718997)

lib/extras/butteraugli.h

@@ -9,14 +9,14 @@
 #ifndef LIB_EXTRAS_BUTTERAUGLI_H_
 #define LIB_EXTRAS_BUTTERAUGLI_H_
 
-#include <stdlib.h>
-#include <string.h>
-
 #include <atomic>
 #include <cstddef>
+#include <cstdlib>
+#include <cstring>
 #include <memory>
 
 #include "lib/base/compiler_specific.h"
+#include "lib/base/memory_manager.h"
 #include "lib/base/status.h"
 #include "lib/extras/image.h"
 
@@ -150,9 +150,11 @@ struct PsychoImage {
 // Hold it here and only allocate on demand to reduce memory usage.
 struct BlurTemp {
   Status GetTransposed(const ImageF &in, ImageF **out) {
+    JxlMemoryManager *memory_manager = in.memory_manager();
     if (transposed_temp.xsize() == 0) {
-      JXL_ASSIGN_OR_RETURN(transposed_temp,
-                           ImageF::Create(in.ysize(), in.xsize()));
+      JXL_ASSIGN_OR_RETURN(
+          transposed_temp,
+          ImageF::Create(memory_manager, in.ysize(), in.xsize()));
     }
     *out = &transposed_temp;
     return true;

lib/extras/butteraugli_test.cc

@@ -12,6 +12,7 @@
 #include <cstdint>
 #include <utility>
 
+#include "lib/base/memory_manager.h"
 #include "lib/base/random.h"
 #include "lib/base/status.h"
 #include "lib/base/testing.h"
@@ -22,6 +23,7 @@
 #include "lib/extras/packed_image.h"
 #include "lib/extras/packed_image_convert.h"
 #include "lib/extras/test_image.h"
+#include "lib/extras/test_utils.h"
 
 namespace jxl {
 namespace {
@@ -31,16 +33,19 @@ using extras::PackedPixelFile;
 using test::TestImage;
 
 Image3F SinglePixelImage(float red, float green, float blue) {
-  JXL_ASSIGN_OR_DIE(Image3F img, Image3F::Create(1, 1));
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
+  JXL_ASSIGN_OR_DIE(Image3F img, Image3F::Create(memory_manager, 1, 1));
   img.PlaneRow(0, 0)[0] = red;
   img.PlaneRow(1, 0)[0] = green;
   img.PlaneRow(2, 0)[0] = blue;
   return img;
 }
 
 Image3F GetColorImage(const PackedPixelFile& ppf) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   JXL_CHECK(!ppf.frames.empty());
-  JXL_ASSIGN_OR_DIE(Image3F color, Image3F::Create(ppf.xsize(), ppf.ysize()));
+  JXL_ASSIGN_OR_DIE(Image3F color,
+                    Image3F::Create(memory_manager, ppf.xsize(), ppf.ysize()));
   JXL_CHECK(ConvertPackedPixelFileToImage3F(ppf, &color, nullptr));
   return color;
 }
@@ -84,12 +89,14 @@ TEST(ButteraugliInPlaceTest, SinglePixel) {
 }
 
 TEST(ButteraugliInPlaceTest, LargeImage) {
+  JxlMemoryManager* memory_manager = jxl::test::MemoryManager();
   const size_t xsize = 1024;
   const size_t ysize = 1024;
   TestImage img;
   img.SetDimensions(xsize, ysize).AddFrame().RandomFill(777);
   Image3F rgb0 = GetColorImage(img.ppf());
-  JXL_ASSIGN_OR_DIE(Image3F rgb1, Image3F::Create(xsize, ysize));
+  JXL_ASSIGN_OR_DIE(Image3F rgb1,
+                    Image3F::Create(memory_manager, xsize, ysize));
   CopyImageTo(rgb0, &rgb1);
   AddUniformNoise(&rgb1, 0.02f, 7777);
   AddEdge(&rgb1, 0.1f, xsize / 2, xsize / 2);

lib/extras/dec/decode.h

@@ -9,11 +9,9 @@
 
 // Facade for image decoders (PNG, PNM, ...).
 
-#include <stddef.h>
-#include <stdint.h>
-
+#include <cstddef>
+#include <cstdint>
 #include <string>
-#include <vector>
 
 #include "lib/base/span.h"
 #include "lib/base/status.h"

lib/extras/image.cc

@@ -10,6 +10,7 @@
 #include <cstddef>
 #include <cstdint>
 
+#include "lib/base/memory_manager.h"
 #include "lib/base/status.h"
 #include "lib/extras/cache_aligned.h"
 #include "lib/extras/simd_util.h"
@@ -99,6 +100,7 @@ PlaneBase::PlaneBase(const size_t xsize, const size_t ysize,
       orig_xsize_(static_cast<uint32_t>(xsize)),
       orig_ysize_(static_cast<uint32_t>(ysize)),
       bytes_per_row_(BytesPerRow(xsize_, sizeof_t)),
+      memory_manager_(nullptr),
       bytes_(nullptr),
       sizeof_t_(sizeof_t) {
   // TODO(eustas): turn to error instead of abort.
@@ -108,7 +110,9 @@ PlaneBase::PlaneBase(const size_t xsize, const size_t ysize,
   JXL_ASSERT(sizeof_t == 1 || sizeof_t == 2 || sizeof_t == 4 || sizeof_t == 8);
 }
 
-Status PlaneBase::Allocate() {
+Status PlaneBase::Allocate(JxlMemoryManager* memory_manager) {
+  JXL_CHECK(memory_manager_ == nullptr);
+  JXL_CHECK(memory_manager != nullptr);
   JXL_CHECK(!bytes_.get());
 
   // Dimensions can be zero, e.g. for lazily-allocated images. Only allocate
@@ -122,6 +126,7 @@ Status PlaneBase::Allocate() {
     // TODO(eustas): use specialized OOM error code
     return JXL_FAILURE("Failed to allocate memory for image surface");
   }
+  memory_manager_ = memory_manager;
   InitializePadding(*this, sizeof_t_);
 
   return true;
@@ -133,6 +138,7 @@ void PlaneBase::Swap(PlaneBase& other) {
   std::swap(orig_xsize_, other.orig_xsize_);
   std::swap(orig_ysize_, other.orig_ysize_);
   std::swap(bytes_per_row_, other.bytes_per_row_);
+  std::swap(memory_manager_, other.memory_manager_);
   std::swap(bytes_, other.bytes_);
 }
 

lib/extras/image.h

@@ -21,6 +21,7 @@
 #include <utility>  // std::move
 
 #include "lib/base/compiler_specific.h"
+#include "lib/base/memory_manager.h"
 #include "lib/base/status.h"
 #include "lib/extras/cache_aligned.h"
 
@@ -38,6 +39,7 @@ struct PlaneBase {
         orig_xsize_(0),
         orig_ysize_(0),
         bytes_per_row_(0),
+        memory_manager_(nullptr),
         bytes_(nullptr),
         sizeof_t_(0) {}
 
@@ -52,6 +54,12 @@ struct PlaneBase {
   // Move assignment (required for std::vector)
   PlaneBase& operator=(PlaneBase&& other) noexcept = default;
 
+  ~PlaneBase() {
+    if (bytes_.get()) {
+      JXL_CHECK(memory_manager_);
+    }
+  }
+
   void Swap(PlaneBase& other);
 
   // Useful for pre-allocating image with some padding for alignment purposes
@@ -74,6 +82,10 @@ struct PlaneBase {
   // NOTE: do not use this for copying rows - the valid xsize may be much less.
   JXL_INLINE size_t bytes_per_row() const { return bytes_per_row_; }
 
+  JXL_INLINE JxlMemoryManager* memory_manager() const {
+    return memory_manager_;
+  }
+
   // Raw access to byte contents, for interfacing with other libraries.
   // Unsigned char instead of char to avoid surprises (sign extension).
   JXL_INLINE uint8_t* bytes() {
@@ -87,7 +99,7 @@ struct PlaneBase {
 
  protected:
   PlaneBase(size_t xsize, size_t ysize, size_t sizeof_t);
-  Status Allocate();
+  Status Allocate(JxlMemoryManager* memory_manager);
 
   // Returns pointer to the start of a row.
   JXL_INLINE void* VoidRow(const size_t y) const {
@@ -109,6 +121,7 @@ struct PlaneBase {
   uint32_t orig_xsize_;
   uint32_t orig_ysize_;
   size_t bytes_per_row_;  // Includes padding.
+  JxlMemoryManager* memory_manager_;
   CacheAlignedUniquePtr bytes_;
   size_t sizeof_t_;
 };
@@ -144,9 +157,10 @@ class Plane : public detail::PlaneBase {
 
   Plane() = default;
 
-  static StatusOr<Plane> Create(const size_t xsize, const size_t ysize) {
+  static StatusOr<Plane> Create(JxlMemoryManager* memory_manager,
+                                const size_t xsize, const size_t ysize) {
     Plane plane(xsize, ysize, sizeof(T));
-    JXL_RETURN_IF_ERROR(plane.Allocate());
+    JXL_RETURN_IF_ERROR(plane.Allocate(memory_manager));
     return plane;
   }
 
@@ -227,12 +241,13 @@ class Image3 {
     return *this;
   }
 
-  static StatusOr<Image3> Create(const size_t xsize, const size_t ysize) {
-    StatusOr<PlaneT> plane0 = PlaneT::Create(xsize, ysize);
+  static StatusOr<Image3> Create(JxlMemoryManager* memory_manager,
+                                 const size_t xsize, const size_t ysize) {
+    StatusOr<PlaneT> plane0 = PlaneT::Create(memory_manager, xsize, ysize);
     JXL_RETURN_IF_ERROR(plane0.status());
-    StatusOr<PlaneT> plane1 = PlaneT::Create(xsize, ysize);
+    StatusOr<PlaneT> plane1 = PlaneT::Create(memory_manager, xsize, ysize);
     JXL_RETURN_IF_ERROR(plane1.status());
-    StatusOr<PlaneT> plane2 = PlaneT::Create(xsize, ysize);
+    StatusOr<PlaneT> plane2 = PlaneT::Create(memory_manager, xsize, ysize);
     JXL_RETURN_IF_ERROR(plane2.status());
     return Image3(std::move(plane0).value(), std::move(plane1).value(),
                   std::move(plane2).value());
@@ -283,6 +298,9 @@ class Image3 {
   }
 
   // Sizes of all three images are guaranteed to be equal.
+  JXL_INLINE JxlMemoryManager* memory_manager() const {
+    return planes_[0].memory_manager();
+  }
   JXL_INLINE size_t xsize() const { return planes_[0].xsize(); }
   JXL_INLINE size_t ysize() const { return planes_[0].ysize(); }
   // Returns offset [bytes] from one row to the next row of the same plane.

lib/extras/image_color_transform.cc

@@ -9,6 +9,7 @@
 #include <atomic>
 #include <utility>
 
+#include "lib/base/memory_manager.h"
 #include "lib/base/rect.h"
 #include "lib/base/status.h"
 #include "lib/cms/cms_interface.h"
@@ -26,8 +27,10 @@ Status ApplyColorTransform(const ColorEncoding& c_current,
   // Changing IsGray is probably a bug.
   JXL_CHECK(c_current.IsGray() == c_desired.IsGray());
   bool is_gray = c_current.IsGray();
+  JxlMemoryManager* memory_amanger = color.memory_manager();
   if (out->xsize() < rect.xsize() || out->ysize() < rect.ysize()) {
-    JXL_ASSIGN_OR_RETURN(*out, Image3F::Create(rect.xsize(), rect.ysize()));
+    JXL_ASSIGN_OR_RETURN(
+        *out, Image3F::Create(memory_amanger, rect.xsize(), rect.ysize()));
   } else {
     out->ShrinkTo(rect.xsize(), rect.ysize());
   }

lib/extras/image_ops.h

@@ -15,6 +15,7 @@
 #include <limits>
 
 #include "lib/base/compiler_specific.h"
+#include "lib/base/memory_manager.h"
 #include "lib/base/rect.h"
 #include "lib/base/status.h"
 #include "lib/extras/image.h"
@@ -118,7 +119,9 @@ StatusOr<Plane<T>> LinComb(const T lambda1, const Plane<T>& image1,
   const size_t ysize = image1.ysize();
   JXL_CHECK(xsize == image2.xsize());
   JXL_CHECK(ysize == image2.ysize());
-  JXL_ASSIGN_OR_RETURN(Plane<T> out, Plane<T>::Create(xsize, ysize));
+  JxlMemoryManager* memory_manager = image1.memory_manager();
+  JXL_ASSIGN_OR_RETURN(Plane<T> out,
+                       Plane<T>::Create(memory_manager, xsize, ysize));
   for (size_t y = 0; y < ysize; ++y) {
     const T* const JXL_RESTRICT row1 = image1.Row(y);
     const T* const JXL_RESTRICT row2 = image2.Row(y);