Add support for WGSL subgroup operations (#6213)

* initial work

* more work

* more work on glsl intrinsics

* add subgroup broadcast for glsl

* wip add wgsl extension tracking

* enable tests, enable extensions and added some todos

* format and warning fixes

* fix wgsl extension tracker

---------

Co-authored-by: Yong He <yonghe@outlook.com>

1 files changed, 36 insertions, 0 deletions

diff --git a/source/slang/slang-emit-wgsl.cpp b/source/slang/slang-emit-wgsl.cpp
index ce60cc2a0..aea766f9f 100644
--- a/source/slang/slang-emit-wgsl.cpp
+++ b/source/slang/slang-emit-wgsl.cpp
@@ -49,6 +49,14 @@ fn _slang_getNan() -> f32
 }
 )";
 
+WGSLSourceEmitter::WGSLSourceEmitter(const Desc& desc)
+    : CLikeSourceEmitter(desc)
+{
+    m_extensionTracker =
+        dynamicCast<ShaderExtensionTracker>(desc.codeGenContext->getExtensionTracker());
+    SLANG_ASSERT(m_extensionTracker);
+}
+
 void WGSLSourceEmitter::emitSwitchCaseSelectorsImpl(
     const SwitchRegion::Case* const currentCase,
     const bool isDefault)
@@ -1556,6 +1564,10 @@ void WGSLSourceEmitter::emitFrontMatterImpl(TargetRequest* /* targetReq */)
         m_writer->emit("enable f16;\n");
         m_writer->emit("\n");
     }
+
+    StringBuilder builder;
+    m_extensionTracker->appendExtensionRequireLinesForWGSL(builder);
+    m_writer->emit(builder.getUnownedSlice());
 }
 
 void WGSLSourceEmitter::emitIntrinsicCallExprImpl(
@@ -1626,4 +1638,28 @@ void WGSLSourceEmitter::emitInterpolationModifiersImpl(
     // https://www.w3.org/TR/WGSL/#interpolation
 }
 
+void WGSLSourceEmitter::_requireExtension(const UnownedStringSlice& name)
+{
+    m_extensionTracker->requireExtension(name);
+}
+
+void WGSLSourceEmitter::handleRequiredCapabilitiesImpl(IRInst* inst)
+{
+    for (auto decoration : inst->getDecorations())
+    {
+        if (const auto extensionDecoration = as<IRRequireWGSLExtensionDecoration>(decoration))
+        {
+            _requireExtension(extensionDecoration->getExtensionName());
+
+            // TODO: Make this cleaner and only enable this extension if f16 is actually used on the
+            // subgroup intrinsic. Check float type in meta file.
+            if (m_f16ExtensionEnabled && extensionDecoration->getExtensionName() == "subgroups")
+            {
+                String extName = "subgroups_f16";
+                _requireExtension(extName.getUnownedSlice());
+            }
+        }
+    }
+}
+
 } // namespace Slang