support WaveGetLane* for WGSL and Metal

Author: fairywreath

source/slang/hlsl.meta.slang

@@ -6,6 +6,15 @@ typedef uint UINT;
 __intrinsic_op($(kIROp_RequireGLSLExtension))
 void __requireGLSLExtension(String extensionName);
 
+__intrinsic_op($(kIROp_RequireWGSLExtension))
+void __requireWGSLExtension(String extensionName);
+
+/// Built-in values or system value semantics represented as in/out global variables.
+/// This allows the built-ins to be arbitrarily used from a global scope without being
+/// explicitly passed as entry point parameters.
+in uint __builtinWaveLaneIndex : SV_WaveLaneIndex;
+in uint __builtinWaveLaneCount : SV_WaveLaneCount;
+
 //@public:
 /// Represents an interface for buffer data layout.
 /// This interface is used as a base for defining specific data layouts for buffers.
@@ -15037,7 +15046,8 @@ uint WaveActiveCountBits(bool value)
 __glsl_extension(GL_KHR_shader_subgroup_basic)
 __spirv_version(1.3)
 [NonUniformReturn]
-[require(cuda_glsl_hlsl_spirv, subgroup_basic)]
+[ForceInline]
+[require(cuda_glsl_hlsl_metal_spirv_wgsl, subgroup_basic)]
 uint WaveGetLaneCount()
 {
     __target_switch
@@ -15051,14 +15061,20 @@ uint WaveGetLaneCount()
             OpCapability GroupNonUniform;
             result:$$uint = OpLoad builtin(SubgroupSize:uint)
         };
+    case metal:
+        return __builtinWaveLaneCount;
+    case wgsl:
+        __requireWGSLExtension("subgroups");
+        return __builtinWaveLaneCount;
     }
 }
 
 /// @category wave
 __glsl_extension(GL_KHR_shader_subgroup_basic)
 __spirv_version(1.3)
 [NonUniformReturn]
-[require(cuda_glsl_hlsl_spirv, subgroup_basic)]
+[ForceInline]
+[require(cuda_glsl_hlsl_metal_spirv_wgsl, subgroup_basic)]
 uint WaveGetLaneIndex()
 {
     __target_switch
@@ -15072,6 +15088,11 @@ uint WaveGetLaneIndex()
             OpCapability GroupNonUniform;
             result:$$uint = OpLoad builtin(SubgroupLocalInvocationId:uint)
         };
+    case metal:
+        return __builtinWaveLaneIndex;
+    case wgsl:
+        __requireWGSLExtension("subgroups");
+        return __builtinWaveLaneIndex;
     }
 }
 

source/slang/slang-emit-c-like.cpp

@@ -3055,6 +3055,11 @@ void CLikeSourceEmitter::defaultEmitInstExpr(IRInst* inst, const EmitOpInfo& inO
             emitOperand(as<IRGlobalValueRef>(inst)->getOperand(0), getInfo(EmitOp::General));
             break;
         }
+    case kIROp_RequireWGSLExtension:
+        {
+            emitRequireExtension(inst);
+            break;
+        }
     default:
         diagnoseUnhandledInst(inst);
         break;

source/slang/slang-emit-c-like.h

@@ -678,6 +678,8 @@ class CLikeSourceEmitter : public SourceEmitterBase
     void _emitCallArgList(IRCall* call, int startingOperandIndex = 1);
     virtual void emitCallArg(IRInst* arg);
 
+    virtual void emitRequireExtension(IRInst* inst) { SLANG_UNUSED(inst); }
+
     String _generateUniqueName(const UnownedStringSlice& slice);
 
     // Sort witnessTable entries according to the order defined in the witnessed interface type.

source/slang/slang-emit-wgsl.cpp

@@ -1696,4 +1696,9 @@ void WGSLSourceEmitter::handleRequiredCapabilitiesImpl(IRInst* inst)
     }
 }
 
+void WGSLSourceEmitter::emitRequireExtension(IRInst* inst)
+{
+    _requireExtension(as<IRRequireWGSLExtension>(inst)->getExtensionName());
+}
+
 } // namespace Slang

