Add `set` to spirv_instruction (#2597)

5 files changed, 66 insertions, 13 deletions

diff --git a/source/slang/core.meta.slang b/source/slang/core.meta.slang
index edb98b3c2..97b9a227d 100644
--- a/source/slang/core.meta.slang
+++ b/source/slang/core.meta.slang
@@ -2801,7 +2801,7 @@ __attributeTarget(VarDeclBase)
 attribute_syntax [vk_index(index : int)] : GLSLIndexAttribute;
 
 __attributeTarget(FuncDecl)
-attribute_syntax [vk_spirv_instruction(op : int)]                   : SPIRVInstructionOpAttribute;
+attribute_syntax [vk_spirv_instruction(op : int, set : String = "")]     : SPIRVInstructionOpAttribute;
 
 // Statement Attributes
 
diff --git a/source/slang/slang-check-modifier.cpp b/source/slang/slang-check-modifier.cpp
index f3623f19f..e73f04301 100644
--- a/source/slang/slang-check-modifier.cpp
+++ b/source/slang/slang-check-modifier.cpp
@@ -438,8 +438,24 @@ namespace Slang
                 getSink()->diagnose(attr, Diagnostics::expectedSingleStringArg, attr->keywordName);
             }
         }
-        else if (as<OutputControlPointsAttribute>(attr) ||
-            as<SPIRVInstructionOpAttribute>(attr))
+        else if (auto opAttr = as<SPIRVInstructionOpAttribute>(attr))
+        {
+            auto sink = getSink();
+            const auto argsCount = opAttr->args.getCount();
+            if (argsCount < 1 || argsCount > 2)
+            {
+                sink->diagnose(attr, Diagnostics::attributeArgumentCountMismatch, attr->keywordName, "1...2", argsCount);
+            }
+            else if (!as<IntegerLiteralExpr>(opAttr->args[0]))
+            {
+                sink->diagnose(attr, Diagnostics::attributeExpectedIntArg, attr->keywordName, 0);
+            }
+            else if (argsCount > 1 && !as<StringLiteralExpr>(opAttr->args[1]))
+            {
+                sink->diagnose(attr, Diagnostics::attributeExpectedStringArg, attr->keywordName, 1);
+            }
+        }
+        else if (as<OutputControlPointsAttribute>(attr))
         {
             // Let it go thru iff single integral attribute
             if (!hasIntArgs(attr, 1))
diff --git a/source/slang/slang-diagnostic-defs.h b/source/slang/slang-diagnostic-defs.h
index 8f9327c53..329bf615b 100644
--- a/source/slang/slang-diagnostic-defs.h
+++ b/source/slang/slang-diagnostic-defs.h
@@ -321,6 +321,9 @@ DIAGNOSTIC(31005, Error, expectedSingleStringArg, "attribute '$0' expects a sing
 
 DIAGNOSTIC(31006, Error, attributeFunctionNotFound, "Could not find function '$0' for attribute'$1'")
 
+DIAGNOSTIC(31007, Error, attributeExpectedIntArg, "attribute '$0' expects argument $1 to be int")
+DIAGNOSTIC(31008, Error, attributeExpectedStringArg, "attribute '$0' expects argument $1 to be string")
+
 DIAGNOSTIC(31100, Error, unknownStageName, "unknown stage name '$0'")
 DIAGNOSTIC(31101, Error, unknownImageFormatName, "unknown image format '$0'")
 DIAGNOSTIC(31101, Error, unknownDiagnosticName, "unknown diagnostic '$0'")
diff --git a/source/slang/slang-emit-glsl.cpp b/source/slang/slang-emit-glsl.cpp
index b1b92e7e7..13eb6e2f6 100644
--- a/source/slang/slang-emit-glsl.cpp
+++ b/source/slang/slang-emit-glsl.cpp
@@ -1999,6 +1999,22 @@ void GLSLSourceEmitter::emitFuncDecorationImpl(IRDecoration* decoration)
 
         m_writer->emit("spirv_instruction(id = ");
         emitSimpleValue(decoration->getOperand(0));
+
+        if (decoration->getOperandCount() >= 2)
+        {
+            if (auto stringLit = as<IRStringLit>(decoration->getOperand(1)))
+            {
+                m_writer->emit(toSlice(", set = "));
+
+                auto handler = StringEscapeUtil::getHandler(StringEscapeUtil::Style::Cpp);
+
+                StringBuilder buf;
+                StringEscapeUtil::appendQuoted(handler, stringLit->getStringSlice(), buf);
+            
+                m_writer->emitRawTextSpan(buf.begin(), buf.end());
+            }
+        }
+
         m_writer->emit(")\n");
     }
     else
diff --git a/source/slang/slang-lower-to-ir.cpp b/source/slang/slang-lower-to-ir.cpp
index 383067363..ec51c7bfa 100644
--- a/source/slang/slang-lower-to-ir.cpp
+++ b/source/slang/slang-lower-to-ir.cpp
@@ -7894,21 +7894,19 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
         }
     }
 
