Add Mesh Shader Support to Vulkan

include/API/CommandBuffer.h

@@ -19,12 +19,22 @@
 #include "API/API.h"
 #include "API/Encoder.h"
 
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Error.h"
 
 #include <memory>
 
 namespace offloadtest {
 
+class RenderPass;
+class Texture;
+
+struct RenderPassBeginDesc {
+  RenderPass *Pass = nullptr;
+  llvm::SmallVector<Texture *, 8> ColorAttachments;
+  Texture *DepthStencil = nullptr;
+};
+
 class CommandBuffer {
   GPUAPI Kind;
 
@@ -44,6 +54,15 @@ class CommandBuffer {
         std::errc::not_supported,
         "createComputeEncoder not implemented for this backend");
   }
+
+  /// Create a render command encoder for recording draw commands.
+  virtual llvm::Expected<std::unique_ptr<RenderEncoder>>
+  createRenderEncoder(const RenderPassBeginDesc &Desc) {
+    (void)Desc;
+    return llvm::createStringError(
+        std::errc::not_supported,
+        "createRenderEncoder not implemented for this backend");
+  }
 };
 
 } // namespace offloadtest

include/API/Device.h

@@ -18,6 +18,7 @@
 #include "API/Buffer.h"
 #include "API/Capabilities.h"
 #include "API/CommandBuffer.h"
+#include "API/RenderPass.h"
 #include "API/Texture.h"
 
 #include "Support/Pipeline.h"
@@ -175,6 +176,9 @@ class Device {
   virtual llvm::Expected<std::unique_ptr<Texture>>
   createTexture(std::string Name, const TextureCreateDesc &Desc) = 0;
 
+  virtual llvm::Expected<std::unique_ptr<RenderPass>>
+  createRenderPass(const RenderPassDesc &Desc) = 0;
+
   virtual void printExtra(llvm::raw_ostream &OS) {}
 
   virtual llvm::Expected<std::unique_ptr<CommandBuffer>>

include/API/Encoder.h

@@ -20,6 +20,7 @@
 namespace offloadtest {
 
 class Buffer;
+class PipelineState;
 
 /// Base class for all command encoders. An encoder records commands into a
 /// command buffer. Call endEncoding() when done recording. Barriers are
@@ -42,12 +43,6 @@ class CommandEncoder {
   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) {}
@@ -80,6 +75,44 @@ class ComputeEncoder : public CommandEncoder {
   /// pipeline state (e.g. the shader's numthreads attribute).
   virtual llvm::Error dispatch(uint32_t GroupCountX, uint32_t GroupCountY,
                                uint32_t GroupCountZ) = 0;
+
+  /// 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;
+};
+
+struct Viewport {
+  float X = 0.0f, Y = 0.0f;
+  float Width = 0.0f, Height = 0.0f;
+  float MinDepth = 0.0f, MaxDepth = 1.0f;
+};
+
+struct ScissorRect {
+  int32_t X = 0, Y = 0;
+  uint32_t Width = 0, Height = 0;
+};
+
+class RenderEncoder : public CommandEncoder {
+public:
+  using CommandEncoder::CommandEncoder;
+
+  virtual void setViewport(const Viewport &VP) = 0;
+  virtual void setScissor(const ScissorRect &Rect) = 0;
+
+  virtual void setVertexBuffer(uint32_t Slot, Buffer *VB, size_t Offset,
+                               uint32_t Stride) = 0;
+
+  virtual llvm::Error drawInstanced(const PipelineState &PSO,
+                                    uint32_t VertexCount,
+                                    uint32_t InstanceCount,
+                                    uint32_t FirstVertex = 0,
+                                    uint32_t FirstInstance = 0) = 0;
+
+  virtual llvm::Error dispatchMesh(const PipelineState &PSO,
+                                   uint32_t GroupCountX, uint32_t GroupCountY,
+                                   uint32_t GroupCountZ) = 0;
 };
 
 } // namespace offloadtest

include/API/Enums.h

@@ -31,6 +31,19 @@ enum ShaderContainerType {
   Metal,
 };
 
+/// Action applied to an attachment when a render pass begins.
+enum class LoadAction {
+  Load,     ///< Preserve existing contents.
+  Clear,    ///< Clear to the texture's OptimizedClearValue at encoder time.
+  DontCare, ///< Contents are undefined; the driver may discard.
+};
+
+/// Action applied to an attachment when a render pass ends.
+enum class StoreAction {
+  Store,    ///< Write the rendered contents back to memory.
+  DontCare, ///< Contents may be discarded after the pass.
+};
+
 } // namespace offloadtest
 
 #endif // OFFLOADTEST_API_ENUMS_H

