From fdc17a974970559d8ff76d52c3ce40aaa056d441 Mon Sep 17 00:00:00 2001
From: Yong He <yonghe@outlook.com>
Date: Fri, 19 Jan 2024 18:02:40 -0800
Subject: Add `-fspv-reflect` support. (#3464)

* Add `-fspv-reflect` support.

Closes #3462.

* Fix.

* Fix.

* Remove use of `SPV_GOOGLE_hlsl_functionality1`.

* Fix spirv validation error.

* Fix test.

* Update typename hints.

* Update commandline options doc.

* Remove superfluous empty lines.

---------

Co-authored-by: Yong He <yhe@nvidia.com>
---
 source/slang/slang-ir-legalize-types.cpp | 57 +++++++++++++++++++++++++++-----
 1 file changed, 49 insertions(+), 8 deletions(-)

(limited to 'source/slang/slang-ir-legalize-types.cpp')

diff --git a/source/slang/slang-ir-legalize-types.cpp b/source/slang/slang-ir-legalize-types.cpp
index bd5c45ff4..6c5ab1223 100644
--- a/source/slang/slang-ir-legalize-types.cpp
+++ b/source/slang/slang-ir-legalize-types.cpp
@@ -2518,6 +2518,24 @@ static LegalVal legalizeFunc(
     return builder.build(irFunc);
 }
 
+static void cloneDecorationToVar(IRInst* srcInst, IRInst* varInst)
+{
+    for (auto decoration : srcInst->getDecorations())
+    {
+        switch (decoration->getOp())
+        {
+        case kIROp_FormatDecoration:
+        case kIROp_UserTypeNameDecoration:
+        case kIROp_SemanticDecoration:
+            cloneDecoration(decoration, varInst);
+            break;
+
+        default:
+            break;
+        }
+    }
+}
+
 static LegalVal declareSimpleVar(
     IRTypeLegalizationContext*  context,
     IROp                        op,
@@ -2601,16 +2619,17 @@ static LegalVal declareSimpleVar(
 
         if( leafVar )
         {
-            for( auto decoration : leafVar->getDecorations() )
+            cloneDecorationToVar(leafVar, irVar);
+            if (as<IRStructKey>(leafVar))
             {
-                switch( decoration->getOp() )
+                // Find the struct field and clone any decorations on the field over.
+                for (auto use = leafVar->firstUse; use; use = use->nextUse)
                 {
-                case kIROp_FormatDecoration:
-                    cloneDecoration(decoration, irVar);
-                    break;
-
-                default:
-                    break;
+                    if (auto field = as<IRStructField>(use->getUser()))
+                    {
+                        cloneDecorationToVar(field, irVar);
+                        break;
+                    }
                 }
             }
         }
@@ -3330,6 +3349,28 @@ static LegalVal declareVars(
                 tupleVal->elements.add(element);
             }
 
+            if (tupleVal->elements.getCount() == 2 &&
+                tupleVal->elements[0].key &&
+                tupleVal->elements[0].key->findDecorationImpl(kIROp_CounterBufferDecoration))
+            {
+                // If this is a lowered struct from a structured buffer type that has an atomic counter,
+                // insert decorations to each element var to associate the element buffer with the atomic buffer.
+                // This decoration is inserted to all lowered structs in the slang-ir-lower-append-consume-structured-buffer
+                // pass.
+                //
+                if (tupleVal->elements[0].val.flavor == LegalVal::Flavor::simple &&
+                    tupleVal->elements[1].val.flavor == LegalVal::Flavor::simple)
+                {
+                    auto simpleElementVar = tupleVal->elements[0].val.getSimple();
+                    auto simpleCounterVar = tupleVal->elements[1].val.getSimple();
+                    IRBuilder builder(simpleElementVar);
+                    builder.addDecoration(simpleElementVar, kIROp_CounterBufferDecoration, simpleCounterVar);
+                    // Clone decorations from leafVar to both element and counter var.
+                    cloneDecorationToVar(leafVar, simpleElementVar);
+                    cloneDecorationToVar(leafVar, simpleCounterVar);
+                }
+            }
+
             return LegalVal::tuple(tupleVal);
         }
         break;
-- 
cgit v1.2.3