From 0b0db2ebacaf82c8937278bbf650b6e737a558fb Mon Sep 17 00:00:00 2001
From: clach <clach@users.noreply.github.com>
Date: Sun, 20 Feb 2022 22:51:00 -0800
Subject: [PATCH] [hdx, hgi, hgiGL, hgiVulkan] For hgi backends such as Metal
 that do not support depth stencil readback, only read the depth buffer in
 HdxPickTask.

Contribution: Jon Creighton, David Yu

(Internal change: 2217079)
---
 pxr/imaging/hdx/pickTask.cpp           | 35 ++++++++++++++++----------
 pxr/imaging/hdx/pickTask.h             |  1 +
 pxr/imaging/hgi/enums.h                |  5 +++-
 pxr/imaging/hgiGL/capabilities.cpp     |  2 ++
 pxr/imaging/hgiVulkan/capabilities.cpp |  1 +
 5 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/pxr/imaging/hdx/pickTask.cpp b/pxr/imaging/hdx/pickTask.cpp
index 156a9d711a..93717f50d8 100644
--- a/pxr/imaging/hdx/pickTask.cpp
+++ b/pxr/imaging/hdx/pickTask.cpp
@@ -84,16 +84,6 @@ _IsStormRenderer(HdRenderDelegate *renderDelegate)
     return true;
 }
 
-// Generated renderbuffers
-static TfTokenVector _aovOutputs {
-    HdAovTokens->primId,
-    HdAovTokens->instanceId,
-    HdAovTokens->elementId,
-    HdAovTokens->edgeId,
-    HdAovTokens->pointId,
-    HdAovTokens->Neye,
-    HdAovTokens->depthStencil
-};
 TF_DEFINE_PRIVATE_TOKENS(_tokens,
     (widgetDepthStencil)
 );
@@ -116,6 +106,7 @@ HdxPickTask::HdxPickTask(HdSceneDelegate* delegate, SdfPath const& id)
     , _index(nullptr)
     , _hgi(nullptr)
     , _pickableDepthIndex(0)
+    , _depthToken(HdAovTokens->depthStencil)
 {
 }
 
@@ -172,11 +163,28 @@ HdxPickTask::_CreateAovBindings()
         std::static_pointer_cast<HdStResourceRegistry>(
             _index->GetResourceRegistry());
 
-    HdRenderDelegate *renderDelegate = _index->GetRenderDelegate();
+    HdRenderDelegate const * renderDelegate = _index->GetRenderDelegate();
 
     GfVec3i dimensions(_contextParams.resolution[0],
                        _contextParams.resolution[1], 1);
 
