Implement an abstract Fence type to expose waiting for command buffer completion.

include/API/Device.h

@@ -56,6 +56,13 @@ class Buffer {
   Buffer() = default;
 };
 
+class Fence {
+public:
+  virtual ~Fence() {}
+  virtual uint64_t getFenceValue() = 0;
+  virtual llvm::Error waitForCompletion(uint64_t SignalValue) = 0;
+};
+
 class Queue {
 public:
   virtual ~Queue() = 0;
@@ -77,6 +84,9 @@ class Device {
 
   virtual Queue &getGraphicsQueue() = 0;
 
+  virtual llvm::Expected<std::shared_ptr<Fence>>
+  createFence(llvm::StringRef Name) = 0;
+
   virtual llvm::Expected<std::shared_ptr<Buffer>>
   createBuffer(std::string Name, BufferCreateDesc &Desc,
                size_t SizeInBytes) = 0;

lib/API/DX/Device.cpp

@@ -289,6 +289,50 @@ class DXBuffer : public offloadtest::Buffer {
       : Buffer(Buffer), Name(Name), Desc(Desc), SizeInBytes(SizeInBytes) {}
 };
 
+class DXFence : public offloadtest::Fence {
+public:
+  std::string Name;
+  ComPtr<ID3D12Fence> Fence;
+  HANDLE Event;
+
+  uint64_t getFenceValue() override { return Fence->GetCompletedValue(); }
+
+  llvm::Error waitForCompletion(uint64_t SignalValue) override {
+
+    if (Fence->GetCompletedValue() >= SignalValue) {
+      return llvm::Error::success();
+    }
+
+    if (auto Err = HR::toError(Fence->SetEventOnCompletion(SignalValue, Event),
+                               "Failed to register end event."))
+      return Err;
+
+#ifdef _WIN32
+    WaitForSingleObject(Event, INFINITE);
+#else // WSL
+    pollfd PollEvent;
+    PollEvent.fd = (int)Event;
+    PollEvent.events = POLLIN;
+    PollEvent.revents = 0;
+    if (poll(&PollEvent, 1, -1) == -1)
+      return llvm::createStringError(
+          std::error_code(errno, std::system_category()), strerror(errno));
+#endif
+    return llvm::Error::success();
+  }
+
+  DXFence(ComPtr<ID3D12Fence> Fence, HANDLE Event, llvm::StringRef Name)
+      : Fence(Fence), Event(Event), Name(Name) {}
+
+  ~DXFence() {
+#ifdef _WIN32
+    CloseHandle(Event);
+#else // WSL
+    close((int)Event);
+#endif
+  }
+};
+
 class DXQueue : public offloadtest::Queue {
 public:
   ComPtr<ID3D12CommandQueue> Queue;
@@ -342,12 +386,7 @@ class DXDevice : public offloadtest::Device {
     ComPtr<ID3D12PipelineState> PSO;
     ComPtr<ID3D12CommandAllocator> Allocator;
     ComPtr<ID3D12GraphicsCommandList> CmdList;
-    ComPtr<ID3D12Fence> Fence;
-#ifdef _WIN32
-    HANDLE Event;
-#else // WSL
-    int Event;
-#endif
+    std::shared_ptr<offloadtest::Fence> Fence;
 
     // Resources for graphics pipelines.
     ComPtr<ID3D12Resource> RT;
@@ -374,6 +413,27 @@ class DXDevice : public offloadtest::Device {
 
   Queue &getGraphicsQueue() override { return GraphicsQueue; }
 
+  llvm::Expected<std::shared_ptr<offloadtest::Fence>>
+  createFence(llvm::StringRef Name) override {
+    ComPtr<ID3D12Fence> Fence;
+    if (auto Err = HR::toError(
+            Device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&Fence)),
+            "Failed to create Fence."))
+      return Err;
+
+#ifdef _WIN32
+    HANDLE Event = CreateEventA(nullptr, false, false, nullptr);
+    if (!Event)
+#else // WSL
+    int Event = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK);
+    if (Event == -1)
+#endif
+      return llvm::createStringError(std::errc::device_or_resource_busy,
+                                     "Failed to create event.");
+
+    return std::make_shared<DXFence>(Fence, Event, Name);
+  }
+
   llvm::Expected<std::shared_ptr<offloadtest::Buffer>>
   createBuffer(std::string Name, BufferCreateDesc &Desc,
                size_t SizeInBytes) override {
@@ -1132,56 +1192,20 @@ class DXDevice : public offloadtest::Device {
     IS.CmdList->ResourceBarrier(1, &Barrier);
   }
 
-  llvm::Error createEvent(InvocationState &IS) {
-    if (auto Err = HR::toError(Device->CreateFence(0, D3D12_FENCE_FLAG_NONE,
-                                                   IID_PPV_ARGS(&IS.Fence)),
-                               "Failed to create fence."))
-      return Err;
-#ifdef _WIN32
-    IS.Event = CreateEventA(nullptr, false, false, nullptr);
-    if (!IS.Event)
-#else // WSL
-    IS.Event = eventfd(0, EFD_CLOEXEC | EFD_NONBLOCK);
-    if (IS.Event == -1)
-#endif
-      return llvm::createStringError(std::errc::device_or_resource_busy,
-                                     "Failed to create event.");
-    return llvm::Error::success();
-  }
-
   llvm::Error waitForSignal(InvocationState &IS) {
     // This is a hack but it works since this is all single threaded code.
     static uint64_t FenceCounter = 0;
     const uint64_t CurrentCounter = FenceCounter + 1;
+    auto *F = static_cast<DXFence *>(IS.Fence.get());
 
     if (auto Err = HR::toError(
-            GraphicsQueue.Queue->Signal(IS.Fence.Get(), CurrentCounter),
+            GraphicsQueue.Queue->Signal(F->Fence.Get(), CurrentCounter),
             "Failed to add signal."))
       return Err;
 
-    if (IS.Fence->GetCompletedValue() < CurrentCounter) {
-#ifdef _WIN32
-      HANDLE Event = IS.Event;
-#else // WSL
-      HANDLE Event = reinterpret_cast<HANDLE>(IS.Event);
-#endif
-      if (auto Err =
-              HR::toError(IS.Fence->SetEventOnCompletion(CurrentCounter, Event),
-                          "Failed to register end event."))
-        return Err;
+    if (auto Err = IS.Fence->waitForCompletion(CurrentCounter))
+      return Err;
 
-#ifdef _WIN32
-      WaitForSingleObject(IS.Event, INFINITE);
-#else // WSL
-      pollfd PollEvent;
-      PollEvent.fd = IS.Event;
-      PollEvent.events = POLLIN;
-      PollEvent.revents = 0;
-      if (poll(&PollEvent, 1, -1) == -1)
-        return llvm::createStringError(
-            std::error_code(errno, std::system_category()), strerror(errno));
-#endif
-    }
     FenceCounter = CurrentCounter;
     return llvm::Error::success();
   }
@@ -1686,9 +1710,10 @@ class DXDevice : public offloadtest::Device {
       return Err;
     llvm::outs() << "Command structures created.\n";
 
-    if (auto Err = createEvent(State))
-      return Err;
-    llvm::outs() << "Event prepared.\n";
+    auto FenceOrErr = this->createFence("Fence");
+    if (!FenceOrErr)
+      return FenceOrErr.takeError();
+    State.Fence = *FenceOrErr;
 
     if (auto Err = createBuffers(P, State))
       return Err;

lib/API/MTL/MTLDevice.cpp

@@ -83,6 +83,29 @@ class MTLQueue : public offloadtest::Queue {
   }
 };
 
+class MTLFence : public offloadtest::Fence {
+public:
+  std::string Name;
+  MTL::SharedEvent *Event;
+
+  uint64_t getFenceValue() override { return Event->signaledValue(); }
+
+  llvm::Error waitForCompletion(uint64_t SignalValue) override {
+    if (!Event->waitUntilSignaledValue(SignalValue, UINT64_MAX))
+      return llvm::createStringError(std::errc::timed_out,
+                                     "Timed out waiting on shared event.");
+    return llvm::Error::success();
+  }
+
+  MTLFence(MTL::SharedEvent *Event, llvm::StringRef Name)
+      : Name(Name), Event(Event) {}
+
+  ~MTLFence() {
+    if (Event)
+      Event->release();
+  }
+};
+
 class MTLBuffer : public offloadtest::Buffer {
 public:
   MTL::Buffer *Buf;
@@ -130,6 +153,7 @@ class MTLDevice : public offloadtest::Device {
     llvm::SmallVector<MTL::Buffer *> Buffers;
     MTL::Texture *FrameBufferTexture = nullptr;
     MTL::CommandBuffer *CmdBuffer = nullptr;
+    std::shared_ptr<offloadtest::Fence> Fence;
   };
 
   llvm::Error setupVertexShader(InvocationState &IS, const Pipeline &P,
@@ -487,14 +511,23 @@ class MTLDevice : public offloadtest::Device {
   }
 
   llvm::Error executeCommands(InvocationState &IS) {
+    // This is a hack but it works since this is all single threaded code.
+    static uint64_t FenceCounter = 0;
+    const uint64_t CurrentCounter = FenceCounter + 1;
+    auto *F = static_cast<MTLFence *>(IS.Fence.get());
+
+    IS.CmdBuffer->encodeSignalEvent(F->Event, CurrentCounter);
     IS.CmdBuffer->commit();
-    IS.CmdBuffer->waitUntilCompleted();
+
+    if (auto Err = IS.Fence->waitForCompletion(CurrentCounter))
+      return Err;
 
     // Check and surface any errors that occurred during execution.
     NS::Error *CBErr = IS.CmdBuffer->error();
     if (CBErr)
       return toError(CBErr);
 
+    FenceCounter = CurrentCounter;
     return llvm::Error::success();
   }
 
@@ -564,6 +597,15 @@ class MTLDevice : public offloadtest::Device {
 
   Queue &getGraphicsQueue() override { return GraphicsQueue; }
 
+  llvm::Expected<std::shared_ptr<offloadtest::Fence>>
+  createFence(llvm::StringRef Name) override {
+    MTL::SharedEvent *Event = Device->newSharedEvent();
+    if (!Event)
+      return llvm::createStringError(std::errc::device_or_resource_busy,
+                                     "Failed to create shared event.");
+    return std::make_shared<MTLFence>(Event, Name);
+  }
+
   llvm::Expected<std::shared_ptr<offloadtest::Buffer>>
   createBuffer(std::string Name, BufferCreateDesc &Desc,
                size_t SizeInBytes) override {
@@ -588,6 +630,11 @@ class MTLDevice : public offloadtest::Device {
   llvm::Error executeProgram(Pipeline &P) override {
     InvocationState IS;
 
+    auto FenceOrErr = this->createFence("Fence");
+    if (!FenceOrErr)
+      return FenceOrErr.takeError();
+    IS.Fence = *FenceOrErr;
+
     if (auto Err = createBuffers(P, IS))
       return Err;