Add `set` to spirv_instruction (#2597)

7 files changed, 123 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>())
diff --git a/tests/spirv/spirv-debug-break.slang b/tests/spirv/spirv-debug-break.slang
new file mode 100644
index 000000000..18e8c8597
--- /dev/null
+++ b/tests/spirv/spirv-debug-break.slang
@@ -0,0 +1,22 @@
+// spirv-instruction.slang
+//TEST(compute, vulkan):SIMPLE:-target glsl -entry computeMain -stage compute
+
+[[vk::spirv_instruction(1, "NonSemantic.DebugBreak")]]
+void _spvDebugBreak(int v);
+
+[ForceInline] 
+void _debugBreak() { _spvDebugBreak(1); }
+
+
+//TEST_INPUT:set resultBuffer = out ubuffer(data=[0 0 0 0], stride=4)
+RWStructuredBuffer<uint> resultBuffer;
+
+[numthreads(4,1,1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    uint threadId = dispatchThreadID.x;
+    
+    _debugBreak();
+    
+    resultBuffer[threadId] = threadId + threadId;
+}
diff --git a/tests/spirv/spirv-debug-break.slang.expected b/tests/spirv/spirv-debug-break.slang.expected
new file mode 100644
index 000000000..d8ac8e010
--- /dev/null
+++ b/tests/spirv/spirv-debug-break.slang.expected
@@ -0,0 +1,35 @@
+result code = 0
+standard error = {
+}
+standard output = {
+#version 450
+#extension GL_EXT_spirv_intrinsics : require
+layout(row_major) uniform;
+layout(row_major) buffer;
+
+#line 5 0
+spirv_instruction(id = 1, set = "NonSemantic.DebugBreak")
+void _spvDebugBreak_0(int _0);
+
+
+#line 12
+layout(std430, binding = 0) buffer _S1 {
+    uint _data[];
+} resultBuffer_0;
+layout(local_size_x = 4, local_size_y = 1, local_size_z = 1) in;
+void main()
+{
+
+#line 17
+    uint threadId_0 = gl_GlobalInvocationID.x;
+
+    _spvDebugBreak_0(1);
+
+    uint _S2 = threadId_0 + threadId_0;
+
+#line 21
+    ((resultBuffer_0)._data[(threadId_0)]) = _S2;
+    return;
+}
+
+}