third_party/jpeg-xl/lib/jpegli/memory_manager.cc

   1 // Copyright (c) the JPEG XL Project Authors. All rights reserved.
   2 //
   3 // Use of this source code is governed by a BSD-style
   4 // license that can be found in the LICENSE file.
   5
   6 #include "lib/jpegli/memory_manager.h"
   7
   8 #include <string.h>
   9
  10 #include <hwy/aligned_allocator.h>
  11 #include <vector>
  12
  13 #include "lib/jpegli/common_internal.h"
  14 #include "lib/jpegli/error.h"
  15
  16 struct jvirt_sarray_control {
  17   JSAMPARRAY full_buffer;
  18   size_t numrows;
  19   JDIMENSION maxaccess;
  20 };
  21
  22 struct jvirt_barray_control {
  23   JBLOCKARRAY full_buffer;
  24   size_t numrows;
  25   JDIMENSION maxaccess;
  26 };
  27
  28 namespace jpegli {
  29
  30 namespace {
  31
  32 struct MemoryManager {
  33   struct jpeg_memory_mgr pub;
  34   std::vector<void*> owned_ptrs[2 * JPOOL_NUMPOOLS];
  35   uint64_t pool_memory_usage[2 * JPOOL_NUMPOOLS];
  36   uint64_t total_memory_usage;
  37   uint64_t peak_memory_usage;
  38 };
  39
  40 void* Alloc(j_common_ptr cinfo, int pool_id, size_t sizeofobject) {
  41   MemoryManager* mem = reinterpret_cast<MemoryManager*>(cinfo->mem);
  42   if (pool_id < 0 || pool_id >= 2 * JPOOL_NUMPOOLS) {
  43     JPEGLI_ERROR("Invalid pool id %d", pool_id);
  44   }
  45   if (mem->pub.max_memory_to_use > 0 &&
  46       mem->total_memory_usage + static_cast<uint64_t>(sizeofobject) >
  47           static_cast<uint64_t>(mem->pub.max_memory_to_use)) {
  48     JPEGLI_ERROR("Total memory usage exceeding %ld",
  49                  mem->pub.max_memory_to_use);
  50   }
  51   void* p;
  52   if (pool_id < JPOOL_NUMPOOLS) {
  53     p = malloc(sizeofobject);
  54   } else {
  55     p = hwy::AllocateAlignedBytes(sizeofobject, nullptr, nullptr);
  56   }
  57   if (p == nullptr) {
  58     JPEGLI_ERROR("Out of memory");
  59   }
  60   mem->owned_ptrs[pool_id].push_back(p);
  61   mem->pool_memory_usage[pool_id] += sizeofobject;
  62   mem->total_memory_usage += sizeofobject;
  63   mem->peak_memory_usage =
  64       std::max(mem->peak_memory_usage, mem->total_memory_usage);
  65   return p;
  66 }
  67
  68 constexpr size_t gcd(size_t a, size_t b) { return b == 0 ? a : gcd(b, a % b); }
  69 constexpr size_t lcm(size_t a, size_t b) { return (a * b) / gcd(a, b); }
  70
  71 template <typename T>
  72 T** Alloc2dArray(j_common_ptr cinfo, int pool_id, JDIMENSION samplesperrow,
  73                  JDIMENSION numrows) {
  74   T** array = Allocate<T*>(cinfo, numrows, pool_id);
  75   // Always use aligned allocator for large 2d arrays.
  76   if (pool_id < JPOOL_NUMPOOLS) {
  77     pool_id += JPOOL_NUMPOOLS;
  78   }
  79   size_t alignment = lcm(sizeof(T), HWY_ALIGNMENT);
  80   size_t memstride = RoundUpTo(samplesperrow * sizeof(T), alignment);
  81   size_t stride = memstride / sizeof(T);
  82   T* buffer = Allocate<T>(cinfo, numrows * stride, pool_id);
  83   for (size_t i = 0; i < numrows; ++i) {
  84     array[i] = &buffer[i * stride];
  85   }
  86   return array;
  87 }
  88
  89 template <typename Control, typename T>
  90 Control* RequestVirtualArray(j_common_ptr cinfo, int pool_id, boolean pre_zero,
  91                              JDIMENSION samplesperrow, JDIMENSION numrows,
  92                              JDIMENSION maxaccess) {
  93   if (pool_id != JPOOL_IMAGE) {
  94     JPEGLI_ERROR("Only image lifetime virtual arrays are supported.");
  95   }
  96   Control* p = Allocate<Control>(cinfo, 1, pool_id);
  97   p->full_buffer = Alloc2dArray<T>(cinfo, pool_id, samplesperrow, numrows);
  98   p->numrows = numrows;
  99   p->maxaccess = maxaccess;
 100   if (pre_zero) {
 101     for (size_t i = 0; i < numrows; ++i) {
 102       memset(p->full_buffer[i], 0, samplesperrow * sizeof(T));
 103     }
 104   }
 105   return p;
 106 }
 107
 108 void RealizeVirtualArrays(j_common_ptr cinfo) {
 109   // Nothing to do, the full arrays were realized at request time already.
 110 }
 111
 112 template <typename Control, typename T>
 113 T** AccessVirtualArray(j_common_ptr cinfo, Control* ptr, JDIMENSION start_row,
 114                        JDIMENSION num_rows, boolean writable) {
 115   if (num_rows > ptr->maxaccess) {
 116     JPEGLI_ERROR("Invalid virtual array access, num rows %u vs max rows %u",
 117                  num_rows, ptr->maxaccess);
 118   }
 119   if (start_row + num_rows > ptr->numrows) {
 120     JPEGLI_ERROR("Invalid virtual array access, %u vs %u total rows",
 121                  start_row + num_rows, ptr->numrows);
 122   }
 123   if (ptr->full_buffer == nullptr) {
 124     JPEGLI_ERROR("Invalid virtual array access, array not realized.");
 125   }
 126   return ptr->full_buffer + start_row;
 127 }
 128
 129 void ClearPool(j_common_ptr cinfo, int pool_id) {
 130   MemoryManager* mem = reinterpret_cast<MemoryManager*>(cinfo->mem);
 131   mem->owned_ptrs[pool_id].clear();
 132   mem->total_memory_usage -= mem->pool_memory_usage[pool_id];
 133   mem->pool_memory_usage[pool_id] = 0;
 134 }
 135
 136 void FreePool(j_common_ptr cinfo, int pool_id) {
 137   MemoryManager* mem = reinterpret_cast<MemoryManager*>(cinfo->mem);
 138   if (pool_id < 0 || pool_id >= JPOOL_NUMPOOLS) {
 139     JPEGLI_ERROR("Invalid pool id %d", pool_id);
 140   }
 141   for (void* ptr : mem->owned_ptrs[pool_id]) {
 142     free(ptr);
 143   }
 144   ClearPool(cinfo, pool_id);
 145   for (void* ptr : mem->owned_ptrs[JPOOL_NUMPOOLS + pool_id]) {
 146     hwy::FreeAlignedBytes(ptr, nullptr, nullptr);
 147   }
 148   ClearPool(cinfo, JPOOL_NUMPOOLS + pool_id);
 149 }
 150
 151 void SelfDestruct(j_common_ptr cinfo) {
 152   MemoryManager* mem = reinterpret_cast<MemoryManager*>(cinfo->mem);
 153   for (int pool_id = 0; pool_id < JPOOL_NUMPOOLS; ++pool_id) {
 154     FreePool(cinfo, pool_id);
 155   }
 156   delete mem;
 157   cinfo->mem = nullptr;
 158 }
 159
 160 }  // namespace
 161
 162 void InitMemoryManager(j_common_ptr cinfo) {
 163   MemoryManager* mem = new MemoryManager;
 164   mem->pub.alloc_small = jpegli::Alloc;
 165   mem->pub.alloc_large = jpegli::Alloc;
 166   mem->pub.alloc_sarray = jpegli::Alloc2dArray<JSAMPLE>;
 167   mem->pub.alloc_barray = jpegli::Alloc2dArray<JBLOCK>;
 168   mem->pub.request_virt_sarray =
 169       jpegli::RequestVirtualArray<jvirt_sarray_control, JSAMPLE>;
 170   mem->pub.request_virt_barray =
 171       jpegli::RequestVirtualArray<jvirt_barray_control, JBLOCK>;
 172   mem->pub.realize_virt_arrays = jpegli::RealizeVirtualArrays;
 173   mem->pub.access_virt_sarray =
 174       jpegli::AccessVirtualArray<jvirt_sarray_control, JSAMPLE>;
 175   mem->pub.access_virt_barray =
 176       jpegli::AccessVirtualArray<jvirt_barray_control, JBLOCK>;
 177   mem->pub.free_pool = jpegli::FreePool;
 178   mem->pub.self_destruct = jpegli::SelfDestruct;
 179   mem->pub.max_memory_to_use = 0;
 180   mem->total_memory_usage = 0;
 181   mem->peak_memory_usage = 0;
 182   memset(mem->pool_memory_usage, 0, sizeof(mem->pool_memory_usage));
 183   cinfo->mem = reinterpret_cast<struct jpeg_memory_mgr*>(mem);
 184 }
 185
 186 }  // namespace jpegli