Add ComputeEncoder abstraction with parallel/serial barrier modes (llvm#1100)

Introduces a generic command encoder abstraction for recording GPU
commands to a command buffer. `ComputeEncoder` provides `dispatch()` and
`barrier()` operations with two modes: `Parallel` (no automatic
barriers, caller manages synchronization) and `Serial` (auto-inserts
barriers between commands using tracked destination scope as next
source).

Each backend implements barrier tracking natively:
- Metal: `MTL::BarrierScope` accumulated on the encoder
- Vulkan: `VkPipelineStageFlags`/`VkAccessFlags` on the command buffer
- DX12: UAV barrier flag on the command buffer

VK/DX store barrier state on the command buffer (encoders hold a
back-reference) so it persists across encoder lifetimes. Metal stores it
on the encoder since each `MTL::ComputeCommandEncoder` is a separate
native object with implicit inter-encoder ordering.

Creating a parallel encoder flushes a full barrier on VK/DX to ensure
prior work is visible. endEncoding() flushes any pending barriers.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: MarijnS95

include/API/CommandBuffer.h

@@ -17,6 +17,11 @@
 #define OFFLOADTEST_API_COMMANDBUFFER_H
 
 #include "API/API.h"
+#include "API/Encoder.h"
+
+#include "llvm/Support/Error.h"
+
+#include <memory>
 
 namespace offloadtest {
 
@@ -30,6 +35,15 @@ class CommandBuffer {
   CommandBuffer &operator=(const CommandBuffer &) = delete;
 
   GPUAPI getKind() const { return Kind; }
+
+  /// Create a compute command encoder for recording dispatch commands.
+  /// Barriers are automatically inserted between commands.
+  virtual llvm::Expected<std::unique_ptr<ComputeEncoder>>
+  createComputeEncoder() {
+    return llvm::createStringError(
+        std::errc::not_supported,
+        "createComputeEncoder not implemented for this backend");
+  }
 };
 
 } // namespace offloadtest

include/API/Device.h

@@ -53,14 +53,9 @@ struct InputLayoutDesc {
   std::optional<uint32_t> InstanceStepRate;
 };
 
-struct DXBinding {
-  uint32_t Register;
-  uint32_t Space;
-};
-
 struct ResourceBindingDesc {
   ResourceKind Kind;
-  DXBinding DXBinding;
+  DirectXBinding DXBinding;
   std::optional<VulkanBinding> VKBinding;
   uint32_t DescriptorCount;
 };

include/API/Encoder.h

@@ -0,0 +1,87 @@
+//===- Encoder.h - Offload API Command Encoder Abstraction ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OFFLOADTEST_API_ENCODER_H
+#define OFFLOADTEST_API_ENCODER_H
+
+#include "API/API.h"
+
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Error.h"
+
+#include <cstddef>
+#include <cstdint>
+
+namespace offloadtest {
+
+class Buffer;
+
+/// Base class for all command encoders. An encoder records commands into a
+/// command buffer. Call endEncoding() when done recording. Barriers are
+/// automatically inserted between commands.
+class CommandEncoder {
+  GPUAPI API;
+  bool Ended = false;
+
+protected:
+  /// Backend-specific cleanup. Called exactly once, either explicitly via
+  /// endEncoding() or implicitly from the most-derived destructor.
+  virtual void endEncodingImpl() = 0;
+
+public:
+  explicit CommandEncoder(GPUAPI API) : API(API) {}
+  virtual ~CommandEncoder();
+  CommandEncoder(const CommandEncoder &) = delete;
+  CommandEncoder &operator=(const CommandEncoder &) = delete;
+
+  GPUAPI getAPI() const { return API; }
+  bool isEnded() const { return Ended; }
+
+  /// Copy \p Size bytes from \p Src at \p SrcOffset to \p Dst at
+  /// \p DstOffset.
+  virtual llvm::Error copyBufferToBuffer(Buffer &Src, size_t SrcOffset,
+                                         Buffer &Dst, size_t DstOffset,
+                                         size_t Size) = 0;
+
+  /// Begin a named debug group. Visible in GPU debuggers (PIX, RenderDoc,
+  /// Xcode). Must be balanced by a corresponding popDebugGroup() call.
+  virtual void pushDebugGroup(llvm::StringRef Label) {}
+
+  /// End the most recently pushed debug group.
+  virtual void popDebugGroup() {}
+
+  /// Insert a point-in-time debug marker.
+  virtual void insertDebugSignpost(llvm::StringRef Label) {}
+
+  /// Finish recording. No further commands may be recorded after this call.
+  /// Idempotent: safe to call more than once. If not called explicitly, the
+  /// most-derived destructor invokes it as a safeguard against leaked open
+  /// encoders.
+  void endEncoding() {
+    if (Ended)
+      return;
+    endEncodingImpl();
+    Ended = true;
+  }
+};
+
+/// Encoder for recording compute dispatch commands.
+class ComputeEncoder : public CommandEncoder {
+public:
+  using CommandEncoder::CommandEncoder;
+
+  /// Dispatch a compute grid. GroupCount specifies how many workgroups to
+  /// launch in each dimension. The workgroup size is derived from the bound
+  /// pipeline state (e.g. the shader's numthreads attribute).
+  virtual llvm::Error dispatch(uint32_t GroupCountX, uint32_t GroupCountY,
+                               uint32_t GroupCountZ) = 0;
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_ENCODER_H

lib/API/DX/Device.cpp

@@ -33,6 +33,7 @@
 
 #include "API/Capabilities.h"
 #include "API/Device.h"
+#include "API/Encoder.h"
 #include "API/FormatConversion.h"
 #include "DXFeatures.h"
 #include "Support/Pipeline.h"
@@ -43,6 +44,7 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Object/DXContainer.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/Signals.h"
 
 #include <atomic>
@@ -533,6 +535,8 @@ class DXCommandBuffer : public offloadtest::CommandBuffer {
 public:
   ComPtr<ID3D12CommandAllocator> Allocator;
   ComPtr<ID3D12GraphicsCommandList> CmdList;
+  /// Whether a UAV barrier is pending from a prior compute command.
+  bool PendingUAVBarrier = false;
 
   static llvm::Expected<std::unique_ptr<DXCommandBuffer>>
   create(ComPtr<ID3D12Device> Device) {
@@ -557,6 +561,20 @@ class DXCommandBuffer : public offloadtest::CommandBuffer {
     return CB->getKind() == GPUAPI::DirectX;
   }
 
+  void addPendingUAVBarrier() { PendingUAVBarrier = true; }
+
+  void flushBarrier() {
+    if (!PendingUAVBarrier)
+      return;
+    const D3D12_RESOURCE_BARRIER Barrier =
+        CD3DX12_RESOURCE_BARRIER::UAV(nullptr);
+    CmdList->ResourceBarrier(1, &Barrier);
+    PendingUAVBarrier = false;
+  }
+
+  llvm::Expected<std::unique_ptr<offloadtest::ComputeEncoder>>
+  createComputeEncoder() override;
+
 private:
   DXCommandBuffer() : CommandBuffer(GPUAPI::DirectX) {}
 };
@@ -606,6 +624,63 @@ struct DescriptorAllocator {
       : Heap(Heap), DescIncSize(DescIncSize), Capacity(Capacity) {}
 };
 
+class DXComputeEncoder : public offloadtest::ComputeEncoder {
+  DXCommandBuffer &CB;
+
+  void addUAVBarrier() {
+    CB.addPendingUAVBarrier();
+    CB.flushBarrier();
+  }
+
+public:
+  DXComputeEncoder(DXCommandBuffer &CB)
+      : ComputeEncoder(GPUAPI::DirectX), CB(CB) {}
+
+  ~DXComputeEncoder() override { endEncoding(); }
+
+  static bool classof(const CommandEncoder *E) {
+    return E->getAPI() == GPUAPI::DirectX;
+  }
+
+  // D3D12 debug labels require WinPixEventRuntime for the proper event
+  // encoding.  Without it, BeginEvent/EndEvent/SetMarker with metadata type 0
+  // crash the D3D12 debug layer, so leave these as no-ops for now.
+  void pushDebugGroup(llvm::StringRef Label) override {}
+  void popDebugGroup() override {}
+  void insertDebugSignpost(llvm::StringRef Label) override {}
+
+  llvm::Error dispatch(uint32_t GroupCountX, uint32_t GroupCountY,
+                       uint32_t GroupCountZ) override {
+    addUAVBarrier();
+    insertDebugSignpost(llvm::formatv("Dispatch [{0},{1},{2}]", GroupCountX,
+                                      GroupCountY, GroupCountZ)
+                            .str());
+    CB.CmdList->Dispatch(GroupCountX, GroupCountY, GroupCountZ);
+    return llvm::Error::success();
+  }
+
+  llvm::Error copyBufferToBuffer(offloadtest::Buffer &Src, size_t SrcOffset,
+                                 offloadtest::Buffer &Dst, size_t DstOffset,
+                                 size_t Size) override {
+    auto &DXSrc = static_cast<DXBuffer &>(Src);
+    auto &DXDst = static_cast<DXBuffer &>(Dst);
+    addUAVBarrier();
+    insertDebugSignpost(llvm::formatv("CopyBuffer {0}B", Size).str());
+    CB.CmdList->CopyBufferRegion(DXDst.Buffer.Get(), DstOffset,
+                                 DXSrc.Buffer.Get(), SrcOffset, Size);
+    return llvm::Error::success();
+  }
+
+  void endEncodingImpl() override { popDebugGroup(); }
+};
+
+llvm::Expected<std::unique_ptr<offloadtest::ComputeEncoder>>
+DXCommandBuffer::createComputeEncoder() {
+  auto Enc = std::make_unique<DXComputeEncoder>(*this);
+  Enc->pushDebugGroup("ComputeEncoder");
+  return Enc;
+}
+
 class DXDevice : public offloadtest::Device {
 private:
   ComPtr<IDXCoreAdapter> Adapter;
@@ -1594,25 +1669,17 @@ class DXDevice : public offloadtest::Device {
   }
 
   void addUploadBeginBarrier(InvocationState &IS, ComPtr<ID3D12Resource> R) {
-    const D3D12_RESOURCE_BARRIER Barrier = {
-        D3D12_RESOURCE_BARRIER_TYPE_TRANSITION,
-        D3D12_RESOURCE_BARRIER_FLAG_NONE,
-        {D3D12_RESOURCE_TRANSITION_BARRIER{
-            R.Get(), D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES,
-            D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_STATE_COPY_DEST}}};
+    const D3D12_RESOURCE_BARRIER Barrier = CD3DX12_RESOURCE_BARRIER::Transition(
+        R.Get(), D3D12_RESOURCE_STATE_COMMON, D3D12_RESOURCE_STATE_COPY_DEST);
     IS.CB->CmdList->ResourceBarrier(1, &Barrier);
   }
 
   void addUploadEndBarrier(InvocationState &IS, ComPtr<ID3D12Resource> R,
                            bool IsUAV) {
-    const D3D12_RESOURCE_BARRIER Barrier = {
-        D3D12_RESOURCE_BARRIER_TYPE_TRANSITION,
-        D3D12_RESOURCE_BARRIER_FLAG_NONE,
-        {D3D12_RESOURCE_TRANSITION_BARRIER{
-            R.Get(), D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES,
-            D3D12_RESOURCE_STATE_COPY_DEST,
-            IsUAV ? D3D12_RESOURCE_STATE_UNORDERED_ACCESS
-                  : D3D12_RESOURCE_STATE_GENERIC_READ}}};
+    const D3D12_RESOURCE_BARRIER Barrier = CD3DX12_RESOURCE_BARRIER::Transition(
+        R.Get(), D3D12_RESOURCE_STATE_COPY_DEST,
+        IsUAV ? D3D12_RESOURCE_STATE_UNORDERED_ACCESS
+              : D3D12_RESOURCE_STATE_GENERIC_READ);
     IS.CB->CmdList->ResourceBarrier(1, &Barrier);
   }
 
@@ -1715,10 +1782,18 @@ class DXDevice : public offloadtest::Device {
       }
     }
 
-    const llvm::ArrayRef<int> DispatchSize =
-        llvm::ArrayRef<int>(P.Shaders[0].DispatchSize);
-
-    IS.CB->CmdList->Dispatch(DispatchSize[0], DispatchSize[1], DispatchSize[2]);
+    {
+      auto EncoderOrErr = IS.CB->createComputeEncoder();
+      if (!EncoderOrErr)
+        return EncoderOrErr.takeError();
+      auto &Encoder = *EncoderOrErr.get();
+      const llvm::ArrayRef<int> DispatchSize =
+          llvm::ArrayRef<int>(P.Shaders[0].DispatchSize);
+      if (auto Err = Encoder.dispatch(DispatchSize[0], DispatchSize[1],
+                                      DispatchSize[2]))
+        return Err;
+      Encoder.endEncoding();
+    }
 
     auto CopyBackResource = [&IS, this](ResourcePair &R) {
       if (R.first->isTexture()) {

lib/API/Device.cpp

@@ -10,6 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "API/Device.h"
+#include "API/Encoder.h"
 #include "API/FormatConversion.h"
 
 #include "Config.h"
@@ -21,6 +22,8 @@
 
 using namespace offloadtest;
 
+CommandEncoder::~CommandEncoder() {}
+
 Buffer::~Buffer() {}
 
 CommandBuffer::~CommandBuffer() {}