include/API/RenderPass.h

@@ -0,0 +1,65 @@
+//===- RenderPass.h - Offload API Render Pass -----------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Defines the RenderPass abstract base class. A RenderPass describes the
+// formats and load / store actions of one or more attachments. It carries no
+// reference to specific textures: those are bound at encoder creation time.
+// Backends that have a corresponding native object (Vulkan's VkRenderPass)
+// build it once at creation time so it can be reused across encoders.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef OFFLOADTEST_API_RENDERPASS_H
+#define OFFLOADTEST_API_RENDERPASS_H
+
+#include "API/API.h"
+#include "API/Enums.h"
+#include "API/Resources.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include <optional>
+
+namespace offloadtest {
+
+struct ColorAttachmentFormatDesc {
+  Format Fmt;
+  LoadAction Load = LoadAction::Clear;
+  StoreAction Store = StoreAction::Store;
+};
+
+struct DepthStencilAttachmentFormatDesc {
+  Format Fmt;
+  LoadAction DepthLoad = LoadAction::Clear;
+  StoreAction DepthStore = StoreAction::Store;
+  LoadAction StencilLoad = LoadAction::DontCare;
+  StoreAction StencilStore = StoreAction::DontCare;
+};
+
+struct RenderPassDesc {
+  llvm::SmallVector<ColorAttachmentFormatDesc, 8> ColorAttachments;
+  std::optional<DepthStencilAttachmentFormatDesc> DepthStencil;
+};
+
+class RenderPass {
+  GPUAPI API;
+
+public:
+  virtual ~RenderPass();
+  RenderPass(const RenderPass &) = delete;
+  RenderPass &operator=(const RenderPass &) = delete;
+
+  GPUAPI getAPI() const { return API; }
+
+protected:
+  explicit RenderPass(GPUAPI API) : API(API) {}
+};
+
+} // namespace offloadtest
+
+#endif // OFFLOADTEST_API_RENDERPASS_H

include/API/Texture.h

@@ -150,6 +150,7 @@ class Texture {
   Texture &operator=(const Texture &) = delete;
 
   GPUAPI getAPI() const { return API; }
+  virtual const TextureCreateDesc &getDesc() const = 0;
 
 protected:
   explicit Texture(GPUAPI API) : API(API) {}

include/Support/Pipeline.h

@@ -26,15 +26,25 @@
 
 namespace offloadtest {
 
-enum class Stages { Compute, Vertex, Pixel };
+enum class Stages {
+  // Compute
+  Compute,
+
+  // Traditional Raster
+  Vertex,
+  Pixel,
+
+  // Mesh Shader
+  Amplification,
+  Mesh
+};
 inline constexpr std::array AllStages = {
-    Stages::Compute,
-    Stages::Vertex,
-    Stages::Pixel,
+    Stages::Compute,       Stages::Vertex, Stages::Pixel,
+    Stages::Amplification, Stages::Mesh,
 };
 inline constexpr size_t NumStages = AllStages.size();
 
-enum class ShaderPipelineKind { Compute, TraditionalRaster };
+enum class ShaderPipelineKind { Compute, TraditionalRaster, MeshShaderRaster };
 
 enum class Rule { BufferExact, BufferFloatULP, BufferFloatEpsilon };
 
@@ -383,11 +393,11 @@ struct VertexAttribute {
 
 struct IOBindings {
   std::string VertexBuffer;
-  CPUBuffer *VertexBufferPtr;
+  CPUBuffer *VertexBufferPtr = nullptr;
   llvm::SmallVector<VertexAttribute> VertexAttributes;
 
   std::string RenderTarget;
-  CPUBuffer *RTargetBufferPtr;
+  CPUBuffer *RTargetBufferPtr = nullptr;
 
   uint32_t getVertexStride() const {
     uint32_t Stride = 0;
@@ -502,6 +512,12 @@ struct Pipeline {
   bool isTraditionalRaster() const {
     return Kind == ShaderPipelineKind::TraditionalRaster;
   }
+  bool isMeshShaderRaster() const {
+    return Kind == ShaderPipelineKind::MeshShaderRaster;
+  }
+  bool isRaster() const {
+    return isTraditionalRaster() || isMeshShaderRaster();
+  }
 };
 } // namespace offloadtest
 
@@ -712,6 +728,8 @@ template <> struct ScalarEnumerationTraits<offloadtest::Stages> {
     ENUM_CASE(Compute);
     ENUM_CASE(Vertex);
     ENUM_CASE(Pixel);
+    ENUM_CASE(Amplification);
+    ENUM_CASE(Mesh);
 #undef ENUM_CASE
   }
 };