source/slang/slang-emit-wgsl.h

@@ -57,6 +57,8 @@ class WGSLSourceEmitter : public CLikeSourceEmitter
         EmitOpInfo const& inOuterPrec) SLANG_OVERRIDE;
     virtual void emitGlobalParamDefaultVal(IRGlobalParam* varDecl) SLANG_OVERRIDE;
 
+    virtual void emitRequireExtension(IRInst* inst) SLANG_OVERRIDE;
+
     virtual void handleRequiredCapabilitiesImpl(IRInst* inst) SLANG_OVERRIDE;
 
     void emit(const AddressSpace addressSpace);

source/slang/slang-emit.cpp

@@ -100,7 +100,7 @@
 #include "slang-ir-strip-default-construct.h"
 #include "slang-ir-strip-legalization-insts.h"
 #include "slang-ir-synthesize-active-mask.h"
-#include "slang-ir-translate-glsl-global-var.h"
+#include "slang-ir-translate-in-out-global-var.h"
 #include "slang-ir-uniformity.h"
 #include "slang-ir-user-type-hint.h"
 #include "slang-ir-validate.h"
@@ -318,7 +318,7 @@ struct RequiredLoweringPassSet
     bool bindingQuery;
     bool meshOutput;
     bool higherOrderFunc;
-    bool glslGlobalVar;
+    bool inOutGlobalVar;
     bool glslSSBO;
     bool byteAddressBuffer;
     bool dynamicResource;
@@ -422,7 +422,7 @@ void calcRequiredLoweringPassSet(
     case kIROp_GlobalInputDecoration:
     case kIROp_GlobalOutputDecoration:
     case kIROp_GetWorkGroupSize:
-        result.glslGlobalVar = true;
+        result.inOutGlobalVar = true;
         break;
     case kIROp_BindExistentialSlotsDecoration:
         result.bindExistential = true;
@@ -641,8 +641,8 @@ Result linkAndOptimizeIR(
     if (!isKhronosTarget(targetRequest) && requiredLoweringPassSet.glslSSBO)
         lowerGLSLShaderStorageBufferObjectsToStructuredBuffers(irModule, sink);
 
-    if (requiredLoweringPassSet.glslGlobalVar)
-        translateGLSLGlobalVar(codeGenContext, irModule);
+    if (requiredLoweringPassSet.inOutGlobalVar)
+        translateInOutGlobalVar(codeGenContext, irModule);
 
     if (requiredLoweringPassSet.resolveVaryingInputRef)
         resolveVaryingInputRef(irModule);

source/slang/slang-ir-call-graph.h

@@ -1,3 +1,6 @@
+// slang-ir-call-graph.h
+#pragma once
+
 #include "slang-ir-clone.h"
 #include "slang-ir-insts.h"
 

source/slang/slang-ir-inst-defs.h

@@ -407,7 +407,7 @@ INST(WitnessTableEntry, witness_table_entry, 2, 0)
 INST(InterfaceRequirementEntry, interface_req_entry, 2, GLOBAL)
 
 // An inst to represent the workgroup size of the calling entry point.
-// We will materialize this inst during `translateGLSLGlobalVar`.
+// We will materialize this inst during `translateGlobalInOutVar`.
 INST(GetWorkGroupSize, GetWorkGroupSize, 0, HOISTABLE)
 
 // An inst that returns the current stage of the calling entry point.
@@ -667,6 +667,7 @@ INST(discard, discard, 0, 0)
 
 INST(RequirePrelude, RequirePrelude, 1, 0)
 INST(RequireGLSLExtension, RequireGLSLExtension, 1, 0)
+INST(RequireWGSLExtension, RequireWGSLExtension, 1, 0)
 INST(RequireComputeDerivative, RequireComputeDerivative, 0, 0)
 INST(StaticAssert, StaticAssert, 2, 0)
 INST(Printf, Printf, 1, 0)

source/slang/slang-ir-insts.h

@@ -3479,6 +3479,15 @@ struct IRRequireGLSLExtension : IRInst
     }
 };
 
