Move executeProgram out of the DX12 backend

include/API/Buffer.h

@@ -70,14 +70,16 @@ struct BufferCreateDesc {
                             false};
   }
 
-  static BufferCreateDesc scratchBuffer() {
+  static BufferCreateDesc gpuOnlyStorage() {
     return BufferCreateDesc{MemoryLocation::GpuOnly,
                             MemoryBacking::Automatic,
                             BufferUsage::Storage,
                             BufferShaderAccessType::Raw,
                             {},
                             false};
   }
+
+  static BufferCreateDesc scratchBuffer() { return gpuOnlyStorage(); }
 };
 
 class Buffer {

include/API/CommandBuffer.h

@@ -26,6 +26,7 @@
 
 namespace offloadtest {
 
+class DescriptorPool;
 class RenderPass;
 class Texture;
 
@@ -46,6 +47,8 @@ class CommandBuffer {
 
   GPUAPI getKind() const { return Kind; }
 
+  virtual void bindPool(const DescriptorPool &Pool) {};
+
   /// Create a compute command encoder for recording dispatch commands.
   /// Barriers are automatically inserted between commands.
   virtual llvm::Expected<std::unique_ptr<ComputeEncoder>>

include/API/DX/AccelerationStructure.h

@@ -0,0 +1,49 @@
+//===- DX/AccelerationStructure.h - Offload API DX AccelStruct API --------===//
+//
+// 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_DX_ACCELERATIONSTRUCTURE_H
+#define OFFLOADTEST_API_DX_ACCELERATIONSTRUCTURE_H
+
+#include <d3d12.h>
+#include <wrl/client.h>
+
+// The windows headers define these macros which conflict with the C++ standard
+// library. Undefining them before including any LLVM C++ code prevents errors.
+#undef max
+#undef min
+
+#include "API/AccelerationStructure.h"
+
+#include <cstdint>
+#include <string>
+
+namespace offloadtest {
+using Microsoft::WRL::ComPtr;
+
+class DXAccelerationStructure : public offloadtest::AccelerationStructure {
+public:
+  ComPtr<ID3D12Resource> Resource;
+  D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle;
+
+  DXAccelerationStructure(ComPtr<ID3D12Resource> Resource,
+                          const AccelerationStructureSizes &Sizes)
+      : offloadtest::AccelerationStructure(GPUAPI::DirectX, Sizes),
+        Resource(Resource) {}
+
+  D3D12_GPU_VIRTUAL_ADDRESS getGPUVirtualAddress() const {
+    return Resource->GetGPUVirtualAddress();
+  }
+
+  static bool classof(const offloadtest::AccelerationStructure *AS) {
+    return AS->getAPI() == GPUAPI::DirectX;
+  }
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DX_ACCELERATIONSTRUCTURE_H

include/API/DX/Buffer.h

@@ -0,0 +1,73 @@
+//===- DX/Buffer.h - Offload API DX Buffer API ----------------------------===//
+//
+// 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_DX_BUFFER_H
+#define OFFLOADTEST_API_DX_BUFFER_H
+
+#include <d3d12.h>
+#include <wrl/client.h>
+
+// The windows headers define these macros which conflict with the C++ standard
+// library. Undefining them before including any LLVM C++ code prevents errors.
+#undef max
+#undef min
+
+#include "API/Buffer.h"
+
+#include <cstdint>
+#include <string>
+
+namespace offloadtest {
+using Microsoft::WRL::ComPtr;
+
+class DXBuffer : public offloadtest::Buffer {
+public:
+  ComPtr<ID3D12Resource> Buffer;
+  std::string Name;
+  BufferCreateDesc Desc;
+  size_t SizeInBytes;
+  uint64_t CounterOffsetInBytes;
+
+  // Contract: If a command on the command buffer needs a resource to be in a
+  // different state it should always transition it back to the PreferredState
+  // afterwards. The PreferredState is the state of the most common use case for
+  // that resource. This allows us to do state transitions without state
+  // tracking.
+  D3D12_RESOURCE_STATES PreferredState;
+
+  D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle;
+  D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle;
+  D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle;
+
+  DXBuffer(ComPtr<ID3D12Resource> Buffer, llvm::StringRef Name,
+           BufferCreateDesc Desc, size_t SizeInBytes,
+           uint64_t CounterOffsetInBytes, D3D12_RESOURCE_STATES PreferredState,
+           D3D12_CPU_DESCRIPTOR_HANDLE SRVHandle,
+           D3D12_CPU_DESCRIPTOR_HANDLE UAVHandle,
+           D3D12_CPU_DESCRIPTOR_HANDLE CBVHandle);
+  DXBuffer(const DXBuffer &) = delete;
+  DXBuffer(DXBuffer &&) = delete;
+  DXBuffer &operator=(const DXBuffer &) = delete;
+  DXBuffer &operator=(DXBuffer &&) = delete;
+
+  size_t getSizeInBytes() const override;
+
+  size_t querySparseTileSizeInBytes(const Device & /*Dev*/) const override;
+
+  llvm::Expected<void *> map() override;
+
+  void unmap() override;
+
+  const BufferCreateDesc &getDesc() const override;
+
+  static bool classof(const offloadtest::Buffer *B);
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DX_BUFFER_H

include/API/DX/CommandBuffer.h

@@ -0,0 +1,77 @@
+//===- DX/CommandBuffer.h - Offload API DX CommandBuffer API --------------===//
+//
+// 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_DX_COMMANDBUFFER_H
+#define OFFLOADTEST_API_DX_COMMANDBUFFER_H
+
+#include <d3d12.h>
+#include <wrl/client.h>
+
+// The windows headers define these macros which conflict with the C++ standard
+// library. Undefining them before including any LLVM C++ code prevents errors.
+#undef max
+#undef min
+
+#include "API/Buffer.h"
+#include "API/CommandBuffer.h"
+#include "API/DX/Common.h"
+#include "API/DX/Encoder.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstdint>
+#include <string>
+
+namespace offloadtest {
+using Microsoft::WRL::ComPtr;
+
+class DXDevice;
+
+class DXCommandBuffer : public offloadtest::CommandBuffer {
+public:
+  ComPtr<ID3D12CommandAllocator> Allocator;
+  ComPtr<ID3D12GraphicsCommandListX> CmdList;
+  /// Back-pointer to the owning device. Used by encoders that need access to
+  /// device-level resources (e.g. allocating AS scratch buffers).
+  DXDevice *Dev = nullptr;
+  /// Whether a UAV barrier is pending from a prior compute command.
+  bool PendingUAVBarrier = false;
+  llvm::SmallVector<D3D12_RESOURCE_BARRIER> PendingTransitions;
+  /// Buffers that must outlive command-buffer submission (e.g. AS scratch
+  /// and TLAS instance buffers used during builds).
+  llvm::SmallVector<std::unique_ptr<offloadtest::Buffer>> KeepAliveOwned;
+
+  static llvm::Expected<std::unique_ptr<DXCommandBuffer>>
+  create(ComPtr<ID3D12DeviceX> Device);
+
+  ~DXCommandBuffer() override = default;
+
+  static bool classof(const CommandBuffer *CB);
+
+  void addPendingUAVBarrier();
+  void addResourceTransition(ID3D12Resource *Resource,
+                             D3D12_RESOURCE_STATES StateBefore,
+                             D3D12_RESOURCE_STATES StateAfter);
+
+  void flushBarrier();
+
+  void bindPool(const DescriptorPool &Pool) override;
+
+  llvm::Expected<std::unique_ptr<offloadtest::ComputeEncoder>>
+  createComputeEncoder() override;
+
+  llvm::Expected<std::unique_ptr<offloadtest::RenderEncoder>>
+  createRenderEncoder(const offloadtest::RenderPassBeginDesc &Desc) override;
+
+private:
+  DXCommandBuffer() : CommandBuffer(GPUAPI::DirectX) {}
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DX_COMMANDBUFFER_H

include/API/DX/Common.h

@@ -0,0 +1,27 @@
+//===- DX/Common.h - Offload API DX API -----------------------------------===//
+//
+// 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_DX_COMMON_H
+#define OFFLOADTEST_API_DX_COMMON_H
+
+#include <d3d12.h>
+#include <dxcore.h>
+#include <wrl/client.h>
+
+// The windows headers define these macros which conflict with the C++ standard
+// library. Undefining them before including any LLVM C++ code prevents errors.
+#undef max
+#undef min
+
+namespace offloadtest {
+using ID3D12DeviceX = ID3D12Device5;
+using ID3D12GraphicsCommandListX = ID3D12GraphicsCommandList6;
+using Microsoft::WRL::ComPtr;
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DX_COMMON_H

include/API/DX/Descriptors.h

@@ -0,0 +1,130 @@
+//===- DX/Descriptors.h - Offload API DX Descriptors API ------------------===//
+//
+// 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_DX_DESCRIPTORS_H
+#define OFFLOADTEST_API_DX_DESCRIPTORS_H
+
+#include <d3d12.h>
+#include <wrl/client.h>
+
+// The windows headers define these macros which conflict with the C++ standard
+// library. Undefining them before including any LLVM C++ code prevents errors.
+#undef max
+#undef min
+
+#include "API/DX/Common.h"
+#include "API/Descriptors.h"
+
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Error.h"
+
+#include <atomic>
+
+namespace offloadtest {
+
+using Microsoft::WRL::ComPtr;
+
+struct DXDescriptorSet {
+  D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle;
+  D3D12_GPU_DESCRIPTOR_HANDLE CSUHandleGPU;
+
+  D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle;
+  D3D12_GPU_DESCRIPTOR_HANDLE SamplerHandleGPU;
+};
+
+class DXDescriptorSets : public DescriptorSets {
+public:
+  llvm::SmallVector<DXDescriptorSet> Sets;
+
+  DXDescriptorSets(llvm::SmallVector<DXDescriptorSet> Sets);
+
+  static bool classof(const DescriptorSets *S);
+};
+
+class DXDescriptorPool : public DescriptorPool {
+public:
+  ComPtr<ID3D12DescriptorHeap> CSUHeap;
+  ComPtr<ID3D12DescriptorHeap> SamplerHeap;
+
+  uint32_t CSUIncSize;
+  uint32_t SamplerIncSize;
+
+  D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle;
+  D3D12_GPU_DESCRIPTOR_HANDLE CSUHandleGPU;
+
+  D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle;
+  D3D12_GPU_DESCRIPTOR_HANDLE SamplerHandleGPU;
+
+  std::atomic<uint32_t> CSUAllocator = 0;
+  std::atomic<uint32_t> SamplerAllocator = 0;
+
+  DXDescriptorPool(ComPtr<ID3D12DescriptorHeap> CSUHeap,
+                   ComPtr<ID3D12DescriptorHeap> SamplerHeap,
+                   uint32_t CSUIncSize, uint32_t SamplerIncSize);
+
+  static llvm::Expected<std::unique_ptr<DescriptorPool>>
+  create(ComPtr<ID3D12DeviceX> Dev);
+
+  void allocateDescriptors(uint32_t Count, D3D12_CPU_DESCRIPTOR_HANDLE &CPU,
+                           D3D12_GPU_DESCRIPTOR_HANDLE &GPU);
+
+  void allocateSamplers(uint32_t Count, D3D12_CPU_DESCRIPTOR_HANDLE &CPU,
+                        D3D12_GPU_DESCRIPTOR_HANDLE &GPU);
+
+  void reset() override;
+
+  static bool classof(const DescriptorPool *P);
+};
+
+class DXDescriptorSetsBuilder : public DescriptorSetsBuilder {
+public:
+  struct SetState {
+    D3D12_CPU_DESCRIPTOR_HANDLE CSUHandle;
+    D3D12_CPU_DESCRIPTOR_HANDLE SamplerHandle;
+  };
+
+  ComPtr<ID3D12DeviceX> Dev;
+  llvm::SmallVector<DXDescriptorSet> Sets;
+  llvm::SmallVector<SetState> SetStates;
+  uint32_t CSUIncSize;
+  uint32_t SamplerIncSize;
+
+  DXDescriptorSetsBuilder(ComPtr<ID3D12DeviceX> Dev,
+                          llvm::SmallVector<DXDescriptorSet> Sets,
+                          uint32_t CSUIncSize, uint32_t SamplerIncSize);
+
+  DescriptorSetsBuilder &constant(uint32_t SetIndex,
+                                  llvm::ArrayRef<const Buffer *> B,
+                                  VKBind) override;
+
+  DescriptorSetsBuilder &
+  read(uint32_t SetIndex, llvm::ArrayRef<const Buffer *> B, VKBind) override;
+  DescriptorSetsBuilder &read(uint32_t SetIndex,
+                              llvm::ArrayRef<const Texture *> T,
+                              llvm::ArrayRef<const Sampler *> S,
+                              VKBind) override;
+  DescriptorSetsBuilder &read(uint32_t SetIndex,
+                              llvm::ArrayRef<const AccelerationStructure *> A,
+                              VKBind) override;
+
+  DescriptorSetsBuilder &
+  write(uint32_t SetIndex, llvm::ArrayRef<const Buffer *> B, VKBind) override;
+  DescriptorSetsBuilder &
+  write(uint32_t SetIndex, llvm::ArrayRef<const Texture *> T, VKBind) override;
+
+  DescriptorSetsBuilder &sampler(uint32_t SetIndex,
+                                 llvm::ArrayRef<const Sampler *> S,
+                                 VKBind) override;
+
+  std::unique_ptr<DescriptorSets> build() override;
+
+  static bool classof(const DescriptorSetsBuilder *B);
+};
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DX_DESCRIPTORS_H