-    IRIntLit* _getIntLitFromAttribute(IRBuilder* builder, Attribute* attrib)
+    IRIntLit* _getIntLitFromAttribute(IRBuilder* builder, Attribute* attrib, Index index = 0)
     {
-        attrib->args.getCount();
-        SLANG_ASSERT(attrib->args.getCount() ==1);
-        Expr* expr = attrib->args[0];
+        SLANG_ASSERT(attrib->args.getCount() > index);
+        Expr* expr = attrib->args[index];
         auto intLitExpr = as<IntegerLiteralExpr>(expr);
         SLANG_ASSERT(intLitExpr);
         return as<IRIntLit>(builder->getIntValue(builder->getIntType(), intLitExpr->value));
     }
 
-    IRStringLit* _getStringLitFromAttribute(IRBuilder* builder, Attribute* attrib)
+    IRStringLit* _getStringLitFromAttribute(IRBuilder* builder, Attribute* attrib, Index index = 0)
     {
-        attrib->args.getCount();
-        SLANG_ASSERT(attrib->args.getCount() == 1);
-        Expr* expr = attrib->args[0];
+        SLANG_ASSERT(attrib->args.getCount() > index);
+        Expr* expr = attrib->args[index];
 
         auto stringLitExpr = as<StringLiteralExpr>(expr);
         SLANG_ASSERT(stringLitExpr);
@@ -8332,8 +8330,28 @@ struct DeclLoweringVisitor : DeclVisitor<DeclLoweringVisitor, LoweredValInfo>
 
         if (auto attr = decl->findModifier<SPIRVInstructionOpAttribute>())
         {
-            IRIntLit* intLit = _getIntLitFromAttribute(getBuilder(), attr);
-            getBuilder()->addDecoration(irFunc, kIROp_SPIRVOpDecoration, intLit);
+            auto builder = getBuilder();
+            IRIntLit* intLit = _getIntLitFromAttribute(builder, attr, 0);
+
+            IRStringLit* setStringLit = nullptr;
+            if (attr->args.getCount() > 1)
+            {
+                IRStringLit* checkSetStringLit = _getStringLitFromAttribute(builder, attr, 1);
+                if (checkSetStringLit && checkSetStringLit->getStringSlice().getLength() > 0)
+                {
+                    setStringLit = checkSetStringLit;
+                }
+            }
+
+            // If it has a `set` defined, set it on the decoration
+            if (setStringLit)
+            {
+                builder->addDecoration(irFunc, kIROp_SPIRVOpDecoration, intLit, setStringLit);
+            }
+            else
+            {
+                builder->addDecoration(irFunc, kIROp_SPIRVOpDecoration, intLit);
+            }
         }
 
         if (decl->findModifier<UnsafeForceInlineEarlyAttribute>())