+struct IRRequireWGSLExtension : IRInst
+{
+    IR_LEAF_ISA(RequireWGSLExtension)
+    UnownedStringSlice getExtensionName()
+    {
+        return as<IRStringLit>(getOperand(0))->getStringSlice();
+    }
+};
+
 struct IRRequireComputeDerivative : IRInst
 {
     IR_LEAF_ISA(RequireComputeDerivative)

source/slang/slang-ir-legalize-varying-params.cpp

@@ -3228,6 +3228,20 @@ class LegalizeMetalEntryPointContext : public LegalizeShaderEntryPointContext
                 result.permittedTypes.add(builder.getBasicType(BaseType::UInt));
                 break;
             }
+        case SystemValueSemanticName::WaveLaneCount:
+            {
+                result.systemValueName = toSlice("threads_per_simdgroup");
+                result.permittedTypes.add(builder.getUIntType());
+                result.permittedTypes.add(builder.getUInt16Type());
+                break;
+            }
+        case SystemValueSemanticName::WaveLaneIndex:
+            {
+                result.systemValueName = toSlice("thread_index_in_simdgroup");
+                result.permittedTypes.add(builder.getUIntType());
+                result.permittedTypes.add(builder.getUInt16Type());
+                break;
+            }
         default:
             m_sink->diagnose(
                 parentVar,
@@ -3845,6 +3859,20 @@ class LegalizeWGSLEntryPointContext : public LegalizeShaderEntryPointContext
                 break;
             }
 
+        case SystemValueSemanticName::WaveLaneCount:
+            {
+                result.systemValueName = toSlice("subgroup_size");
+                result.permittedTypes.add(builder.getUIntType());
+                break;
+            }
+
+        case SystemValueSemanticName::WaveLaneIndex:
+            {
+                result.systemValueName = toSlice("subgroup_invocation_id");
+                result.permittedTypes.add(builder.getUIntType());
+                break;
+            }
+
         case SystemValueSemanticName::ViewID:
         case SystemValueSemanticName::ViewportArrayIndex:
         case SystemValueSemanticName::StartVertexLocation:
@@ -3853,7 +3881,6 @@ class LegalizeWGSLEntryPointContext : public LegalizeShaderEntryPointContext
                 result.isUnsupported = true;
                 break;
             }
-
         default:
             {
                 m_sink->diagnose(

source/slang/slang-ir-legalize-varying-params.h

@@ -68,6 +68,8 @@ void depointerizeInputParams(IRFunc* entryPoint);
     M(Target, SV_Target)                                 \
     M(StartVertexLocation, SV_StartVertexLocation)       \
     M(StartInstanceLocation, SV_StartInstanceLocation)   \
+    M(WaveLaneCount, SV_WaveLaneCount)                   \
+    M(WaveLaneIndex, SV_WaveLaneIndex)                   \
     /* end */
 
 /// A known system-value semantic name that can be applied to a parameter

source/slang/slang-ir-translate-glsl-global-var.h

@@ -1,4 +1,4 @@
-#include "slang-ir-translate-glsl-global-var.h"
+#include "slang-ir-translate-in-out-global-var.h"
 
 #include "slang-ir-call-graph.h"
 #include "slang-ir-insts.h"
@@ -373,7 +373,7 @@ struct GlobalVarTranslationContext
     }
 };
 
-void translateGLSLGlobalVar(CodeGenContext* context, IRModule* module)
+void translateInOutGlobalVar(CodeGenContext* context, IRModule* module)
 {
     GlobalVarTranslationContext ctx;
     ctx.context = context;

source/slang/slang-ir-translate-glsl-global-var.cpp source/slang/slang-ir-translate-in-out-global-var.cpp

@@ -0,0 +1,14 @@
+// slang-ir-translate-in-out-global-var.h
+// #pragma once
+
+namespace Slang
+{
+
+struct IRModule;
+struct CodeGenContext;
+
+/// Translate GLSL-flavored global in/out variables into
+/// entry point parameters with system value semantics.
+void translateInOutGlobalVar(CodeGenContext* context, IRModule* module);
+
+} // namespace Slang