Sow MemoryManager (#3561)

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);