rewrite depth stencil using createTexture and align usages of depth stencil target through all backends (vk was using depth stencil, dx and metal weren't)

EmilioLaiso · EmilioLaiso · commit 6da45021c3bc · 2026-04-07T09:43:54.000+02:00
diff --git a/include/API/Device.h b/include/API/Device.h
@@ -92,8 +92,6 @@ llvm::Expected<std::shared_ptr<Texture>>
 createRenderTarget(Device &Dev, const CPUBuffer &Buf);
 
 // Creates a depth/stencil texture matching the dimensions of a render target.
-// TODO: Currently always created for graphics pipelines. Consider making this
-// conditional on the pipeline definition (e.g. a flag in Pipeline).
 llvm::Expected<std::shared_ptr<Texture>>
 createDepthStencil(Device &Dev, uint32_t Width, uint32_t Height);
 
diff --git a/lib/API/DX/Device.cpp b/lib/API/DX/Device.cpp
@@ -369,7 +369,9 @@ class DXDevice : public offloadtest::Device {
     // Resources for graphics pipelines.
     std::shared_ptr<DXTexture> RT;
     std::shared_ptr<DXBuffer> RTReadback;
+    std::shared_ptr<DXTexture> DS;
     ComPtr<ID3D12DescriptorHeap> RTVHeap;
+    ComPtr<ID3D12DescriptorHeap> DSVHeap;
     ComPtr<ID3D12Resource> VB;
 
     llvm::SmallVector<DescriptorTable> DescTables;
@@ -1482,6 +1484,16 @@ class DXDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
+  llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) {
+    auto TexOrErr = offloadtest::createDepthStencil(
+        *this, P.Bindings.RTargetBufferPtr->OutputProps.Width,
+        P.Bindings.RTargetBufferPtr->OutputProps.Height);
+    if (!TexOrErr)
+      return TexOrErr.takeError();
+    IS.DS = std::static_pointer_cast<DXTexture>(*TexOrErr);
+    return llvm::Error::success();
+  }
+
   llvm::Error createVertexBuffer(Pipeline &P, InvocationState &IS) {
     if (!P.Bindings.VertexBufferPtr)
       return llvm::createStringError(
@@ -1557,8 +1569,11 @@ class DXDevice : public offloadtest::Device {
 
     PSODesc.RasterizerState = CD3DX12_RASTERIZER_DESC(D3D12_DEFAULT);
     PSODesc.BlendState = CD3DX12_BLEND_DESC(D3D12_DEFAULT);
-    PSODesc.DepthStencilState.DepthEnable = false;
+    PSODesc.DepthStencilState.DepthEnable = true;
+    PSODesc.DepthStencilState.DepthWriteMask = D3D12_DEPTH_WRITE_MASK_ALL;
+    PSODesc.DepthStencilState.DepthFunc = D3D12_COMPARISON_FUNC_LESS;
     PSODesc.DepthStencilState.StencilEnable = false;
+    PSODesc.DSVFormat = getDXGIFormat(IS.DS->Desc.Format);
     PSODesc.SampleMask = UINT_MAX;
     PSODesc.PrimitiveTopologyType = D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE;
     PSODesc.NumRenderTargets = 1;
@@ -1590,6 +1605,19 @@ class DXDevice : public offloadtest::Device {
         IS.RTVHeap->GetCPUDescriptorHandleForHeapStart();
     Device->CreateRenderTargetView(IS.RT->Resource.Get(), nullptr, RTVHandle);
 
+    // Create descriptor heap and view for the depth/stencil buffer.
+    D3D12_DESCRIPTOR_HEAP_DESC DSVHeapDesc = {};
+    DSVHeapDesc.NumDescriptors = 1;
+    DSVHeapDesc.Type = D3D12_DESCRIPTOR_HEAP_TYPE_DSV;
+    DSVHeapDesc.Flags = D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
+    if (auto Err = HR::toError(Device->CreateDescriptorHeap(
+                                   &DSVHeapDesc, IID_PPV_ARGS(&IS.DSVHeap)),
+                               "Failed to create DSV heap"))
+      return Err;
+    const D3D12_CPU_DESCRIPTOR_HANDLE DSVHandle =
+        IS.DSVHeap->GetCPUDescriptorHandleForHeapStart();
+    Device->CreateDepthStencilView(IS.DS->Resource.Get(), nullptr, DSVHandle);
+
     IS.CmdList->SetGraphicsRootSignature(IS.RootSig.Get());
     if (IS.DescHeap) {
       ID3D12DescriptorHeap *const Heaps[] = {IS.DescHeap.Get()};
@@ -1599,7 +1627,17 @@ class DXDevice : public offloadtest::Device {
     }
     IS.CmdList->SetPipelineState(IS.PSO.Get());
 
-    IS.CmdList->OMSetRenderTargets(1, &RTVHandle, false, nullptr);
+    IS.CmdList->OMSetRenderTargets(1, &RTVHandle, false, &DSVHandle);
+
+    const auto *DepthCV =
+        std::get_if<ClearDepthStencil>(&*IS.DS->Desc.OptimizedClearValue);
+    if (!DepthCV)
+      return llvm::createStringError(
+          std::errc::invalid_argument,
+          "Depth/stencil clear value must be a ClearDepthStencil.");
+    IS.CmdList->ClearDepthStencilView(
+        DSVHandle, D3D12_CLEAR_FLAG_DEPTH | D3D12_CLEAR_FLAG_STENCIL,
+        DepthCV->Depth, DepthCV->Stencil, 0, nullptr);
 
     D3D12_VIEWPORT VP = {};
     VP.Width =
@@ -1736,10 +1774,15 @@ class DXDevice : public offloadtest::Device {
       llvm::outs() << "Compute command list created.\n";
 
     } else {
-      // Create render target, readback and vertex buffer and PSO.
+      // Create render target, depth/stencil, readback and vertex buffer and PSO.
       if (auto Err = createRenderTarget(P, State))
         return Err;
       llvm::outs() << "Render target created.\n";
+      // TODO: Always created for graphics pipelines. Consider making this
+      // conditional on the pipeline definition.
+      if (auto Err = createDepthStencil(P, State))
+        return Err;
+      llvm::outs() << "Depth stencil created.\n";
       if (auto Err = createVertexBuffer(P, State))
         return Err;
       llvm::outs() << "Vertex buffer created.\n";
diff --git a/lib/API/Device.cpp b/lib/API/Device.cpp
@@ -68,7 +68,7 @@ offloadtest::createRenderTarget(Device &Dev, const CPUBuffer &Buf) {
 }
 
 llvm::Expected<std::shared_ptr<Texture>>
-createDepthStencil(Device &Dev, uint32_t Width, uint32_t Height) {
+offloadtest::createDepthStencil(Device &Dev, uint32_t Width, uint32_t Height) {
   TextureCreateDesc Desc = {};
   Desc.Location = MemoryLocation::GpuOnly;
   Desc.Usage = TextureUsage::DepthStencil;
diff --git a/lib/API/MTL/MTLDevice.cpp b/lib/API/MTL/MTLDevice.cpp
@@ -146,6 +146,7 @@ class MTLDevice : public offloadtest::Device {
     llvm::SmallVector<MTL::Buffer *> Buffers;
     std::shared_ptr<MTLTexture> FrameBufferTexture;
     std::shared_ptr<MTLBuffer> FrameBufferReadback;
+    std::shared_ptr<MTLTexture> DepthStencil;
     MTL::CommandBuffer *CmdBuffer = nullptr;
   };
 
@@ -283,6 +284,12 @@ class MTLDevice : public offloadtest::Device {
             MTL::RenderPipelineColorAttachmentDescriptor::alloc()->init();
         RPCA->setPixelFormat(PF);
         Desc->colorAttachments()->setObject(RPCA, 0);
+
+        // Set the depth/stencil format on the pipeline descriptor.
+        const MTL::PixelFormat DepthFmt =
+            getMetalFormat(TextureFormat::D32FloatS8Uint);
+        Desc->setDepthAttachmentPixelFormat(DepthFmt);
+        Desc->setStencilAttachmentPixelFormat(DepthFmt);
       }
 
       IS.RenderPipeline = Device->newRenderPipelineState(Desc, &Error);
@@ -476,31 +483,82 @@ class MTLDevice : public offloadtest::Device {
     return llvm::Error::success();
   }
 
+  llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) {
+    auto TexOrErr = offloadtest::createDepthStencil(
+        *this, P.Bindings.RTargetBufferPtr->OutputProps.Width,
+        P.Bindings.RTargetBufferPtr->OutputProps.Height);
+    if (!TexOrErr)
+      return TexOrErr.takeError();
+    IS.DepthStencil = std::static_pointer_cast<MTLTexture>(*TexOrErr);
+    return llvm::Error::success();
+  }
+
   llvm::Error createGraphicsCommands(Pipeline &P, InvocationState &IS) {
     IS.CmdBuffer = GraphicsQueue.Queue->commandBuffer();
 
     if (auto Err = createRenderTarget(P, IS))
       return Err;
+    // TODO: Always created for graphics pipelines. Consider making this
+    // conditional on the pipeline definition.
+    if (auto Err = createDepthStencil(P, IS))
+      return Err;
 
     MTL::RenderPassDescriptor *Desc =
         MTL::RenderPassDescriptor::alloc()->init();
 
     const uint64_t Width = P.Bindings.RTargetBufferPtr->OutputProps.Width;
     const uint64_t Height = P.Bindings.RTargetBufferPtr->OutputProps.Height;
 
+    // Color attachment.
     auto *CADesc = MTL::RenderPassColorAttachmentDescriptor::alloc()->init();
     CADesc->setTexture(IS.FrameBufferTexture->Tex);
     CADesc->setLoadAction(MTL::LoadActionClear);
-    const auto &CV =
-        std::get<ClearColor>(*IS.FrameBufferTexture->Desc.OptimizedClearValue);
-    CADesc->setClearColor(MTL::ClearColor(CV.R, CV.G, CV.B, CV.A));
+    const auto *ColorCV =
+        std::get_if<ClearColor>(&*IS.FrameBufferTexture->Desc.OptimizedClearValue);
+    if (!ColorCV)
+      return llvm::createStringError(
+          std::errc::invalid_argument,
+          "Render target clear value must be a ClearColor.");
+
+    CADesc->setClearColor(MTL::ClearColor(ColorCV->R, ColorCV->G, ColorCV->B, ColorCV->A));
     CADesc->setStoreAction(MTL::StoreActionStore);
     Desc->colorAttachments()->setObject(CADesc, 0);
 
+    // Depth/stencil attachment.
+    const auto *DepthCV =
+        std::get_if<ClearDepthStencil>(&*IS.DepthStencil->Desc.OptimizedClearValue);
+    if (!DepthCV)
+      return llvm::createStringError(
+          std::errc::invalid_argument,
+          "Depth/stencil clear value must be a ClearDepthStencil.");
+          
+    auto *DADesc = Desc->depthAttachment();
+    DADesc->setTexture(IS.DepthStencil->Tex);
+    DADesc->setLoadAction(MTL::LoadActionClear);
+    DADesc->setClearDepth(DepthCV->Depth);
+    DADesc->setStoreAction(MTL::StoreActionDontCare);
+
+    auto *SADesc = Desc->stencilAttachment();
+    SADesc->setTexture(IS.DepthStencil->Tex);
+    SADesc->setLoadAction(MTL::LoadActionClear);
+    SADesc->setClearStencil(DepthCV->Stencil);
+    SADesc->setStoreAction(MTL::StoreActionDontCare);
+
     MTL::RenderCommandEncoder *CmdEncoder =
         IS.CmdBuffer->renderCommandEncoder(Desc);
 
     CmdEncoder->setRenderPipelineState(IS.RenderPipeline);
+
+    // Configure depth stencil state: depth test enabled, write all, less.
+    MTL::DepthStencilDescriptor *DSDesc =
+        MTL::DepthStencilDescriptor::alloc()->init();
+    DSDesc->setDepthCompareFunction(MTL::CompareFunctionLess);
+    DSDesc->setDepthWriteEnabled(true);
+    MTL::DepthStencilState *DSState = Device->newDepthStencilState(DSDesc);
+    CmdEncoder->setDepthStencilState(DSState);
+    DSDesc->release();
+    DSState->release();
+
     // Explicitly set viewport to texture dimensions.
     CmdEncoder->setViewport(
         MTL::Viewport{0.0, 0.0, (double)Width, (double)Height, 0.0, 1.0});
diff --git a/lib/API/VK/Device.cpp b/lib/API/VK/Device.cpp
@@ -510,7 +510,7 @@ class VulkanDevice : public offloadtest::Device {
     VkFramebuffer FrameBuffer = VK_NULL_HANDLE;
     std::shared_ptr<VulkanTexture> RenderTarget;
     std::shared_ptr<VulkanBuffer> RTReadback;
-    ImageRef DepthStencil = {0, 0, 0};
+    std::shared_ptr<VulkanTexture> DepthStencil;
     std::optional<ResourceRef> VertexBuffer = std::nullopt;
 
     VkRenderPass RenderPass = VK_NULL_HANDLE;
@@ -1118,46 +1118,12 @@ class VulkanDevice : public offloadtest::Device {
   }
 
   llvm::Error createDepthStencil(Pipeline &P, InvocationState &IS) {
-    // Create an optimal image used as the depth stencil attachment
-    VkImageCreateInfo ImageCi = {};
-    ImageCi.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
-    ImageCi.imageType = getVKImageType(ResourceKind::Texture2D);
-    ImageCi.format = VK_FORMAT_D32_SFLOAT_S8_UINT;
-    // Use example's height and width
-    ImageCi.extent = {
-        static_cast<uint32_t>(P.Bindings.RTargetBufferPtr->OutputProps.Width),
-        static_cast<uint32_t>(P.Bindings.RTargetBufferPtr->OutputProps.Height),
-        1};
-    ImageCi.mipLevels = 1;
-    ImageCi.arrayLayers = 1;
-    ImageCi.samples = VK_SAMPLE_COUNT_1_BIT;
-    ImageCi.tiling = VK_IMAGE_TILING_OPTIMAL;
-    ImageCi.usage = VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT;
-    ImageCi.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
-    if (vkCreateImage(Device, &ImageCi, nullptr, &IS.DepthStencil.Image))
-      return llvm::createStringError(std::errc::device_or_resource_busy,
-                                     "Depth stencil creation failed.");
-
-    // Allocate memory for the image (device local) and bind it to our image
-    VkMemoryAllocateInfo MemAlloc{};
-    MemAlloc.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
-    VkMemoryRequirements MemReqs;
-    vkGetImageMemoryRequirements(Device, IS.DepthStencil.Image, &MemReqs);
-    MemAlloc.allocationSize = MemReqs.size;
-    llvm::Expected<uint32_t> MemIdx =
-        getMemoryIndex(PhysicalDevice, MemReqs.memoryTypeBits,
-                       VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
-    if (!MemIdx)
-      return MemIdx.takeError();
-
-    MemAlloc.memoryTypeIndex = *MemIdx;
-    if (vkAllocateMemory(Device, &MemAlloc, nullptr, &IS.DepthStencil.Memory))
-      return llvm::createStringError(std::errc::not_enough_memory,
-                                     "Depth stencil memory allocation failed.");
-    if (vkBindImageMemory(Device, IS.DepthStencil.Image, IS.DepthStencil.Memory,
-                          0))
-      return llvm::createStringError(std::errc::not_enough_memory,
-                                     "Depth stencil memory binding failed.");
+    auto TexOrErr = offloadtest::createDepthStencil(
+        *this, P.Bindings.RTargetBufferPtr->OutputProps.Width,
+        P.Bindings.RTargetBufferPtr->OutputProps.Height);
+    if (!TexOrErr)
+      return TexOrErr.takeError();
+    IS.DepthStencil = std::static_pointer_cast<VulkanTexture>(*TexOrErr);
     return llvm::Error::success();
   }
 
@@ -1172,6 +1138,8 @@ class VulkanDevice : public offloadtest::Device {
     if (P.isGraphics()) {
       if (auto Err = createRenderTarget(P, IS))
         return Err;
+      // TODO: Always created for graphics pipelines. Consider making this
+      // conditional on the pipeline definition.
       if (auto Err = createDepthStencil(P, IS))
         return Err;
 
@@ -1655,7 +1623,7 @@ class VulkanDevice : public offloadtest::Device {
     DepthStencilViewCi.subresourceRange.levelCount = 1;
     DepthStencilViewCi.subresourceRange.baseArrayLayer = 0;
     DepthStencilViewCi.subresourceRange.layerCount = 1;
-    DepthStencilViewCi.image = IS.DepthStencil.Image;
+    DepthStencilViewCi.image = IS.DepthStencil->Image;
     if (vkCreateImageView(Device, &DepthStencilViewCi, nullptr, &Views[1]))
       return llvm::createStringError(
           std::errc::device_or_resource_busy,
@@ -2041,8 +2009,9 @@ class VulkanDevice : public offloadtest::Device {
                              VkImageLayout OldLayout,
                              VkAccessFlags SrcAccessMask,
                              VkPipelineStageFlags SrcStageMask) {
-    VkImageAspectFlags AspectMask =
-        isDepth(Tex.Desc.Format) ? VK_IMAGE_ASPECT_DEPTH_BIT : VK_IMAGE_ASPECT_COLOR_BIT;
+    VkImageAspectFlags AspectMask = isDepth(Tex.Desc.Format)
+                                        ? VK_IMAGE_ASPECT_DEPTH_BIT
+                                        : VK_IMAGE_ASPECT_COLOR_BIT;
 
     // Transition texture to transfer source.
     VkImageSubresourceRange SubRange = {};
@@ -2213,11 +2182,21 @@ class VulkanDevice : public offloadtest::Device {
       copyResourceDataToDevice(IS, R);
 
     if (P.isGraphics()) {
-      const auto &CV =
-          std::get<ClearColor>(*IS.RenderTarget->Desc.OptimizedClearValue);
+      const auto *ColorCV =
+          std::get_if<ClearColor>(&*IS.RenderTarget->Desc.OptimizedClearValue);
+      if (!ColorCV)
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Render target clear value must be a ClearColor.");
+      const auto *DepthCV = std::get_if<ClearDepthStencil>(
+          &*IS.DepthStencil->Desc.OptimizedClearValue);
+      if (!DepthCV)
+        return llvm::createStringError(
+            std::errc::invalid_argument,
+            "Depth/stencil clear value must be a ClearDepthStencil.");
       VkClearValue ClearValues[2] = {};
-      ClearValues[0].color = {{CV.R, CV.G, CV.B, CV.A}};
-      ClearValues[1].depthStencil = {1.0f, 0};
+      ClearValues[0].color = {{ColorCV->R, ColorCV->G, ColorCV->B, ColorCV->A}};
+      ClearValues[1].depthStencil = {DepthCV->Depth, DepthCV->Stencil};
 
       VkRenderPassBeginInfo RenderPassBeginInfo = {};
       RenderPassBeginInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
@@ -2399,9 +2378,8 @@ class VulkanDevice : public offloadtest::Device {
         vkFreeMemory(Device, IS.VertexBuffer->Host.Memory, nullptr);
       }
       // Render target image and readback buffer are owned by
-      // IS.RenderTarget and IS.RTReadback (shared_ptrs).
-      vkDestroyImage(Device, IS.DepthStencil.Image, nullptr);
-      vkFreeMemory(Device, IS.DepthStencil.Memory, nullptr);
+      // Render target, readback buffer, and depth stencil are owned by
+      // shared_ptrs (IS.RenderTarget, IS.RTReadback, IS.DepthStencil).
       vkDestroyFramebuffer(Device, IS.FrameBuffer, nullptr);
       vkDestroyRenderPass(Device, IS.RenderPass, nullptr);
     }

Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ offloadtest::createRenderTarget(Device &Dev, const CPUBuffer &Buf) {`
`68`	`68`	`}`
`69`	`69`
`70`	`70`	`llvm::Expected<std::shared_ptr<Texture>>`
`71`		`-createDepthStencil(Device &Dev, uint32_t Width, uint32_t Height) {`
	`71`	`+offloadtest::createDepthStencil(Device &Dev, uint32_t Width, uint32_t Height) {`
`72`	`72`	`TextureCreateDesc Desc = {};`
`73`	`73`	`Desc.Location = MemoryLocation::GpuOnly;`
`74`	`74`	`Desc.Usage = TextureUsage::DepthStencil;`