+    bool const stencilReadback = _hgi->GetCapabilities()->
+        IsSet(HgiDeviceCapabilitiesBitsStencilReadback);
+
+    _depthToken = stencilReadback ? HdAovTokens->depthStencil
+                                  : HdAovTokens->depth;
+
+    // Generated renderbuffers
+    TfTokenVector const _aovOutputs {
+        HdAovTokens->primId,
+        HdAovTokens->instanceId,
+        HdAovTokens->elementId,
+        HdAovTokens->edgeId,
+        HdAovTokens->pointId,
+        HdAovTokens->Neye,
+        _depthToken
+    };
+
     // Add the new renderbuffers.
     for (size_t i = 0; i < _aovOutputs.size(); ++i) {
         TfToken const & aovOutput = _aovOutputs[i];
@@ -201,7 +209,8 @@ HdxPickTask::_CreateAovBindings()
 
         _pickableAovBindings.push_back(binding);
 
-        if (HdAovHasDepthStencilSemantic(aovOutput)) {
+        if (HdAovHasDepthSemantic(aovOutput) ||
+            HdAovHasDepthStencilSemantic(aovOutput)) {
             _pickableDepthIndex = i;
             _occluderAovBinding = binding;
         }
@@ -661,7 +670,7 @@ HdxPickTask::Execute(HdTaskContext* ctx)
 
     std::vector<uint8_t> depths;
     float const * depthPtr =
-        _ReadAovBuffer<float>(HdAovTokens->depthStencil, &depths);
+        _ReadAovBuffer<float>(_depthToken, &depths);
 
     // For un-projection, get the current depth range.
     GLfloat p[2];
diff --git a/pxr/imaging/hdx/pickTask.h b/pxr/imaging/hdx/pickTask.h
index 0fb33071c1..d0e57f3d00 100644
--- a/pxr/imaging/hdx/pickTask.h
+++ b/pxr/imaging/hdx/pickTask.h
@@ -260,6 +260,7 @@ class HdxPickTask : public HdTask
     HdRenderPassAovBindingVector _pickableAovBindings;
     HdRenderPassAovBinding _occluderAovBinding;
     size_t _pickableDepthIndex;
+    TfToken _depthToken;
     std::unique_ptr<HdStRenderBuffer> _widgetDepthStencilBuffer;
     HdRenderPassAovBindingVector _widgetAovBindings;
 
diff --git a/pxr/imaging/hgi/enums.h b/pxr/imaging/hgi/enums.h
index 510136b56b..b8a3b1b12a 100644
--- a/pxr/imaging/hgi/enums.h
+++ b/pxr/imaging/hgi/enums.h
@@ -63,6 +63,8 @@ using HgiBits = uint32_t;
 ///   Use CPP padding for shader language structures</li>
 /// <li>HgiDeviceCapabilitiesBitsConservativeRaster:
 ///   The device supports conservative rasterization</li>
+/// <li>HgiDeviceCapabilitiesBitsStencilReadback:
+///   Supports reading back the stencil buffer from GPU to CPU.</li>
 /// </ul>
 ///
 enum HgiDeviceCapabilitiesBits : HgiBits
@@ -78,7 +80,8 @@ enum HgiDeviceCapabilitiesBits : HgiBits
     HgiDeviceCapabilitiesBitsShaderDoublePrecision   = 1 << 8,
     HgiDeviceCapabilitiesBitsDepthRangeMinusOnetoOne = 1 << 9,
     HgiDeviceCapabilitiesBitsCppShaderPadding        = 1 << 10,
-    HgiDeviceCapabilitiesBitsConservativeRaster      = 1 << 11
+    HgiDeviceCapabilitiesBitsConservativeRaster      = 1 << 11,
+    HgiDeviceCapabilitiesBitsStencilReadback         = 1 << 12
 };
 
 using HgiDeviceCapabilities = HgiBits;
diff --git a/pxr/imaging/hgiGL/capabilities.cpp b/pxr/imaging/hgiGL/capabilities.cpp
index 9a08767db5..b353003e40 100644
--- a/pxr/imaging/hgiGL/capabilities.cpp
+++ b/pxr/imaging/hgiGL/capabilities.cpp
@@ -195,6 +195,8 @@ HgiGLCapabilities::_LoadCapabilities()
         true);
     _SetFlag(HgiDeviceCapabilitiesBitsConservativeRaster, 
         conservativeRasterEnabled);
+    _SetFlag(HgiDeviceCapabilitiesBitsStencilReadback,
+        true);
 
     if (TfDebug::IsEnabled(HGI_DEBUG_DEVICE_CAPABILITIES)) {
         std::cout
diff --git a/pxr/imaging/hgiVulkan/capabilities.cpp b/pxr/imaging/hgiVulkan/capabilities.cpp
index 751b559826..5e2ac024c2 100755
--- a/pxr/imaging/hgiVulkan/capabilities.cpp
+++ b/pxr/imaging/hgiVulkan/capabilities.cpp
@@ -87,6 +87,7 @@ HgiVulkanCapabilities::HgiVulkanCapabilities(HgiVulkanDevice* device)
     _SetFlag(HgiDeviceCapabilitiesBitsDepthRangeMinusOnetoOne, false);
     _SetFlag(HgiDeviceCapabilitiesBitsConservativeRaster, 
         conservativeRasterEnabled);
+    _SetFlag(HgiDeviceCapabilitiesBitsStencilReadback, true);
 }
 
 HgiVulkanCapabilities::~HgiVulkanCapabilities() = default;