Add an API generic way to build descriptor sets with a DX12 implementation

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/Descriptors.h

@@ -0,0 +1,92 @@
+//===- Descriptors.h - Offload API Buffer ---------------------------------===//
+//
+// 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_DESCRIPTORS_H
+#define OFFLOADTEST_API_DESCRIPTORS_H
+
+#include "API/API.h"
+
+#include "llvm/ADT/ArrayRef.h"
+
+namespace offloadtest {
+
+class AccelerationStructure;
+class Buffer;
+class Sampler;
+class Texture;
+
+struct VKBind {
+  uint32_t Binding;
+  uint32_t CounterBinding;
+};
+
+class DescriptorPool {
+  GPUAPI API;
+
+public:
+  virtual ~DescriptorPool() {} // TOOD(manon): Move this def to a .cpp file
+  GPUAPI getAPI() const { return API; }
+
+  virtual void reset() = 0;
+
+protected:
+  explicit DescriptorPool(GPUAPI API) : API(API) {}
+};
+
+class DescriptorSets {
+  GPUAPI API;
+
+public:
+  virtual ~DescriptorSets() {} // TOOD(manon): Move this def to a .cpp file
+
+  GPUAPI getAPI() const { return API; }
+
+protected:
+  explicit DescriptorSets(GPUAPI API) : API(API) {}
+};
+
+class DescriptorSetsBuilder {
+  GPUAPI API;
+
+public:
+  virtual ~DescriptorSetsBuilder() {} // TODO(manon): move to .cpp file
+
+  GPUAPI getAPI() const { return API; }
+
+  virtual DescriptorSetsBuilder &
+  constant(uint32_t SetIndex, llvm::ArrayRef<const Buffer *>, VKBind) = 0;
+
+  virtual DescriptorSetsBuilder &
+  read(uint32_t SetIndex, llvm::ArrayRef<const Buffer *>, VKBind) = 0;
+  virtual DescriptorSetsBuilder &read(uint32_t SetIndex,
+                                      llvm::ArrayRef<const Texture *>,
+                                      llvm::ArrayRef<const Sampler *>,
+                                      VKBind) = 0;
+  virtual DescriptorSetsBuilder &
+  read(uint32_t SetIndex, llvm::ArrayRef<const AccelerationStructure *>,
+       VKBind) = 0;
+
+  virtual DescriptorSetsBuilder &
+  write(uint32_t SetIndex, llvm::ArrayRef<const Buffer *>, VKBind) = 0;
+  virtual DescriptorSetsBuilder &
+  write(uint32_t SetIndex, llvm::ArrayRef<const Texture *>, VKBind) = 0;
+
+  virtual DescriptorSetsBuilder &
+  sampler(uint32_t SetIndex, llvm::ArrayRef<const Sampler *>, VKBind) = 0;
+
+  virtual std::unique_ptr<DescriptorSets> build() = 0;
+
+protected:
+  explicit DescriptorSetsBuilder(GPUAPI API) : API(API) {}
+};
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_DESCRIPTORS_H

include/API/Device.h

@@ -19,7 +19,9 @@
 #include "API/Buffer.h"
 #include "API/Capabilities.h"
 #include "API/CommandBuffer.h"
+#include "API/Descriptors.h"
 #include "API/RenderPass.h"
+#include "API/Sampler.h"
 #include "API/ShaderBindingTable.h"
 #include "API/Texture.h"
 
@@ -346,6 +348,9 @@ class Device {
   virtual llvm::Expected<std::unique_ptr<Texture>>
   createTexture(std::string Name, const TextureCreateDesc &Desc) = 0;
 
+  virtual llvm::Expected<std::unique_ptr<Sampler>>
+  createSampler(std::string Name, const SamplerCreateDesc &Desc) = 0;
+
   virtual llvm::Expected<std::unique_ptr<MemoryHeap>>
   createMemoryHeap(std::string Name, size_t SizeInBytes) = 0;
 
@@ -354,6 +359,12 @@ class Device {
   virtual uint32_t
   getTextureUploadRowStrideInBytes(const TextureCreateDesc &Desc) const = 0;
 
+  // The layout an upload buffer must have to feed createTextureWithData /
+  // copyBufferToTexture for the given texture description. Encodes per-mip
+  // offsets, row pitch, and total size in the backend's required alignment.
+  virtual TextureUploadLayout
+  getTextureUploadLayout(const TextureCreateDesc &Desc) const = 0;
+
   virtual llvm::Expected<std::unique_ptr<RenderPass>>
   createRenderPass(const RenderPassDesc &Desc) = 0;
 
@@ -382,7 +393,14 @@ class Device {
   createBLAS(const AccelerationStructureSizes &Sizes) = 0;
 
   virtual llvm::Expected<std::unique_ptr<AccelerationStructure>>
-  createTLAS(const AccelerationStructureSizes &Sizes) = 0;
+  createTLAS(const AccelerationStructureSizes &Sizes,
+             uint32_t InstanceCount) = 0;
+
+  virtual llvm::Expected<std::unique_ptr<DescriptorPool>>
+  createDescriptorPool() = 0;
+  virtual llvm::Expected<std::unique_ptr<DescriptorSetsBuilder>>
+  createDescriptorSetsBuilder(DescriptorPool &Pool,
+                              const PipelineState &Pipeline) = 0;
 
   virtual ~Device() = 0;
 
@@ -421,6 +439,9 @@ createBufferWithData(Device &Dev, std::string Name,
                      size_t SizeInBytes, ComputeEncoder *Encoder,
                      std::unique_ptr<offloadtest::Buffer> *OutUploadBuffer);
 
+// Create a texture and upload `Data` (tightly-packed across mip levels) into
+// it via a staging buffer recorded on `Encoder`. The staging buffer is handed
+// back through `OutUploadBuffer` and must outlive command-buffer submission.
 llvm::Expected<std::unique_ptr<offloadtest::Texture>>
 createTextureWithData(Device &Dev, std::string Name,
                       const TextureCreateDesc &Desc, const void *Data,

include/API/Encoder.h

@@ -21,12 +21,13 @@
 
 namespace offloadtest {
 
+class AccelerationStructure;
 class Buffer;
-class Texture;
+class DescriptorSets;
 class PipelineState;
-class Texture;
-class AccelerationStructure;
 class ShaderBindingTable;
+class Texture;
+
 struct BLASBuildRequest;
 struct TLASBuildRequest;
 
@@ -87,6 +88,12 @@ class ComputeEncoder : public CommandEncoder {
 public:
   using CommandEncoder::CommandEncoder;
 
+  virtual void bindComputeDescriptorSets(const PipelineState &PSO,
+                                         const DescriptorSets &DSets) = 0;
+
+  virtual void bindRayTracingDescriptorSets(const PipelineState &PSO,
+                                            const DescriptorSets &DSets) = 0;
+
   /// Dispatch a compute grid. GroupCount specifies how many workgroups to
   /// launch in each dimension. The workgroup size is derived from \p PSO
   /// (e.g. the shader's numthreads attribute), which is also bound for the
@@ -150,6 +157,9 @@ class RenderEncoder : public CommandEncoder {
   virtual void setVertexBuffer(uint32_t Slot, Buffer *VB, size_t Offset,
                                uint32_t Stride) = 0;
 
+  virtual void bindDescriptorSets(const PipelineState &PSO,
+                                  const DescriptorSets &DSets) = 0;
+
   virtual llvm::Error drawInstanced(const PipelineState &PSO,
                                     uint32_t VertexCount,
                                     uint32_t InstanceCount,

include/API/Sampler.h

@@ -0,0 +1,67 @@
+//===- Sampler.h - Offload API Texture ------------------------------------===//
+//
+// 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_SAMPLER_H
+#define OFFLOADTEST_API_SAMPLER_H
+
+#include "API/API.h"
+#include "API/Resources.h"
+
+namespace offloadtest {
+
+enum class FilterMode { Nearest, Linear };
+
+enum class AddressMode { Clamp, Repeat, Mirror, Border, MirrorOnce };
+
+enum class CompareFunction {
+  Never,
+  Less,
+  Equal,
+  LessEqual,
+  Greater,
+  NotEqual,
+  GreaterEqual,
+  Always
+};
+
+enum class SamplerKind { Sampler, SamplerComparison };
+
+struct SamplerCreateDesc {
+  FilterMode MinFilter = FilterMode::Linear;
+  FilterMode MagFilter = FilterMode::Linear;
+  AddressMode Address = AddressMode::Clamp;
+  float MinLOD = 0.0f;
+  float MaxLOD = std::numeric_limits<float>::max();
+  float MipLODBias = 0.0f;
+  CompareFunction ComparisonOp = CompareFunction::Never;
+  SamplerKind Kind = SamplerKind::Sampler;
+};
+
+class Sampler {
+  GPUAPI API;
+
+public:
+  virtual ~Sampler();
+  Sampler(const Sampler &) = delete;
+  // Sampler(Sampler &&) = delete;
+  Sampler &operator=(const Sampler &) = delete;
+  // Sampler &operator=(Sampler &&) = delete;
+
+  GPUAPI getAPI() const { return API; }
+  virtual const SamplerCreateDesc &getDesc() const = 0;
+
+protected:
+  explicit Sampler(GPUAPI API) : API(API) {}
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_SAMPLER_H

include/API/Texture.h

@@ -16,6 +16,7 @@
 #include "API/Resources.h"
 
 #include "llvm/ADT/BitmaskEnum.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Error.h"
@@ -105,12 +106,6 @@ inline llvm::Error validateTextureCreateDesc(const TextureCreateDesc &Desc) {
         std::errc::not_supported,
         "DepthStencil combined with Storage is not yet supported.");
 
-  // Depth formats require DepthStencil usage; non-depth formats forbid it.
-  if (IsDepth && !IsDS)
-    return llvm::createStringError(
-        std::errc::invalid_argument,
-        "Depth format '%s' requires DepthStencil usage.",
-        getFormatName(Desc.Fmt).data());
   if (!IsDepth && IsDS)
     return llvm::createStringError(
         std::errc::invalid_argument,
@@ -154,6 +149,24 @@ struct TileShape {
   uint32_t Depth = 1;
 };
 
+struct SubresourceFootprint {
+  uint64_t Offset = 0; // Byte offset of this subresource in the buffer.
+  uint32_t RowPitchInBytes = 0; // Destination row stride (may include padding).
+  uint32_t RowSizeInBytes = 0;  // Tightly-packed bytes per row to copy.
+  uint32_t NumRows = 0;         // Number of rows in this subresource.
+};
+
+struct TextureUploadLayout {
+  llvm::SmallVector<SubresourceFootprint> Subresources; // One entry per mip.
+  uint64_t TotalSizeInBytes = 0;
+};
+
+// Compute a tightly-packed upload layout (no row or subresource padding) for
+// the given texture description. Suitable for backends whose buffer-to-texture
+// copy consumes a tightly-packed staging buffer (e.g. Vulkan, Metal).
+TextureUploadLayout
+computeTightTextureUploadLayout(const TextureCreateDesc &Desc);
+
 class Texture {
   GPUAPI API;