Add WaveGetLane* support for Metal and WGSL

source/slang/glsl.meta.slang

@@ -6334,7 +6334,7 @@ void requireGLSLExtForSubgroupBasicBuiltin() {
     }
 }
 
-[require(cpp_cuda_glsl_hlsl_spirv_wgsl, subgroup_basic)]
+[require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, subgroup_basic)]
 void setupExtForSubgroupBasicBuiltIn() {
     __target_switch
     {
@@ -6409,7 +6409,8 @@ public property uint gl_SubgroupID
 
 public property uint gl_SubgroupSize
 {
-    [require(cpp_cuda_glsl_hlsl_spirv_wgsl, subgroup_basic)]
+    [ForceInline]
+    [require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, subgroup_basic)]
     get {
         setupExtForSubgroupBasicBuiltIn();
         return WaveGetLaneCount();
@@ -6418,7 +6419,8 @@ public property uint gl_SubgroupSize
 
 public property uint gl_SubgroupInvocationID
 {
-    [require(cpp_cuda_glsl_hlsl_spirv_wgsl, subgroup_basic)]
+    [ForceInline]
+    [require(cpp_cuda_glsl_hlsl_metal_spirv_wgsl, subgroup_basic)]
     get {
         setupExtForSubgroupBasicBuiltIn();
         return WaveGetLaneIndex();

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.cpp → source/slang/slang-ir-translate-in-out-global-var.cpp

@@ -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-in-out-global-var.h

@@ -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

tests/glsl-intrinsic/shader-subgroup/shader-subgroup-builtin-variables.slang

@@ -1,5 +1,5 @@
-//TEST:SIMPLE(filecheck=CHECK_GLSL):  -allow-glsl -stage compute -entry computeMain -target glsl
-//TEST:SIMPLE(filecheck=CHECK_SPV):  -allow-glsl -stage compute -entry computeMain -target spirv -emit-spirv-directly
+//TEST:SIMPLE(filecheck=CHECK_GLSL):  -allow-glsl -stage compute -entry computeMain -target glsl -DTARGET_VK
+//TEST:SIMPLE(filecheck=CHECK_SPV):  -allow-glsl -stage compute -entry computeMain -target spirv -emit-spirv-directly -DTARGET_VK
 
 // missing implementation of most builtin values due to non trivial translation
 //DISABLE_TEST:SIMPLE(filecheck=CHECK_HLSL): -allow-glsl -stage compute -entry computeMain -target hlsl -DTARGET_HLSL
@@ -8,8 +8,11 @@
 //missing implementation of system (varying?) values
 //DISABLE_TEST:SIMPLE(filecheck=CHECK_CPP):  -allow-glsl -stage compute -entry computeMain -target cpp -DTARGET_CPP
 
-//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl
-//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -emit-spirv-directly
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -xslang -DTARGET_VK
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-vk -compute -entry computeMain -allow-glsl -emit-spirv-directly -xslang -DTARGET_VK
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-wgpu -compute -entry computeMain -allow-glsl
+//TEST(compute, vulkan):COMPARE_COMPUTE(filecheck-buffer=BUF):-metal -compute -entry computeMain -allow-glsl
+
 #version 430
 
 //TEST_INPUT:ubuffer(data=[0], stride=4):out,name=outputBuffer
@@ -24,15 +27,17 @@ void computeMain()
 {
     if (gl_GlobalInvocationID.x == 3) {
         outputBuffer.data[0] = true
-            && gl_NumSubgroups == 1
-            && gl_SubgroupID  == 0 //1 subgroup, 0 based indexing
             && gl_SubgroupSize == 32
             && gl_SubgroupInvocationID == 3
+#if defined(TARGET_VK)
+            && gl_SubgroupID  == 0 //1 subgroup, 0 based indexing
+            && gl_NumSubgroups == 1
             && gl_SubgroupEqMask == uvec4(0b1000,0,0,0)
             && gl_SubgroupGeMask == uvec4(0xFFFFFFF8,0,0,0)
             && gl_SubgroupGtMask == uvec4(0xFFFFFFF0,0,0,0)
             && gl_SubgroupLeMask == uvec4(0b1111,0,0,0)
             && gl_SubgroupLtMask == uvec4(0b111,0,0,0)
+#endif
             ;
     }
     // CHECK_GLSL: void main(

tests/hlsl-intrinsic/quad-control/quad-control-comp-functionality.slang

@@ -8,11 +8,7 @@ RWStructuredBuffer<uint> outputBuffer;
 [numthreads(16, 1, 1)]
 void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
 {
-#if !defined(METAL)
     uint index = WaveGetLaneIndex();
-#else
-    uint index = dispatchThreadID.x;
-#endif
 
     if (index < 4)
     {

tests/hlsl-intrinsic/wave-get-lane-index.slang

@@ -4,6 +4,8 @@
 //TEST:COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0 -shaderobj
 //TEST(vulkan):COMPARE_COMPUTE_EX:-vk -compute -shaderobj
 //TEST:COMPARE_COMPUTE_EX:-cuda -compute -shaderobj
+//TEST:COMPARE_COMPUTE_EX:-wgpu -compute -shaderobj
+//TEST:COMPARE_COMPUTE_EX:-metal -compute -shaderobj
 
 //TEST_INPUT:ubuffer(data=[0 0 0 0], stride=4):out,name outputBuffer
 RWStructuredBuffer<int> outputBuffer;
@@ -17,4 +19,4 @@ void computeMain(int3 dispatchThreadID : SV_DispatchThreadID)
     // For now we'll just check it's not 0.
     uint laneCount = WaveGetLaneCount();
     outputBuffer[idx] = int(((laneCount > 0) ? 0x100 : 0) + laneId);
-}
+}