From 0f5a2ce2ecf79fba79a9d20c9e3bcd4c31ee45bb Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Wed, 18 Dec 2024 15:34:16 -0800
Subject: Fix metadata of register space and varying params. (#5906)

---
 .../compiler-core/slang-artifact-associated-impl.h |  3 +
 source/slang/slang-ir-metadata.cpp                 | 70 +++++++++++++---------
 2 files changed, 45 insertions(+), 28 deletions(-)

(limited to 'source')

diff --git a/source/compiler-core/slang-artifact-associated-impl.h b/source/compiler-core/slang-artifact-associated-impl.h
index 8a89f865b..0418b2018 100644
--- a/source/compiler-core/slang-artifact-associated-impl.h
+++ b/source/compiler-core/slang-artifact-associated-impl.h
@@ -163,6 +163,9 @@ struct ShaderBindingRange
         case slang::UnorderedAccess:
         case slang::SamplerState:
         case slang::DescriptorTableSlot:
+        case slang::VaryingInput:
+        case slang::VaryingOutput:
+        case slang::SpecializationConstant:
             return true;
         default:
             return false;
diff --git a/source/slang/slang-ir-metadata.cpp b/source/slang/slang-ir-metadata.cpp
index 7f0c82471..7348c06fd 100644
--- a/source/slang/slang-ir-metadata.cpp
+++ b/source/slang/slang-ir-metadata.cpp
@@ -43,6 +43,41 @@ static void _insertBinding(
     ranges.add(newRange);
 }
 
+void collectMetadataFromInst(IRInst* param, ArtifactPostEmitMetadata& outMetadata)
+{
+    auto layoutDecoration = param->findDecoration<IRLayoutDecoration>();
+    if (!layoutDecoration)
+        return;
+
+    auto varLayout = as<IRVarLayout>(layoutDecoration->getLayout());
+    if (!varLayout)
+        return;
+
+    for (auto sizeAttr : varLayout->getTypeLayout()->getSizeAttrs())
+    {
+        auto kind = sizeAttr->getResourceKind();
+
+        // Only track resource types that we can reliably track, such as textures.
+        // Do not track individual uniforms, for example.
+        if (!ShaderBindingRange::isUsageTracked(kind))
+            continue;
+
+        if (auto offsetAttr = varLayout->findOffsetAttr(kind))
+        {
+            // Get the binding information from this attribute and insert it into the list
+            auto spaceIndex = offsetAttr->getSpace();
+            if (auto spaceAttr = varLayout->findOffsetAttr(LayoutResourceKind::RegisterSpace))
+            {
+                spaceIndex += spaceAttr->getOffset();
+            }
+            auto registerIndex = offsetAttr->getOffset();
+            auto size = sizeAttr->getSize();
+            auto count = size.isFinite() ? size.getFiniteValue() : 0;
+            _insertBinding(outMetadata.m_usedBindings, kind, spaceIndex, registerIndex, count);
+        }
+    }
+}
+
 // Collects the metadata from the provided IR module, saves it in outMetadata.
 void collectMetadata(const IRModule* irModule, ArtifactPostEmitMetadata& outMetadata)
 {
@@ -57,39 +92,18 @@ void collectMetadata(const IRModule* irModule, ArtifactPostEmitMetadata& outMeta
                 auto name = func->findDecoration<IRExportDecoration>()->getMangledName();
                 outMetadata.m_exportedFunctionMangledNames.add(name);
             }
+
+            // Collect metadata from entrypoint params.
+            for (auto param : func->getParams())
+            {
+                collectMetadataFromInst(param, outMetadata);
+            }
         }
 
         auto param = as<IRGlobalParam>(inst);
         if (!param)
             continue;
-
-        auto layoutDecoration = param->findDecoration<IRLayoutDecoration>();
-        if (!layoutDecoration)
-            continue;
-
-        auto varLayout = as<IRVarLayout>(layoutDecoration->getLayout());
-        if (!varLayout)
-            continue;
-
-        for (auto sizeAttr : varLayout->getTypeLayout()->getSizeAttrs())
-        {
-            auto kind = sizeAttr->getResourceKind();
-
-            // Only track resource types that we can reliably track, such as textures.
-            // Do not track individual uniforms, for example.
-            if (!ShaderBindingRange::isUsageTracked(kind))
-                continue;
-
-            if (auto offsetAttr = varLayout->findOffsetAttr(kind))
-            {
-                // Get the binding information from this attribute and insert it into the list
-                auto spaceIndex = offsetAttr->getSpace();
-                auto registerIndex = offsetAttr->getOffset();
-                auto size = sizeAttr->getSize();
-                auto count = size.isFinite() ? size.getFiniteValue() : 0;
-                _insertBinding(outMetadata.m_usedBindings, kind, spaceIndex, registerIndex, count);
-            }
-        }
+        collectMetadataFromInst(param, outMetadata);
     }
 }
 
-- 
cgit v1.2.3