fixed autoexposure bug

core/CMakeLists.txt

@@ -54,6 +54,7 @@ set(OIDN_CORE_SOURCES
   pool.cpp
   progress.h
   progress.cpp
+  record.h
   ref.h
   rt_filter.h
   rt_filter.cpp

core/autoexposure.h

@@ -6,7 +6,7 @@
 #if !defined(OIDN_COMPILE_METAL_DEVICE)
   #include "op.h"
   #include "image.h"
-  #include "tensor.h"
+  #include "record.h"
 #endif
 
 OIDN_NAMESPACE_BEGIN
@@ -38,14 +38,13 @@ OIDN_NAMESPACE_BEGIN
       this->src = src;
     }
 
-    // The destination value is assumed to be stored in the scratch buffer!
-    void setDst(float* dst) { this->dst = dst; }
-    float* getDst() const { return dst; }
+    void setDst(const Ref<Record<float>>& dst) { this->dst = dst; }
+    float* getDstPtr() const { return dst->getPtr(); }
 
   protected:
     ImageDesc srcDesc;
     Ref<Image> src;
-    float* dst;
+    Ref<Record<float>> dst;
 
     int numBinsH;
     int numBinsW;

core/record.h

@@ -0,0 +1,31 @@
+// Copyright 2023 Intel Corporation
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "buffer.h"
+#include "exception.h"
+
+OIDN_NAMESPACE_BEGIN
+
+  // Plain value or structure stored in a buffer
+  template<typename T>
+  class Record final : public Memory
+  {
+    static_assert(std::is_trivial<T>::value, "record can be used only for trivial types");
+
+  public:
+    Record(const Ref<Buffer>& buffer, size_t byteOffset = 0)
+      : Memory(buffer, byteOffset)
+    {
+      if (byteOffset + sizeof(T) > buffer->getByteSize())
+        throw Exception(Error::InvalidArgument, "buffer region is out of bounds");
+    }
+
+    T* getPtr() const
+    {
+      return (T*)((char*)buffer->getPtr() + byteOffset);
+    }
+  };
+
+OIDN_NAMESPACE_END

core/unet_filter.cpp

@@ -164,7 +164,7 @@ OIDN_NAMESPACE_BEGIN
           autoexposure->setSrc(color);
           autoexposure->submit();
           progress.update(mainEngine, 1);
-          transferFunc->setInputScale(autoexposure->getDst());
+          transferFunc->setInputScale(autoexposure->getDstPtr());
         }
         else
         {
@@ -544,7 +544,7 @@ OIDN_NAMESPACE_BEGIN
       if (instanceID == 0 && hdr)
       {
         autoexposure->setScratch(scratch);
-        autoexposure->setDst((float*)((char*)scratch->getPtr() + autoexposureDstOffset));
+        autoexposure->setDst(makeRef<Record<float>>(scratch, autoexposureDstOffset));
       }
 
       // Finalize the network

devices/cpu/cpu_autoexposure.cpp

@@ -57,7 +57,7 @@ OIDN_NAMESPACE_BEGIN
         [](Sum a, Sum b) -> Sum { return Sum(a.first+b.first, a.second+b.second); }
       );
 
-    *getDst() = (sum.second > 0) ? (key / math::exp2(sum.first / float(sum.second))) : 1.f;
+    *getDstPtr() = (sum.second > 0) ? (key / math::exp2(sum.first / float(sum.second))) : 1.f;
   }
 
 OIDN_NAMESPACE_END

devices/gpu/gpu_autoexposure.h

@@ -208,6 +208,8 @@ OIDN_NAMESPACE_BEGIN
         throw std::logic_error("autoexposure destination not set");
       if (!scratch)
         throw std::logic_error("autoexposure scratch not set");
+      if (dst->getBuffer() != scratch)
+        throw std::invalid_argument("autoexposure result must be stored in the scratch buffer");
 
       float* bins = (float*)scratch->getPtr();
       float* sums = (float*)((char*)bins + numBins * sizeof(float));
@@ -227,7 +229,7 @@ OIDN_NAMESPACE_BEGIN
       reduceFinal.sums   = sums;
       reduceFinal.counts = counts;
       reduceFinal.size   = numGroups;
-      reduceFinal.result = getDst();
+      reduceFinal.result = getDstPtr();
 
     #if defined(OIDN_COMPILE_METAL)
       engine->submitKernel(WorkDim<2>(numBinsH, numBinsW), WorkDim<2>(maxBinSize, maxBinSize), downsample,