progress monitoring cleanup, fix cancellation issues

Author: atafra

CHANGELOG.md

@@ -1,6 +1,12 @@
 Version History
 ---------------
 
+-   Fixed issues with cancellation through progress monitor callbacks:
+        -   Fixed cancellation requests almost never being fulfilled on CPU
+            devices since `v2.3.0`
+        -   Fixed not calling the callback anymore after requesting cancellation,
+            while the operation is still being executed
+
 ### Changes in v2.3.0:
 
 -   Significantly improved image quality of the `RT` filter in *high* quality

api/api.cpp

@@ -550,7 +550,7 @@ OIDN_API_NAMESPACE_BEGIN
       checkHandle(hDevice);
       OIDN_LOCK_DEVICE(device);
       device->checkCommitted();
-      device->wait();
+      device->waitAndThrow();
     OIDN_CATCH_DEVICE(device)
   }
 

apps/oidnTest.cpp

@@ -941,23 +941,11 @@ bool progressCallback(void* userPtr, double n)
   return n < progress->nMax; // cancel if reached nMax
 }
 
-void progressTest(DeviceRef& device, double nMax = 1000)
+void progressTest(DeviceRef& device, FilterRef& filter, double nMax = 1000)
 {
-  const int W = 1283;
-  const int H = 727;
-
-  FilterRef filter = device.newFilter("RT");
-  REQUIRE(bool(filter));
-
-  auto image = makeConstImage(device, W, H);
-  setFilterImage(filter, "color",  image);
-  setFilterImage(filter, "output", image); // in-place
-
   Progress progress(nMax);
   filter.setProgressMonitorFunction(progressCallback, &progress);
 
-  filter.set("maxMemoryMB", 0); // make sure there will be multiple tiles
-
   filter.commit();
   REQUIRE(device.getError() == Error::None);
 
@@ -968,8 +956,7 @@ void progressTest(DeviceRef& device, double nMax = 1000)
     // Execution should be cancelled but it's not guaranteed
     Error error = device.getError();
     REQUIRE((error == Error::None || error == Error::Cancelled));
-    // Check whether the callback has not been called after requesting cancellation
-    REQUIRE(progress.n >= nMax);
+    REQUIRE((progress.n >= nMax && progress.n <= 1));
   }
   else
   {
@@ -981,26 +968,50 @@ void progressTest(DeviceRef& device, double nMax = 1000)
 
 TEST_CASE("progress monitor", "[progress]")
 {
+  const int W = 1283;
+  const int H = 727;
+
   DeviceRef device = makeAndCommitDevice();
 
-  SECTION("progress monitor: complete")
+  FilterRef filter = device.newFilter("RT");
+  REQUIRE(bool(filter));
+
+  auto image = makeConstImage(device, W, H);
+  setFilterImage(filter, "color",  image);
+  setFilterImage(filter, "output", image); // in-place
+
+  filter.set("maxMemoryMB", 0); // make sure there will be multiple tiles
+
+  SECTION("progress monitor: finish")
   {
-    progressTest(device);
+    progressTest(device, filter);
   }
 
   SECTION("progress monitor: cancel at the middle")
   {
-    progressTest(device, 0.5);
+    progressTest(device, filter, 0.5);
   }
 
   SECTION("progress monitor: cancel at the beginning")
   {
-    progressTest(device, 0);
+    progressTest(device, filter, 0);
   }
 
   SECTION("progress monitor: cancel at the end")
   {
-    progressTest(device, 1);
+    progressTest(device, filter, 1);
+  }
+
+  SECTION("progress monitor: cancel around the middle, finish")
+  {
+    progressTest(device, filter, 0.4);
+    progressTest(device, filter);
+  }
+
+  SECTION("progress monitor: finish, cancel around the middle")
+  {
+    progressTest(device, filter);
+    progressTest(device, filter, 0.6);
   }
 }
 

core/CMakeLists.txt

@@ -48,6 +48,7 @@ set(OIDN_CORE_SOURCES
   module.h
   module.cpp
   op.h
+  op.cpp
   output_process.h
   output_process.cpp
   pool.h

core/autoexposure.h

@@ -20,7 +20,7 @@ OIDN_NAMESPACE_BEGIN
 
 #if !defined(OIDN_COMPILE_METAL_DEVICE)
 
-  class Autoexposure : public Op, public AutoexposureParams
+  class Autoexposure : public BaseOp, public AutoexposureParams
   {
   public:
     explicit Autoexposure(const ImageDesc& srcDesc)

core/concat_conv.h

@@ -18,7 +18,7 @@ OIDN_NAMESPACE_BEGIN
     bool fastMath; // prefer performance over accuracy
   };
 
-  class ConcatConv : public Op, protected ConcatConvDesc
+  class ConcatConv : public BaseOp, protected ConcatConvDesc
   {
   public:
     ConcatConv(const ConcatConvDesc& desc);

core/concat_conv_chw.h

@@ -14,13 +14,15 @@ OIDN_NAMESPACE_BEGIN
   public:
     ConcatConvCHW(Engine* engine, const ConcatConvDesc& desc);
 
-    size_t getScratchByteSize() const override { return conv->getScratchByteSize(); }
+    Engine* getEngine() const override { return conv->getEngine(); }
+
+    size_t getScratchByteSize() override { return conv->getScratchByteSize(); }
     void setScratch(const Ref<Buffer>& scratch) override { conv->setScratch(scratch); }
 
     void setWeight(const Ref<Tensor>& weight) { conv->setWeight(weight); }
 
     void finalize() override { conv->finalize(); }
-    void submit() override { conv->submit(); }
+    void submitKernels(const Ref<CancellationToken>& ct) override { conv->submitKernels(ct); }
 
   private:
     void updateSrc() override;

core/concat_conv_hwc.cpp

@@ -36,7 +36,7 @@ OIDN_NAMESPACE_BEGIN
     return conv1->isSupported() && conv2->isSupported();
   }
 
-  size_t ConcatConvHWC::getScratchByteSize() const
+  size_t ConcatConvHWC::getScratchByteSize()
   {
     return max(conv1->getScratchByteSize(), conv2->getScratchByteSize());
   }
@@ -73,10 +73,10 @@ OIDN_NAMESPACE_BEGIN
     conv2->finalize();
   }
 
-  void ConcatConvHWC::submit()
+  void ConcatConvHWC::submitKernels(const Ref<CancellationToken>& ct)
   {
-    conv1->submit();
-    conv2->submit();
+    conv1->submitKernels(ct);
+    conv2->submitKernels(ct);
   }
 
 OIDN_NAMESPACE_END

core/concat_conv_hwc.h

@@ -15,17 +15,18 @@ OIDN_NAMESPACE_BEGIN
   public:
     ConcatConvHWC(Engine* engine, const ConcatConvDesc& desc);
 
+    Engine* getEngine() const override { return conv1->getEngine(); }
     bool isSupported() const override;
 
-    size_t getScratchByteSize() const override;
+    size_t getScratchByteSize() override;
     void setScratch(const Ref<Buffer>& scratch) override;
 
     TensorDesc getWeight1Desc() const { return weight1Desc; }
     TensorDesc getWeight2Desc() const { return weight2Desc; }
     void setWeight(const Ref<Tensor>& weight1, const Ref<Tensor>& weight2);
 
     void finalize() override;
-    void submit() override;
+    void submitKernels(const Ref<CancellationToken>& ct) override;
 
   private:
     void updateSrc() override;

core/conv.h

@@ -34,7 +34,7 @@ OIDN_NAMESPACE_BEGIN
   };
 
   // Convolution
