From 388a291430c0a7f7fbd241811d40fe5cf0916c63 Mon Sep 17 00:00:00 2001
From: Steve Urquhart <steve.urquhart@ntd.nintendo.com>
Date: Mon, 1 Dec 2025 18:21:00 -0500
Subject: [PATCH] [SPIR-V] Create CounterVar for ResourceDescriptorHeap when
 directly returned

---
 tools/clang/lib/SPIRV/SpirvEmitter.cpp           | 11 +++++++++++
 .../sm6_6.descriptorheap.return.counter.hlsl     | 16 ++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.return.counter.hlsl

diff --git a/tools/clang/lib/SPIRV/SpirvEmitter.cpp b/tools/clang/lib/SPIRV/SpirvEmitter.cpp
index 0251e4eb5e..89b4291989 100644
--- a/tools/clang/lib/SPIRV/SpirvEmitter.cpp
+++ b/tools/clang/lib/SPIRV/SpirvEmitter.cpp
@@ -2926,6 +2926,17 @@ void SpirvEmitter::doReturnStmt(const ReturnStmt *stmt) {
   bool returnsVoid = curFunction->getReturnType().getTypePtr()->isVoidType();
   if (!returnsVoid) {
     assert(retVal);
+    const Expr *srcExpr = retVal->IgnoreParenCasts();
+    if (isDescriptorHeap(srcExpr)) {
+      const Expr *base = nullptr;
+      getDescriptorHeapOperands(srcExpr, &base, /* index= */ nullptr);
+      const Expr *parentExpr = cast<CastExpr>(parentMap->getParent(srcExpr));
+      QualType resourceType = parentExpr->getType();
+      const auto *declRefExpr = dyn_cast<DeclRefExpr>(base->IgnoreCasts());
+      auto *decl = cast<VarDecl>(declRefExpr->getDecl());
+      declIdMapper.createResourceHeap(decl, resourceType);
+    }
+
     // Update counter variable associated with function returns
     tryToAssignCounterVar(curFunction, retVal);
 
diff --git a/tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.return.counter.hlsl b/tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.return.counter.hlsl
new file mode 100644
index 0000000000..fbd4c190bf
--- /dev/null
+++ b/tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.return.counter.hlsl
@@ -0,0 +1,16 @@
+// RUN: %dxc -T cs_6_6 -E main -fcgl %s -spirv | FileCheck %s
+
+RWStructuredBuffer<uint> GetBindlessResource_UIntBuffer() {
+     return  ResourceDescriptorHeap[0];
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap DescriptorSet 0
+// CHECK-DAG: OpDecorate %ResourceDescriptorHeap Binding 0
+// CHECK-DAG: OpDecorate %counter_var_ResourceDescriptorHeap DescriptorSet 0
+// CHECK-DAG: OpDecorate %counter_var_ResourceDescriptorHeap Binding 1
+  RWStructuredBuffer<uint> a = GetBindlessResource_UIntBuffer();
+
+  a.IncrementCounter();
+}