[SPIR-V] Fix counter in direct return stmt (#8413)

There is a logic to create counter variables, and find them from a
source expression.
For return statements, we were first looking for the referenced counter
before dealing with the RHS of the expression: the call. The fact that
it worked so far was just luck:
- if the function decl was handled/seen before the return statement, the
counter variable association was handled. But in case of a direct
return, this was not the case.

Fixes #8215

Author: Keenuts

tools/clang/lib/SPIRV/SpirvEmitter.cpp

@@ -2967,13 +2967,13 @@ void SpirvEmitter::doReturnStmt(const ReturnStmt *stmt) {
       declIdMapper.createResourceHeap(decl, resourceType);
     }
 
-    // Update counter variable associated with function returns
-    tryToAssignCounterVar(curFunction, retVal);
-
     auto *retInfo = loadIfGLValue(retVal);
     if (!retInfo)
       return;
 
+    // Update counter variable associated with function returns
+    tryToAssignCounterVar(curFunction, retVal);
+
     auto retType = retVal->getType();
     if (retInfo->getLayoutRule() != SpirvLayoutRule::Void &&
         retType->isStructureType()) {

tools/clang/test/CodeGenSPIRV/sm6_6.descriptorheap.return.counter.nested.call.hlsl

@@ -0,0 +1,27 @@
+// RUN: %dxc -T cs_6_6 -E main -O3 %s -spirv | FileCheck %s
+
+RWStructuredBuffer<uint> GetBindlessResource_UIntBuffer_inner() {
+     return  ResourceDescriptorHeap[1];
+}
+
+RWStructuredBuffer<uint> GetBindlessResource_UIntBuffer_outter() {
+  return GetBindlessResource_UIntBuffer_inner();
+}
+
+RWStructuredBuffer<uint> GetBindlessResource_UIntBuffer() {
+     return  ResourceDescriptorHeap[0];
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+  RWStructuredBuffer<uint> a = GetBindlessResource_UIntBuffer();
+  RWStructuredBuffer<uint> b = GetBindlessResource_UIntBuffer_outter();
+// CHECK-DAG: [[counter_a:%[0-9]+]] = OpAccessChain %_ptr_Uniform_int %counter_var_ResourceDescriptorHeap %uint_0 %uint_0
+// CHECK-DAG: [[counter_b:%[0-9]+]] = OpAccessChain %_ptr_Uniform_int %counter_var_ResourceDescriptorHeap %uint_1 %uint_0
+
+  a.IncrementCounter();
+  b.IncrementCounter();
+
+// CHECK-DAG: {{%[0-9]+}} = OpAtomicIAdd %int [[counter_a]] %uint_1 %uint_0 %int_1
+// CHECK-DAG: {{%[0-9]+}} = OpAtomicIAdd %int [[counter_b]] %uint_1 %uint_0 %int_1
+}

tools/clang/test/CodeGenSPIRV/spirv.legal.sbuffer.usage.hlsl

@@ -87,6 +87,6 @@ float4 useAsStaticRWSBuffer() {
 // CHECK:  %paramRWSBuffer = OpFunctionParameter %_ptr_Function__ptr_Uniform_type_RWStructuredBuffer_S
 RWStructuredBuffer<S> returnRWSBuffer(RWStructuredBuffer<S> paramRWSBuffer) {
 // CHECK:     [[ptr_1:%[0-9]+]] = OpLoad %_ptr_Uniform_type_RWStructuredBuffer_S %paramRWSBuffer
-// CHECK-NEXT:               OpReturnValue [[ptr_1]]
+// CHECK:                         OpReturnValue [[ptr_1]]
     return paramRWSBuffer;
 }

tools/clang/test/CodeGenSPIRV/type.rwstructured-buffer.return.counter.hlsl

@@ -0,0 +1,28 @@
+// RUN: %dxc -T cs_6_6 -E main -O3 %s -spirv | FileCheck %s
+
+RWStructuredBuffer<uint> resource;
+
+RWStructuredBuffer<uint> get_resource_inner() {
+     return resource;
+}
+
+RWStructuredBuffer<uint> get_resource_outter() {
+  return get_resource_inner();
+}
+
+RWStructuredBuffer<uint> get_resource() {
+     return resource;
+}
+
+[numthreads(1, 1, 1)]
+void main() {
+  RWStructuredBuffer<uint> a = get_resource();
+  RWStructuredBuffer<uint> b = get_resource_outter();
+// CHECK-DAG: [[counter:%[0-9]+]] = OpAccessChain %_ptr_Uniform_int %counter_var_resource %uint_0
+
+  a.IncrementCounter();
+  b.IncrementCounter();
+
+// CHECK-DAG: {{%[0-9]+}} = OpAtomicIAdd %int [[counter]] %uint_1 %uint_0 %int_1
+// CHECK-DAG: {{%[0-9]+}} = OpAtomicIAdd %int [[counter]] %uint_1 %uint_0 %int_1
+}