-  class Conv : public Op, protected ConvDesc
+  class Conv : public BaseOp, protected ConvDesc
   {
   public:
     Conv(const ConvDesc& desc);

core/device.cpp

@@ -3,6 +3,7 @@
 
 #include "device.h"
 #include "subdevice.h"
+#include "engine.h"
 #include "context.h"
 #include "rt_filter.h"
 #include "rtlightmap_filter.h"
@@ -159,6 +160,17 @@ OIDN_NAMESPACE_BEGIN
     }
   }
 
+  void Device::setAsyncError(Error code, const std::string& message)
+  {
+    // Update the stored async error only if the previous error was thrown
+    std::lock_guard<std::mutex> lock(asyncErrorMutex);
+    if (asyncError.code == Error::None)
+    {
+      asyncError.code = code;
+      asyncError.message = message;
+    }
+  }
+
   void Device::setErrorFunction(ErrorFunction func, void* userPtr)
   {
     errorFunc = func;
@@ -291,4 +303,19 @@ OIDN_NAMESPACE_BEGIN
       subdevice->trimScratch();
   }
 
+  void Device::waitAndThrow()
+  {
+    wait();
+
+    // If an asynchronous error was stored, throw it now
+    std::lock_guard<std::mutex> asyncErrorLock(asyncErrorMutex);
+    if (asyncError.code != Error::None)
+    {
+      const Error code = asyncError.code;
+      const std::string message = std::move(asyncError.message);
+      asyncError = {}; // clear the error
+      throw Exception(code, message);
+    }
+  }
+
 OIDN_NAMESPACE_END

