Return SubmitResult from Queue::submit() instead of blocking

Change `Queue::submit()` to return `Expected<SubmitResult>` (a fence +
signal value pair) instead of blocking until completion.  Each backend
(Vulkan, DirectX, Metal) now returns its queue fence and counter value,
and callers explicitly wait via `SubmitResult::waitForCompletion()`
before reading back results.

This allows command buffers submitted between the first and last submit
to remain in-flight on the GPU without host-side synchronization.

- Vulkan: `cleanup()` waits on the queue fence for error-path safety
- DirectX: tile mapping sync reuses the queue's `SubmitFence`
- Metal: `MTLQueue` gains a `SharedEvent` fence; the last command buffer
  in each batch encodes a signal on it

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>

Author: MarijnS95

include/API/Device.h

@@ -54,18 +54,29 @@ class Fence {
   Fence() = default;
 };
 
+/// Returned by Queue::submit() so that callers can decide when to block.
+struct SubmitResult {
+  Fence *F;
+  uint64_t Value;
+
+  llvm::Error waitForCompletion() const { return F->waitForCompletion(Value); }
+};
+
 class Queue {
 public:
   virtual ~Queue() = 0;
-
-  /// Submit command buffers for execution and block until completion.
-  /// Command buffers execute in array order, but dependencies between them
-  /// require appropriate barriers within the command buffers themselves.
-  virtual llvm::Error
+  Queue(const Queue &) = delete;
+  Queue &operator=(const Queue &) = delete;
+  Queue(Queue &&) = default;
+  Queue &operator=(Queue &&) = default;
+
+  /// Submit command buffers for GPU execution.  Returns a fence + value that
+  /// the caller can wait on; the call itself does not block.
+  virtual llvm::Expected<SubmitResult>
   submit(llvm::SmallVector<std::unique_ptr<CommandBuffer>> CBs) = 0;
 
   /// Convenience overload for submitting a single command buffer.
-  llvm::Error submit(std::unique_ptr<CommandBuffer> CB) {
+  llvm::Expected<SubmitResult> submit(std::unique_ptr<CommandBuffer> CB) {
     llvm::SmallVector<std::unique_ptr<CommandBuffer>> CBs;
     CBs.push_back(std::move(CB));
     return submit(std::move(CBs));

lib/API/DX/Device.cpp

@@ -421,7 +421,7 @@ class DXQueue : public offloadtest::Queue {
     return DXQueue(CmdQueue, std::move(*FenceOrErr));
   }
 
-  llvm::Error
+  llvm::Expected<offloadtest::SubmitResult>
   submit(llvm::SmallVector<std::unique_ptr<offloadtest::CommandBuffer>> CBs)
       override;
 };
@@ -458,7 +458,7 @@ class DXCommandBuffer : public offloadtest::CommandBuffer {
   DXCommandBuffer() : CommandBuffer(GPUAPI::DirectX) {}
 };
 
-llvm::Error DXQueue::submit(
+llvm::Expected<offloadtest::SubmitResult> DXQueue::submit(
     llvm::SmallVector<std::unique_ptr<offloadtest::CommandBuffer>> CBs) {
   llvm::SmallVector<ID3D12CommandList *> CmdLists;
   CmdLists.reserve(CBs.size());
@@ -488,11 +488,7 @@ llvm::Error DXQueue::submit(
                       "Failed to add signal."))
     return Err;
 
-  // TODO: Return a Fence+value with keepalive lists instead of blocking here.
-  if (auto Err = SubmitFence->waitForCompletion(CurrentCounter))
-    return Err;
-
-  return llvm::Error::success();
+  return offloadtest::SubmitResult{SubmitFence.get(), CurrentCounter};
 }
 class DXDevice : public offloadtest::Device {
 private:
@@ -526,7 +522,6 @@ class DXDevice : public offloadtest::Device {
     ComPtr<ID3D12DescriptorHeap> DescHeap;
     ComPtr<ID3D12PipelineState> PSO;
     std::unique_ptr<DXCommandBuffer> CB;
-    std::unique_ptr<offloadtest::Fence> Fence;
 
     // Resources for graphics pipelines.
     std::shared_ptr<DXTexture> RT;
@@ -947,19 +942,15 @@ class DXDevice : public offloadtest::Device {
         &HeapRangeStartOffset, &RangeTileCount, D3D12_TILE_MAPPING_FLAG_NONE);
 
     // Synchronize after UpdateTileMappings, which is a queue operation (not
-    // recorded into a command list). This is a hack but it works since this
-    // is all single threaded code.
-    // TODO: Replace with a proper fence abstraction.
-    static uint64_t TileMappingFenceCounter = 0;
-    const uint64_t CurrentCounter = ++TileMappingFenceCounter;
-    auto *F = static_cast<DXFence *>(IS.Fence.get());
-
-    if (auto Err =
-            HR::toError(CommandQueue->Signal(F->Fence.Get(), CurrentCounter),
-                        "Failed to add signal."))
+    // recorded into a command list).
+    const uint64_t CurrentCounter = ++GraphicsQueue.FenceCounter;
+    if (auto Err = HR::toError(
+            CommandQueue->Signal(GraphicsQueue.SubmitFence->Fence.Get(),
+                                 CurrentCounter),
+            "Failed to add signal."))
       return Err;
 
-    return IS.Fence->waitForCompletion(CurrentCounter);
+    return GraphicsQueue.SubmitFence->waitForCompletion(CurrentCounter);
   }
 
   llvm::Expected<ResourceBundle> createSRV(Resource &R, InvocationState &IS) {
@@ -1398,7 +1389,8 @@ class DXDevice : public offloadtest::Device {
     IS.CB->CmdList->ResourceBarrier(1, &Barrier);
   }
 
-  llvm::Error executeCommandList(InvocationState &IS) {
+  llvm::Expected<offloadtest::SubmitResult>
+  executeCommandList(InvocationState &IS) {
     return GraphicsQueue.submit(std::move(IS.CB));
   }
 
@@ -1874,11 +1866,6 @@ class DXDevice : public offloadtest::Device {
     State.CB = std::move(*CBOrErr);
     llvm::outs() << "Command buffer created.\n";
 
-    auto FenceOrErr = createFence("Fence");
-    if (!FenceOrErr)
-      return FenceOrErr.takeError();
-    State.Fence = std::move(*FenceOrErr);
-
     if (auto Err = createBuffers(P, State))
       return Err;
     llvm::outs() << "Buffers created.\n";
@@ -1918,9 +1905,12 @@ class DXDevice : public offloadtest::Device {
       llvm::outs() << "Graphics command list created complete.\n";
     }
 
-    if (auto Err = executeCommandList(State))
-      return Err;
+    auto SubmitResult = executeCommandList(State);
+    if (!SubmitResult)
+      return SubmitResult.takeError();
     llvm::outs() << "Compute commands executed.\n";
+    if (auto Err = SubmitResult->waitForCompletion())
+      return Err;
     if (auto Err = readBack(P, State))
       return Err;
     llvm::outs() << "Read data back.\n";

lib/API/MTL/MTLDevice.cpp

@@ -75,22 +75,6 @@ static MTL::VertexFormat getMTLVertexFormat(DataFormat Format, int Channels) {
 }
 
 namespace {
-class MTLQueue : public offloadtest::Queue {
-public:
-  using Queue::submit;
-
-  MTL::CommandQueue *Queue;
-  MTLQueue(MTL::CommandQueue *Queue) : Queue(Queue) {}
-  ~MTLQueue() override {
-    if (Queue)
-      Queue->release();
-  }
-
-  llvm::Error
-  submit(llvm::SmallVector<std::unique_ptr<offloadtest::CommandBuffer>> CBs)
-      override;
-};
-
 class MTLFence : public offloadtest::Fence {
 public:
   MTLFence(MTL::SharedEvent *Event, llvm::StringRef Name)
@@ -122,6 +106,26 @@ class MTLFence : public offloadtest::Fence {
   }
 };
 
+class MTLQueue : public offloadtest::Queue {
+public:
+  using Queue::submit;
+
+  MTL::CommandQueue *Queue;
+  std::unique_ptr<MTLFence> SubmitFence;
+  uint64_t FenceCounter = 0;
+
+  MTLQueue(MTL::CommandQueue *Queue, std::unique_ptr<MTLFence> SubmitFence)
+      : Queue(Queue), SubmitFence(std::move(SubmitFence)) {}
+  ~MTLQueue() override {
+    if (Queue)
+      Queue->release();
+  }
+
+  llvm::Expected<offloadtest::SubmitResult>
+  submit(llvm::SmallVector<std::unique_ptr<offloadtest::CommandBuffer>> CBs)
+      override;
+};
+
 class MTLBuffer : public offloadtest::Buffer {
 public:
   MTL::Buffer *Buf;
@@ -178,23 +182,21 @@ class MTLCommandBuffer : public offloadtest::CommandBuffer {
   MTLCommandBuffer() : CommandBuffer(GPUAPI::Metal) {}
 };
 
-llvm::Error MTLQueue::submit(
+llvm::Expected<offloadtest::SubmitResult> MTLQueue::submit(
     llvm::SmallVector<std::unique_ptr<offloadtest::CommandBuffer>> CBs) {
   // Metal serial queues guarantee that command buffers execute in commit order,
   // so no explicit wait on prior work is needed here.
-  for (auto &CB : CBs)
-    llvm::cast<MTLCommandBuffer>(CB.get())->CmdBuffer->commit();
-
-  // TODO: Return a Fence+value with keepalive lists instead of blocking here.
-  for (auto &CB : CBs) {
-    auto &MCB = *llvm::cast<MTLCommandBuffer>(CB.get());
-    MCB.CmdBuffer->waitUntilCompleted();
-
-    NS::Error *Err = MCB.CmdBuffer->error();
-    if (Err)
-      return toError(Err);
+  const uint64_t SignalValue = ++FenceCounter;
+
+  for (size_t I = 0; I < CBs.size(); ++I) {
+    auto &MCB = *llvm::cast<MTLCommandBuffer>(CBs[I].get());
+    // Signal the submit fence when the last command buffer completes.
+    if (I == CBs.size() - 1)
+      MCB.CmdBuffer->encodeSignalEvent(SubmitFence->Event, SignalValue);
+    MCB.CmdBuffer->commit();
   }
-  return llvm::Error::success();
+
+  return offloadtest::SubmitResult{SubmitFence.get(), SignalValue};
 }
 class MTLDevice : public offloadtest::Device {
   Capabilities Caps;
@@ -228,7 +230,6 @@ class MTLDevice : public offloadtest::Device {
     std::shared_ptr<MTLBuffer> FrameBufferReadback;
     std::shared_ptr<MTLTexture> DepthStencil;
     std::unique_ptr<MTLCommandBuffer> CB;
-    std::unique_ptr<offloadtest::Fence> Fence;
   };
 
   llvm::Error setupVertexShader(InvocationState &IS, const Pipeline &P,
@@ -666,7 +667,8 @@ class MTLDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
-  llvm::Error executeCommands(InvocationState &IS) {
+  llvm::Expected<offloadtest::SubmitResult>
+  executeCommands(InvocationState &IS) {
     return GraphicsQueue.submit(std::move(IS.CB));
   }
 
@@ -720,8 +722,9 @@ class MTLDevice : public offloadtest::Device {
   }
 
 public:
-  MTLDevice(MTL::Device *D, MTL::CommandQueue *Q)
-      : Device(D), GraphicsQueue(MTLQueue(Q)) {
+  MTLDevice(MTL::Device *D, MTL::CommandQueue *Q,
+            std::unique_ptr<MTLFence> SubmitFence)
+      : Device(D), GraphicsQueue(Q, std::move(SubmitFence)) {
     Description = Device->name()->utf8String();
   }
   const Capabilities &getCapabilities() override {
@@ -783,11 +786,6 @@ class MTLDevice : public offloadtest::Device {
       return CBOrErr.takeError();
     IS.CB = std::move(*CBOrErr);
 
-    auto FenceOrErr = createFence("Fence");
-    if (!FenceOrErr)
-      return FenceOrErr.takeError();
-    IS.Fence = std::move(*FenceOrErr);
-
     if (auto Err = createBuffers(P, IS))
       return Err;
 
@@ -804,7 +802,11 @@ class MTLDevice : public offloadtest::Device {
       llvm::outs() << "Created graphics commands.\n";
     }
 
-    if (auto Err = executeCommands(IS))
+    auto SubmitResult = executeCommands(IS);
+    if (!SubmitResult)
+      return SubmitResult.takeError();
+
+    if (auto Err = SubmitResult->waitForCompletion())
       return Err;
 
     if (auto Err = copyBack(P, IS))
@@ -825,7 +827,12 @@ llvm::Error offloadtest::initializeMetalDevices(
   MTL::Device *MetalDevice = MTL::CreateSystemDefaultDevice();
   MTL::CommandQueue *MetalQueue = MetalDevice->newCommandQueue();
 
-  auto DefaultDev = std::make_unique<MTLDevice>(MetalDevice, MetalQueue);
+  auto SubmitFenceOrErr = MTLFence::create(MetalDevice, "QueueSubmitFence");
+  if (!SubmitFenceOrErr)
+    return SubmitFenceOrErr.takeError();
+
+  auto DefaultDev = std::make_unique<MTLDevice>(MetalDevice, MetalQueue,
+                                                std::move(*SubmitFenceOrErr));
   Devices.push_back(std::move(DefaultDev));
 
   return llvm::Error::success();