WGSLUtils: improved atomic buffers detection

Author: TheMostDiligent

Graphics/ShaderTools/src/WGSLUtils.cpp

@@ -112,26 +112,47 @@ std::string ConvertSPIRVtoWGSL(const std::vector<uint32_t>& SPIRV)
 
 static bool IsAtomic(const tint::core::type::Type* WGSLType)
 {
-    if (WGSLType == nullptr || !WGSLType->Is<tint::core::type::Struct>())
+    if (WGSLType == nullptr)
         return false;
 
-    const tint::core::type::Struct* WGSLStruct = WGSLType->As<tint::core::type::Struct>();
-    if (WGSLStruct == nullptr || WGSLStruct->Members().IsEmpty())
-        return false;
+    if (WGSLType->Is<tint::core::type::Struct>())
+    {
+        const tint::core::type::Struct* WGSLStruct = WGSLType->As<tint::core::type::Struct>();
+        if (WGSLStruct == nullptr || WGSLStruct->Members().IsEmpty())
+            return false;
 
-    const tint::core::type::Type* WGSLMember = WGSLStruct->Members().Front()->Type();
-    if (WGSLMember == nullptr || !WGSLMember->Is<tint::core::type::Array>())
-        return false;
+        for (const tint::core::type::StructMember* WGSLMember : WGSLStruct->Members())
+        {
+            const tint::core::type::Type* WGSLMemberType = WGSLMember->Type();
+            if (WGSLMemberType == nullptr)
+                return false;
 
-    const tint::core::type::Array* WGSLArray = WGSLMember->As<tint::core::type::Array>();
-    if (WGSLArray == nullptr)
-        return false;
+            if (WGSLMemberType->Is<tint::core::type::Array>())
+            {
+                const tint::core::type::Array* WGSLArray = WGSLMemberType->As<tint::core::type::Array>();
+                if (WGSLArray == nullptr)
+                    return false;
 
-    const tint::core::type::Type* WGSLArrayMember = WGSLArray->ElemType();
-    if (WGSLArrayMember == nullptr)
-        return false;
+                const tint::core::type::Type* WGSLArrayElemType = WGSLArray->ElemType();
+                if (WGSLArrayElemType == nullptr)
+                    return false;
+
+                if (IsAtomic(WGSLArrayElemType))
+                    return true;
+            }
+            else
+            {
+                if (IsAtomic(WGSLMemberType))
+                    return true;
+            }
+        }
 
-    return WGSLArrayMember->Is<tint::core::type::Atomic>();
+        return false;
+    }
+    else
+    {
+        return WGSLType->Is<tint::core::type::Atomic>();
+    }
 }
 
 std::string GetWGSLResourceAlternativeName(const tint::Program& Program, const tint::inspector::ResourceBinding& Binding)

Tests/DiligentCoreTest/assets/shaders/WGSL/RWStructBuffers.psh

@@ -8,6 +8,19 @@ struct BufferData1
     float4 data;
 };
 
+struct BufferData2
+{
+    int data;
+};
+
+struct BufferData3
+{
+    int data0;
+    int data1;
+    int data2;
+    int data3;
+};
+
 RWStructuredBuffer<BufferData0> g_RWBuff0;
 RWStructuredBuffer<BufferData0> g_RWBuff1;
 RWStructuredBuffer<BufferData1> g_RWBuff2;
@@ -22,6 +35,11 @@ RWStructuredBuffer<int> g_RWBuff1_atomic;
 RWStructuredBuffer<int> g_RWBuff0Atomic_atomic;
 RWStructuredBuffer<int> g_RWBuff1Atomic_atomic;
 
+RWStructuredBuffer<BufferData2> g_RWBuff2Atomic;
+RWStructuredBuffer<BufferData3> g_RWBuff3Atomic;
+RWStructuredBuffer<BufferData2> g_RWBuff4Atomic_atomic;
+RWStructuredBuffer<BufferData3> g_RWBuff5Atomic_atomic;
+
 float4 main() : SV_Target
 {
     BufferData0 Data0;
@@ -43,6 +61,11 @@ float4 main() : SV_Target
 
     InterlockedAdd(g_RWBuff0Atomic_atomic[0], 1, OriginalListIdx);
     InterlockedAdd(g_RWBuff1Atomic_atomic[0], 1, OriginalListIdx);
+
+    InterlockedAdd(g_RWBuff2Atomic[0].data, 1, OriginalListIdx);
+    InterlockedAdd(g_RWBuff3Atomic[0].data1, 1, OriginalListIdx);
+    InterlockedAdd(g_RWBuff4Atomic_atomic[0].data, 1, OriginalListIdx);
+    InterlockedAdd(g_RWBuff5Atomic_atomic[0].data3, 1, OriginalListIdx);
 
     return float4(1.0, 0.0, 0.0, 1.0);
 }

Tests/DiligentCoreTest/src/ShaderTools/WGSLShaderResourcesTest.cpp

@@ -188,6 +188,10 @@ TEST(WGSLShaderResources, RWStructBuffers)
                           {"g_RWBuff1_atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
                           {"g_RWBuff0Atomic_atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
                           {"g_RWBuff1Atomic_atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
+                          {"g_RWBuff2Atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
+                          {"g_RWBuff3Atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
+                          {"g_RWBuff4Atomic_atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
+                          {"g_RWBuff5Atomic_atomic", WGSLResourceType::RWStorageBuffer, 1, RESOURCE_DIM_BUFFER},
                       });
 }
 

Tests/DiligentCoreTest/src/ShaderTools/WGSLUtilsTest.cpp

@@ -251,6 +251,10 @@ TEST(WGSLUtils, RemapRWStructBuffers)
                               {"g_RWBuff1_atomic", {15, 16}},
                               {"g_RWBuff0Atomic_atomic", {17, 18}},
                               {"g_RWBuff1Atomic_atomic", {19, 20}},
+                              {"g_RWBuff2Atomic", {21, 22}},
+                              {"g_RWBuff3Atomic", {23, 24}},
+                              {"g_RWBuff4Atomic_atomic", {25, 26}},
+                              {"g_RWBuff5Atomic_atomic", {27, 28}},
                           });
 }