core/device.h

@@ -54,6 +54,8 @@ OIDN_NAMESPACE_BEGIN
 
     static void setError(Device* device, Error code, const std::string& message);
     static Error getError(Device* device, const char** outMessage);
+
+    void setAsyncError(Error code, const std::string& message);
     void setErrorFunction(ErrorFunction func, void* userPtr);
 
     // Some devices (e.g. CUDA, HIP) need to change some per-thread state, which must be later restored
@@ -123,6 +125,10 @@ OIDN_NAMESPACE_BEGIN
     // Waits for all previously submitted commands to complete (blocks)
     virtual void wait() = 0;
 
+    // Waits for all previously submitted commands to complete, and throws the first asynchronous
+    // error that occured since the previous invocation of this function (blocks)
+    void waitAndThrow();
+
   protected:
     virtual void init() = 0;
 
@@ -156,7 +162,11 @@ OIDN_NAMESPACE_BEGIN
     };
 
     static thread_local ErrorState globalError;
+
     ThreadLocal<ErrorState> error;
+    ErrorState asyncError;
+    std::mutex asyncErrorMutex;
+
     ErrorFunction errorFunc = nullptr;
     void* errorUserPtr = nullptr;
   };

core/engine.h

@@ -9,6 +9,7 @@
 #include "heap.h"
 #include "buffer.h"
 #include "image.h"
+#include "progress.h"
 
 OIDN_NAMESPACE_BEGIN
 
@@ -80,7 +81,8 @@ OIDN_NAMESPACE_BEGIN
     virtual void submitUSMCopy(void* dstPtr, const void* srcPtr, size_t byteSize);
 
     // Enqueues a host function
-    virtual void submitHostFunc(std::function<void()>&& f) = 0;
+    virtual void submitHostFunc(std::function<void()>&& f,
+                                const Ref<CancellationToken>& ct = nullptr) = 0;
 
     // Issues all previously submitted commands (does not block)
     virtual void flush() {}

core/graph.cpp

@@ -332,6 +332,7 @@ OIDN_NAMESPACE_BEGIN
     tensorScratchPlanner.addDepAllocs(opID, srcAllocIDs, concatSrcs);
 
     ops.push_back(op);
+    workAmount += op->getWorkAmount();
     dirty = true;
   }
 
@@ -374,11 +375,6 @@ OIDN_NAMESPACE_BEGIN
     dirty = false;
   }
 
-  double Graph::getWorkAmount() const
-  {
-    return double(ops.size());
-  }
-
   bool Graph::isSupported() const
   {
     for (const auto& opTensorAllocPair : tensorAllocs)
@@ -423,6 +419,7 @@ OIDN_NAMESPACE_BEGIN
     scratch.reset();
     scratchByteSize = 0;
     privateByteSize = 0;
+    workAmount = 0;
     tensorScratchByteOffset = 0;
     dirty = false;
   }
@@ -455,7 +452,7 @@ OIDN_NAMESPACE_BEGIN
     finalized = true;
   }
 
-  void Graph::submit(Progress& progress)
+  void Graph::submit(const Ref<Progress>& progress)
   {
     if (!finalized)
       throw std::logic_error("graph not finalized");
@@ -468,14 +465,14 @@ OIDN_NAMESPACE_BEGIN
 
     for (size_t i = 0; i < ops.size(); ++i)
     {
-      ops[i]->submit();
+      ops[i]->submit(progress);
 
     #if defined(OIDN_MICROBENCH)
       engine->wait();
       const int numRuns = OIDN_MICROBENCH;
       Timer timer;
       for (int j = 0; j < numRuns; ++j)
-        ops[i]->submit();
+        ops[i]->submit(progress);
       engine->wait();
       const double time = timer.query() / numRuns;
       std::cerr << i << "," << ops[i]->getName() << "," << time * 1000 << std::endl;
@@ -505,8 +502,6 @@ OIDN_NAMESPACE_BEGIN
         dst->dump(toString(i) + "_" + ops[i]->getName() + "_");
       }
     #endif
-
-      progress.update(engine, 1);
     }
 
   #if defined(OIDN_MICROBENCH)