First support for 'WaveMask' intrinsics (#1321)

* WIP tests to confirm divergence on CUDA.

* Added wave.slang test that uses masks.
Made all CUDA intrinsic impls take a mask explicitly.
Added initial WaveMaskXXX intrinsics.

* Added WaveMaskSharedSync.

* Improvements aroung WaveMaskSharedSync/WaveMaskSync

* Remove tabs.

3 files changed, 325 insertions, 138 deletions

diff --git a/prelude/slang-cuda-prelude.h b/prelude/slang-cuda-prelude.h
index c23189320..4a91848e4 100644
--- a/prelude/slang-cuda-prelude.h
+++ b/prelude/slang-cuda-prelude.h
@@ -484,6 +484,8 @@ __forceinline__ __device__ uint32_t _getLaneId()
 }
 #endif
 
+typedef int WarpMask;
+
 // It appears that the __activemask() cannot always be used because 
 // threads need to be converged. 
 // 
@@ -500,44 +502,39 @@ __forceinline__ __device__ uint32_t _getLaneId()
 // to ensure they are properly converged before the intrinsic is executed by the hardware. All active threads named 
 // in mask must execute the same intrinsic with the same mask, or the result is undefined.```
 //
-// To get the right results we need to use the __activemask() within _ballot_sync it seems.
-// 
-// Also note that __all_sync and __any_sync are listed with __ballot_sync. That if they have a similar synchronizing behavior
-// we can use __activemask() there (instead of _getConvergedMask), because they will converge too. 
-__forceinline__ __device__ int _getConvergedMask()
-{
-    //return __activemask(); 
-    //return __ballot_sync(SLANG_CUDA_WARP_MASK, true);
-    return __ballot_sync(__activemask(), true);
-}
+// Currently there isn't a mechanism to correctly get the mask without it being passed through.
+// Doing so will most likely require some changes to slang code generation to track masks, for now then we use
+// _getActiveMask. 
 
 // Return mask of all the lanes less than the current lane
-__forceinline__ __device__ int _getLaneLtMask()
+__forceinline__ __device__ WarpMask _getLaneLtMask()
 {
     return (int(1) << _getLaneId()) - 1;
 }    
 
-// Return a mask suitable for the straight 'Prefix' style ops
-__forceinline__ __device__ int _getPrefixMask()
+// TODO(JS): 
+// THIS IS NOT CORRECT! That determining the appropriate active mask requires appropriate
+// mask tracking.
+__forceinline__ __device__ WarpMask _getActiveMask()
 {
-    return __activemask();
+    return __ballot_sync(__activemask(), true);
 }
 
 // Return a mask suitable for the 'MultiPrefix' style functions
-__forceinline__ __device__ int _getMultiPrefixMask(int mask)
+__forceinline__ __device__ WarpMask _getMultiPrefixMask(int mask)
 {
     return mask;
 }
 
 // Note! Note will return true if mask is 0, but thats okay, because there must be one
 // lane active to execute anything
-__inline__ __device__ bool _waveIsSingleLane(int mask)
+__inline__ __device__ bool _waveIsSingleLane(WarpMask mask)
 {
     return (mask & (mask - 1)) == 0;
 }
 
 // Returns the power of 2 size of run of set bits. Returns 0 if not a suitable run.
-__inline__ __device__ int _waveCalcPow2Offset(int mask)
+__inline__ __device__ int _waveCalcPow2Offset(WarpMask mask)
 {
     // This should be the most common case, so fast path it
     if (mask == SLANG_CUDA_WARP_MASK)
@@ -560,7 +557,7 @@ __inline__ __device__ int _waveCalcPow2Offset(int mask)
 
 __inline__ __device__ bool _waveIsFirstLane()
 {
-    const int mask = __activemask();
+    const WarpMask mask = __activemask();
     // We special case bit 0, as that most warps are expected to be fully active. 
     
     // mask & -mask, isolates the lowest set bit.
@@ -665,12 +662,8 @@ struct ElementTypeTrait<Matrix<T, ROWS, COLS> >
 
 // Scalar 
 template <typename INTF, typename T>
-__device__ T _waveReduceScalar(T val)
+__device__ T _waveReduceScalar(WarpMask mask, T val)
 {
-    // The shuffles appear to converge on set bits, so it appears ok to use __activemask()
-    //const int mask = _getConvergedMask();
-    const int mask = __activemask();
-    
     const int offsetSize = _waveCalcPow2Offset(mask);
     if (offsetSize > 0)
     {
@@ -701,12 +694,8 @@ __device__ T _waveReduceScalar(T val)
 
 // Multiple values
 template <typename INTF, typename T, size_t COUNT>
-__device__ void _waveReduceMultiple(T* val)
+__device__ void _waveReduceMultiple(WarpMask mask, T* val)
 {
-    // The shuffles appear to converge on set bits, so it appears ok to use __activemask()
-    //const int mask = _getConvergedMask();
-    const int mask = __activemask();
-    
     const int offsetSize = _waveCalcPow2Offset(mask);
     if (offsetSize > 0)
     {
@@ -747,75 +736,71 @@ __device__ void _waveReduceMultiple(T* val)
 }
 
 template <typename INTF, typename T>
-__device__ void _waveReduceMultiple(T* val)
+__device__ void _waveReduceMultiple(WarpMask mask, T* val)
 {
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _waveReduceMultiple<INTF, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)val);
+    _waveReduceMultiple<INTF, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)val);
 }
 
 template <typename T>
-__inline__ __device__  T _waveOr(T val) { return _waveReduceScalar<WaveOpOr<T>, T>(val); }
+__inline__ __device__  T _waveOr(WarpMask mask, T val) { return _waveReduceScalar<WaveOpOr<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveAnd(T val) { return _waveReduceScalar<WaveOpAnd<T>, T>(val); }
+__inline__ __device__ T _waveAnd(WarpMask mask, T val) { return _waveReduceScalar<WaveOpAnd<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveXor(T val) { return _waveReduceScalar<WaveOpXor<T>, T>(val); }
+__inline__ __device__ T _waveXor(WarpMask mask, T val) { return _waveReduceScalar<WaveOpXor<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveProduct(T val) { return _waveReduceScalar<WaveOpMul<T>, T>(val); }
+__inline__ __device__ T _waveProduct(WarpMask mask, T val) { return _waveReduceScalar<WaveOpMul<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveSum(T val) { return _waveReduceScalar<WaveOpAdd<T>, T>(val); }
+__inline__ __device__ T _waveSum(WarpMask mask, T val) { return _waveReduceScalar<WaveOpAdd<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveMin(T val) { return _waveReduceScalar<WaveOpMin<T>, T>(val); }
+__inline__ __device__ T _waveMin(WarpMask mask, T val) { return _waveReduceScalar<WaveOpMin<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _waveMax(T val) { return _waveReduceScalar<WaveOpMax<T>, T>(val); }
+__inline__ __device__ T _waveMax(WarpMask mask, T val) { return _waveReduceScalar<WaveOpMax<T>, T>(mask, val); }
 
 
 // Multiple
 
 template <typename T>
-__inline__ __device__  T _waveOrMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpOr<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveOrMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpOr<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveAndMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpAnd<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveAndMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpAnd<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveXorMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpXor<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveXorMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpXor<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveProductMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMul<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveProductMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMul<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveSumMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpAdd<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveSumMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpAdd<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveMinMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMin<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveMinMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMin<ElemType> >(mask, &val); return val; }
 
 template <typename T>
-__inline__ __device__  T _waveMaxMultiple(T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMax<ElemType> >(&val); return val; }
+__inline__ __device__  T _waveMaxMultiple(WarpMask mask, T val) { typedef typename ElementTypeTrait<T>::Type ElemType; _waveReduceMultiple<WaveOpMax<ElemType> >(mask, &val); return val; }
 
 
 template <typename T>
-__inline__ __device__ bool _waveAllEqual(T val) 
+__inline__ __device__ bool _waveAllEqual(WarpMask mask, T val) 
 {
-    // __match_all_sync synchronizes so can use __activemask()
-    const int mask = __activemask();
     int pred;
     __match_all_sync(mask, val, &pred);
     return pred != 0;
 }
 
 template <typename T>
-__inline__ __device__ bool _waveAllEqualMultiple(T inVal) 
+__inline__ __device__ bool _waveAllEqualMultiple(WarpMask mask, T inVal) 
 {
     typedef typename ElementTypeTrait<T>::Type ElemType;
     const size_t count = sizeof(T) / sizeof(ElemType);
-    // __match_all_sync synchronizes so can use __activemask()
-    const int mask = __activemask();
     int pred;
     const ElemType* src = (const ElemType*)&inVal;
     for (size_t i = 0; i < count; ++i)
@@ -830,22 +815,20 @@ __inline__ __device__ bool _waveAllEqualMultiple(T inVal)
 }
 
 template <typename T>
-__inline__ __device__ T _waveReadFirst(T val) 
+__inline__ __device__ T _waveReadFirst(WarpMask mask, T val) 
 {
-    const int mask = __activemask();
     const int lowestLaneId = __ffs(mask) - 1;
     return __shfl_sync(mask, val, lowestLaneId);   
 }
 
 template <typename T>
-__inline__ __device__ T _waveReadFirstMultiple(T inVal) 
+__inline__ __device__ T _waveReadFirstMultiple(WarpMask mask, T inVal) 
 {
     typedef typename ElementTypeTrait<T>::Type ElemType;
     const size_t count = sizeof(T) / sizeof(ElemType);
     T outVal;
     const ElemType* src = (const ElemType*)&inVal;
     ElemType* dst = (ElemType*)&outVal;
-    const int mask = __activemask();
     const int lowestLaneId = __ffs(mask) - 1;
     for (size_t i = 0; i < count; ++i)
     {
@@ -855,14 +838,13 @@ __inline__ __device__ T _waveReadFirstMultiple(T inVal)
 }
 
 template <typename T>
-__inline__ __device__ T _waveShuffleMultiple(T inVal, int lane)
+__inline__ __device__ T _waveShuffleMultiple(WarpMask mask, T inVal, int lane)
 {
     typedef typename ElementTypeTrait<T>::Type ElemType;
     const size_t count = sizeof(T) / sizeof(ElemType);
     T outVal;
     const ElemType* src = (const ElemType*)&inVal;
     ElemType* dst = (ElemType*)&outVal;
-    const int mask = __activemask();
     for (size_t i = 0; i < count; ++i)
     {
         dst[i] = __shfl_sync(mask, src[i], lane);   
@@ -875,7 +857,7 @@ __inline__ __device__ T _waveShuffleMultiple(T inVal, int lane)
 // Invertable means that when we get to the end of the reduce, we can remove val (to make exclusive), using 
 // the inverse of the op.
 template <typename INTF, typename T>
-__device__ T _wavePrefixInvertableScalar(T val, const int mask)
+__device__ T _wavePrefixInvertableScalar(WarpMask mask, T val)
 {
     const int offsetSize = _waveCalcPow2Offset(mask);
     
@@ -925,7 +907,7 @@ __device__ T _wavePrefixInvertableScalar(T val, const int mask)
 // This implementation separately tracks the value to be propogated, and the value
 // that is the final result 
 template <typename INTF, typename T>
-__device__ T _wavePrefixScalar(T val, const int mask)
+__device__ T _wavePrefixScalar(WarpMask mask, T val)
 {
     const int offsetSize = _waveCalcPow2Offset(mask);
     
@@ -1000,7 +982,7 @@ __device__ T _waveOpSetInitial(T* out, const T* val)
 } 
 
 template <typename INTF, typename T, size_t COUNT>
-__device__ T _wavePrefixInvertableMultiple(T* val, const int mask)
+__device__ T _wavePrefixInvertableMultiple(WarpMask mask, T* val)
 {
     const int offsetSize = _waveCalcPow2Offset(mask);
     
@@ -1058,7 +1040,7 @@ __device__ T _wavePrefixInvertableMultiple(T* val, const int mask)
 }
  
 template <typename INTF, typename T, size_t COUNT>
-__device__ T _wavePrefixMultiple(T* val, const int mask)
+__device__ T _wavePrefixMultiple(WarpMask mask, T* val)
 {
     const int offsetSize = _waveCalcPow2Offset(mask);
     
@@ -1114,77 +1096,73 @@ __device__ T _wavePrefixMultiple(T* val, const int mask)
 }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixProduct(T val, const int mask = _getPrefixMask()) { return _wavePrefixScalar<WaveOpMul<T>, T>(val, mask); }
+__inline__ __device__ T _wavePrefixProduct(WarpMask mask, T val) { return _wavePrefixScalar<WaveOpMul<T>, T>(mask, val); }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixSum(T val, const int mask = _getPrefixMask()) { return _wavePrefixInvertableScalar<WaveOpAdd<T>, T>(val, mask); }    
+__inline__ __device__ T _wavePrefixSum(WarpMask mask, T val) { return _wavePrefixInvertableScalar<WaveOpAdd<T>, T>(mask, val); }    
 
 template <typename T>
-__inline__ __device__ T _wavePrefixXor(T val, const int mask = _getPrefixMask()) { return _wavePrefixInvertableScalar<WaveOpXor<T>, T>(val, mask); }    
+__inline__ __device__ T _wavePrefixXor(WarpMask mask, T val) { return _wavePrefixInvertableScalar<WaveOpXor<T>, T>(mask, val); }    
     
 template <typename T>
-__inline__ __device__ T _wavePrefixOr(T val, const int mask = _getPrefixMask()) { return _wavePrefixScalar<WaveOpOr<T>, T>(val, mask); }      
+__inline__ __device__ T _wavePrefixOr(WarpMask mask, T val) { return _wavePrefixScalar<WaveOpOr<T>, T>(mask, val); }      
     
 template <typename T>
-__inline__ __device__ T _wavePrefixAnd(T val, const int mask = _getPrefixMask()) { return _wavePrefixScalar<WaveOpAnd<T>, T>(val, mask); }      
+__inline__ __device__ T _wavePrefixAnd(WarpMask mask, T val) { return _wavePrefixScalar<WaveOpAnd<T>, T>(mask, val); }      
     
     
 template <typename T>
-__inline__ __device__ T _wavePrefixProductMultiple(T val, const int mask = _getPrefixMask())  
+__inline__ __device__ T _wavePrefixProductMultiple(WarpMask mask, T val)  
 { 
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _wavePrefixInvertableMultiple<WaveOpMul<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)&val, mask);    
+    _wavePrefixInvertableMultiple<WaveOpMul<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)&val);    
     return val;
 }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixSumMultiple(T val, const int mask = _getPrefixMask()) 
+__inline__ __device__ T _wavePrefixSumMultiple(WarpMask mask, T val) 
 { 
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _wavePrefixInvertableMultiple<WaveOpAdd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)&val, mask);    
+    _wavePrefixInvertableMultiple<WaveOpAdd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)&val);    
     return val;
 }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixXorMultiple(T val, const int mask = _getPrefixMask())  
+__inline__ __device__ T _wavePrefixXorMultiple(WarpMask mask, T val)  
 { 
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _wavePrefixInvertableMultiple<WaveOpXor<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)&val, mask);    
+    _wavePrefixInvertableMultiple<WaveOpXor<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)&val);    
     return val;
 }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixOrMultiple(T val, const int mask = _getPrefixMask()) 
+__inline__ __device__ T _wavePrefixOrMultiple(WarpMask mask, T val) 
 { 
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _wavePrefixMultiple<WaveOpOr<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)&val, mask);    
+    _wavePrefixMultiple<WaveOpOr<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)&val);    
     return val;
 }
 
 template <typename T>
-__inline__ __device__ T _wavePrefixAndMultiple(T val, const int mask = _getPrefixMask())  
+__inline__ __device__ T _wavePrefixAndMultiple(WarpMask mask, T val)  
 { 
     typedef typename ElementTypeTrait<T>::Type ElemType;    
-    _wavePrefixMultiple<WaveOpAnd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>((ElemType*)&val, mask);    
+    _wavePrefixMultiple<WaveOpAnd<ElemType>, ElemType, sizeof(T) / sizeof(ElemType)>(mask, (ElemType*)&val);    
     return val;
 }
 
 template <typename T>
-__inline__ __device__ uint4 _waveMatchScalar(T val) 
+__inline__ __device__ uint4 _waveMatchScalar(WarpMask mask, T val) 
 {
-    // __match_all_sync synchronizes so can use __activemask()
-    const int mask = __activemask();
     int pred;
     return make_uint4(__match_all_sync(mask, val, &pred), 0, 0, 0);
 }
 
 template <typename T>
-__inline__ __device__ uint4 _waveMatchMultiple(const T& inVal) 
+__inline__ __device__ uint4 _waveMatchMultiple(WarpMask mask, const T& inVal) 
 {
     typedef typename ElementTypeTrait<T>::Type ElemType;
     const size_t count = sizeof(T) / sizeof(ElemType);
-    // __match_all_sync synchronizes so can use __activemask()
-    const int mask = __activemask();
     int pred;
     const ElemType* src = (const ElemType*)&inVal;
     uint matchBits = 0xffffffff;
diff --git a/source/slang/hlsl.meta.slang b/source/slang/hlsl.meta.slang
index 4279e4a4e..f096a125e 100644
--- a/source/slang/hlsl.meta.slang
+++ b/source/slang/hlsl.meta.slang
@@ -2479,6 +2479,151 @@ matrix<T, N, M> trunc(matrix<T, N, M> x)
     MATRIX_MAP_UNARY(T, N, M, trunc, x);
 }
 
+// Slang Specific Mask Wave Intrinsics 
+
+typedef uint WaveMask;
+
+__target_intrinsic(cuda, "__activemask()")
+WaveMask WaveGetActiveMask() { return 0xffffffff; }
+
+__glsl_extension(GL_KHR_shader_subgroup_vote)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupAll($1)") 
+__target_intrinsic(cuda, "(__all_sync($0, $1) != 0)")
+__target_intrinsic(hlsl, "WaveActiveAllTrue($1)")
+bool WaveMaskAllTrue(WaveMask mask, bool condition);
+
+__glsl_extension(GL_KHR_shader_subgroup_vote)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupAny($1)") 
+__target_intrinsic(cuda, "(__any_sync($0, $1) != 0)")
+__target_intrinsic(hlsl, "WaveActiveAnyTrue($1)")
+bool WaveMaskAnyTrue(WaveMask mask, bool condition);
+
+__glsl_extension(GL_KHR_shader_subgroup_ballot)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBallot($1).x")
+__target_intrinsic(cuda, "__ballot_sync($0, $1)")
+__target_intrinsic(hlsl, "WaveActiveBallot($1)")
+WaveMask WaveMaskBallot(WaveMask mask, bool condition);
+
+__glsl_extension(GL_KHR_shader_subgroup_ballot)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "bitCount(subgroupBallot($1))")
+__target_intrinsic(cuda, "__popc(__ballot_sync($0, $1))")
+__target_intrinsic(hlsl, "WaveActiveCountBits($1)")
+WaveMask WaveMaskCountBits(WaveMask mask, bool value);
+
+// Waits until all warp lanes named in mask have executed a WaveMaskSharedSync (with the same mask)
+// before resuming execution. Guarantees memory ordering in shared memory among threads participating
+// in the barrier.
+//
+// The CUDA intrinsic says it orders *all* memory accesses, which appears to match most closely subgroupBarrier.
+//
+// TODO(JS):
+// For HLSL it's not clear what to do. There is no explicit mechanism to 'reconverge' threads. In the docs it describes
+// behavior as
+// "These intrinsics are dependent on active lanes and therefore flow control. In the model of this document, implementations
+// must enforce that the number of active lanes exactly corresponds to the programmer’s view of flow control."
+// 
+// It seems this can only mean the active threads are the "threads the program flow would lead to". This implies a lockstep
+// "straight SIMD" style interpretation. That being the case this op on HLSL is just a memory barrier without any Sync.
+
+__target_intrinsic(cuda, "__syncwarp($0)")
+__glsl_extension(GL_KHR_shader_subgroup_basic)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBarrier()")
+__target_intrinsic(hlsl, "AllMemoryBarrier()")
+void WaveMaskSync(WaveMask mask);
+
+// On GLSL, it appears we can't use subgroupMemoryBarrierShared, because it only implies a memory ordering, it does not
+// imply convergence. For subgroupBarrier we have from the docs..
+// "The function subgroupBarrier() enforces that all active invocations within a subgroup must execute this function before any
+// are allowed to continue their execution"
+
+__target_intrinsic(cuda, "__syncwarp($0)")
+__glsl_extension(GL_KHR_shader_subgroup_basic)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBarrier()")
+__target_intrinsic(hlsl, "GroupMemoryBarrier()")
+void WaveMaskSharedSync(WaveMask mask);
+
+// NOTE! WaveMaskBroadcastLaneAt is *NOT* standard HLSL
+// It is provided as access to subgroupBroadcast which can only take a
+// constexpr laneId.
+// https://github.com/KhronosGroup/GLSL/blob/master/extensions/khr/GL_KHR_shader_subgroup.txt
+// Versions SPIR-V greater than 1.4 loosen this restriction, and allow 'dynamic uniform' index
+// If that's the behavior required then client code should use WaveReadLaneAt which works this way.
+
+__generic<T : __BuiltinType>
+__glsl_extension(GL_KHR_shader_subgroup_ballot)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBroadcast($1, $2)")
+__target_intrinsic(cuda, "__shfl_sync($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+T WaveMaskBroadcastLaneAt(WaveMask mask, T value, constexpr int lane);
+__generic<T : __BuiltinType, let N : int>
+__glsl_extension(GL_KHR_shader_subgroup_ballot)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBroadcast($1, $2)")
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+vector<T,N> WaveMaskBroadcastLaneAt(WaveMask mask, vector<T,N> value, constexpr int lane);
+__generic<T : __BuiltinType, let N : int, let M : int>
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+matrix<T,N,M> WaveMaskBroadcastLaneAt(WaveMask mask, matrix<T,N,M> value, constexpr int lane);
+
+// TODO(JS): If it can be determines that the `laneId` is constExpr, then subgroupBroadcast
+// could be used on GLSL. For now we just use subgroupShuffle
+__generic<T : __BuiltinType>
+__glsl_extension(GL_KHR_shader_subgroup_shuffle)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupShuffle($1, $2)")
+__target_intrinsic(cuda, "__shfl_sync($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+T WaveMaskReadLaneAt(WaveMask mask, T value, int lane);
+__generic<T : __BuiltinType, let N : int>
+__spirv_version(1.3)
+__glsl_extension(GL_KHR_shader_subgroup_shuffle)
+__target_intrinsic(glsl, "subgroupShuffle($1, $2)")
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+vector<T,N> WaveMaskReadLaneAt(WaveMask mask, vector<T,N> value, int lane);
+__generic<T : __BuiltinType, let N : int, let M : int>
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+matrix<T,N,M> WaveMaskReadLaneAt(WaveMask mask, matrix<T,N,M> value, int lane);
+
+// NOTE! WaveMaskShuffle is a NON STANDARD HLSL intrinsic! It will map to WaveReadLaneAt on HLSL
+// which means it will only work on hardware which allows arbitrary laneIds which is not true
+// in general because it breaks the HLSL standard, which requires it's 'dynamically uniform' across the Wave.
+__generic<T : __BuiltinType>
+__glsl_extension(GL_KHR_shader_subgroup_shuffle)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupShuffle($1, $2)")
+__target_intrinsic(cuda, "__shfl_sync($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+T WaveMaskShuffle(WaveMask mask, T value, int lane);
+__generic<T : __BuiltinType, let N : int>
+__glsl_extension(GL_KHR_shader_subgroup_shuffle)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupShuffle($1, $2)")
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+vector<T,N> WaveMaskShuffle(WaveMask mask, vector<T,N> value, int lane);
+__generic<T : __BuiltinType, let N : int, let M : int>
+__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1, $2)")
+__target_intrinsic(hlsl, "WaveReadLaneAt($1, $2)")
+matrix<T,N,M> WaveMaskShuffle(WaveMask mask, matrix<T,N,M> value, int lane);
+
+__glsl_extension(GL_KHR_shader_subgroup_ballot)
+__spirv_version(1.3)
+__target_intrinsic(glsl, "subgroupBallotExclusiveBitCount(subgroupBallot($1))")
+__target_intrinsic(cuda, "__popc(__ballot_sync($0, $1)  & _getLaneLtMask())")
+__target_intrinsic(hlsl, "WavePrefixCountBits($1)")
+uint WaveMaskPrefixCountBits(WaveMask mask, bool value);
+
 // Shader model 6.0 stuff
 
 // Information for GLSL wave/subgroup support
@@ -2504,112 +2649,112 @@ __generic<T : __BuiltinIntegerType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAnd($0)")
-__target_intrinsic(cuda, "_waveAnd($0)")
+__target_intrinsic(cuda, "_waveAnd(_getActiveMask(), $0)")
 T WaveActiveBitAnd(T expr);
 __generic<T : __BuiltinIntegerType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAnd($0)")
-__target_intrinsic(cuda, "_waveAndMultiple($0)")
+__target_intrinsic(cuda, "_waveAndMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveBitAnd(vector<T,N> expr);
 __generic<T : __BuiltinIntegerType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveAndMultiple($0)")
+__target_intrinsic(cuda, "_waveAndMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveBitAnd(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinIntegerType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupOr($0)")
-__target_intrinsic(cuda, "_waveOr($0)")
+__target_intrinsic(cuda, "_waveOr(_getActiveMask(), $0)")
 T WaveActiveBitOr(T expr);
 __generic<T : __BuiltinIntegerType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupOr($0)")
-__target_intrinsic(cuda, "_waveOrMultiple($0)")
+__target_intrinsic(cuda, "_waveOrMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveBitOr(vector<T,N> expr);
 __generic<T : __BuiltinIntegerType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveOrMultiple($0)")
+__target_intrinsic(cuda, "_waveOrMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveBitOr(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinIntegerType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupXor($0)")
-__target_intrinsic(cuda, "_waveXor($0)")
+__target_intrinsic(cuda, "_waveXor(_getActiveMask(), $0)")
 T WaveActiveBitXor(T expr);
 __generic<T : __BuiltinIntegerType, let N : int> 
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupXor($0)")
-__target_intrinsic(cuda, "_waveXorMultiple($0)")
+__target_intrinsic(cuda, "_waveXorMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveBitXor(vector<T,N> expr);
 __generic<T : __BuiltinIntegerType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveXorMultiple($0)")
+__target_intrinsic(cuda, "_waveXorMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveBitXor(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMax($0)")
-__target_intrinsic(cuda, "_waveMax($0)")
+__target_intrinsic(cuda, "_waveMax(_getActiveMask(), $0)")
 T WaveActiveMax(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMax($0)")
-__target_intrinsic(cuda, "_waveMaxMultiple($0)")
+__target_intrinsic(cuda, "_waveMaxMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveMax(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveMaxMultiple($0)")
+__target_intrinsic(cuda, "_waveMaxMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveMax(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMin($0)")
-__target_intrinsic(cuda, "_waveMin($0)")
+__target_intrinsic(cuda, "_waveMin(_getActiveMask(), $0)")
 T WaveActiveMin(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMin($0)")
-__target_intrinsic(cuda, "_waveMinMultiple($0)")
+__target_intrinsic(cuda, "_waveMinMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveMin(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveMinMultiple($0)")
+__target_intrinsic(cuda, "_waveMinMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveMin(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMul($0)")
-__target_intrinsic(cuda, "_waveProduct($0)")
+__target_intrinsic(cuda, "_waveProduct(_getActiveMask(), $0)")
 T WaveActiveProduct(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupMul($0)")
-__target_intrinsic(cuda, "_waveProductMultiple($0)")
+__target_intrinsic(cuda, "_waveProductMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveProduct(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveProductMultiple($0)")
+__target_intrinsic(cuda, "_waveProductMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveProduct(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAdd($0)")
-__target_intrinsic(cuda, "_waveSum($0)")
+__target_intrinsic(cuda, "_waveSum(_getActiveMask(), $0)")
 T WaveActiveSum(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAdd($0)")
-__target_intrinsic(cuda, "_waveSumMultiple($0)")
+__target_intrinsic(cuda, "_waveSumMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveActiveSum(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveSumMultiple($0)")
+__target_intrinsic(cuda, "_waveSumMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveActiveSum(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinType>
@@ -2617,18 +2762,18 @@ __glsl_extension(GL_KHR_shader_subgroup_vote)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAllEqual($0)")
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveAllEqual($0)")
+__target_intrinsic(cuda, "_waveAllEqual(_getActiveMask(), $0)")
 bool WaveActiveAllEqual(T value);
 __generic<T : __BuiltinType, let N : int> 
 __glsl_extension(GL_KHR_shader_subgroup_vote)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupAllEqual($0)")
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveAllEqualMultiple($0)")
+__target_intrinsic(cuda, "_waveAllEqualMultiple(_getActiveMask(), $0)")
 bool WaveActiveAllEqual(vector<T,N> value);
 __generic<T : __BuiltinType, let N : int, let M : int>
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveAllEqualMultiple($0)")
+__target_intrinsic(cuda, "_waveAllEqualMultiple(_getActiveMask(), $0)")
 bool WaveActiveAllEqual(matrix<T,N,M> value);
 
 
@@ -2679,48 +2824,48 @@ __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveMul($0)")
-__target_intrinsic(cuda, "_wavePrefixProduct($0)")
+__target_intrinsic(cuda, "_wavePrefixProduct(_getActiveMask(), $0)")
 T WavePrefixProduct(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveMul($0)")
-__target_intrinsic(cuda, "_wavePrefixProductMultiple($0)")
+__target_intrinsic(cuda, "_wavePrefixProductMultiple(_getActiveMask(), $0)")
 vector<T,N> WavePrefixProduct(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_wavePrefixProductMultiple($0)")
+__target_intrinsic(cuda, "_wavePrefixProductMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WavePrefixProduct(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinArithmeticType>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveAdd($0)")
-__target_intrinsic(cuda, "_wavePrefixSum($0)")
+__target_intrinsic(cuda, "_wavePrefixSum(_getActiveMask(), $0)")
 T WavePrefixSum(T expr);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveAdd($0)")
-__target_intrinsic(cuda, "_wavePrefixSumMultiple($0)")
+__target_intrinsic(cuda, "_wavePrefixSumMultiple(_getActiveMask(), $0)")
 vector<T,N> WavePrefixSum(vector<T,N> expr);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
-__target_intrinsic(cuda, "_wavePrefixSumMultiple($0)")
+__target_intrinsic(cuda, "_wavePrefixSumMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WavePrefixSum(matrix<T,N,M> expr);
 
 __generic<T : __BuiltinType>
 __glsl_extension(GL_KHR_shader_subgroup_ballot)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupBroadcastFirst($0)")
-__target_intrinsic(cuda, "_waveReadFirst($0)")
+__target_intrinsic(cuda, "_waveReadFirst(_getActiveMask(), $0)")
 T WaveReadLaneFirst(T expr);
 __generic<T : __BuiltinType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_ballot)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupBroadcastFirst($0)")
-__target_intrinsic(cuda, "_waveReadFirstMultiple($0)")
+__target_intrinsic(cuda, "_waveReadFirstMultiple(_getActiveMask(), $0)")
 vector<T,N> WaveReadLaneFirst(vector<T,N> expr);
 __generic<T : __BuiltinType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveReadFirstMultiple($0)")
+__target_intrinsic(cuda, "_waveReadFirstMultiple(_getActiveMask(), $0)")
 matrix<T,N,M> WaveReadLaneFirst(matrix<T,N,M> expr);
 
 // NOTE! WaveBroadcastLaneAt is *NOT* standard HLSL
@@ -2740,11 +2885,11 @@ __generic<T : __BuiltinType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_ballot)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupBroadcast($0, $1)")
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 __target_intrinsic(hlsl, "WaveReadLaneAt")
 vector<T,N> WaveBroadcastLaneAt(vector<T,N> value, constexpr int lane);
 __generic<T : __BuiltinType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 __target_intrinsic(hlsl, "WaveReadLaneAt")
 matrix<T,N,M> WaveBroadcastLaneAt(matrix<T,N,M> value, constexpr int lane);
 
@@ -2760,10 +2905,10 @@ __generic<T : __BuiltinType, let N : int>
 __spirv_version(1.3)
 __glsl_extension(GL_KHR_shader_subgroup_shuffle)
 __target_intrinsic(glsl, "subgroupShuffle($0, $1)")
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 vector<T,N> WaveReadLaneAt(vector<T,N> value, int lane);
 __generic<T : __BuiltinType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 matrix<T,N,M> WaveReadLaneAt(matrix<T,N,M> value, int lane);
 
 // NOTE! WaveShuffle is a NON STANDARD HLSL intrinsic! It will map to WaveReadLaneAt on HLSL
@@ -2780,11 +2925,11 @@ __generic<T : __BuiltinType, let N : int>
 __glsl_extension(GL_KHR_shader_subgroup_shuffle)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupShuffle($0, $1)")
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 __target_intrinsic(hlsl, "WaveReadLaneAt")
 vector<T,N> WaveShuffle(vector<T,N> value, int lane);
 __generic<T : __BuiltinType, let N : int, let M : int>
-__target_intrinsic(cuda, "_waveShuffleMultiple($0, $1)")
+__target_intrinsic(cuda, "_waveShuffleMultiple(_getActiveMask(), $0, $1)")
 __target_intrinsic(hlsl, "WaveReadLaneAt")
 matrix<T,N,M> WaveShuffle(matrix<T,N,M> value, int lane);
 
@@ -2800,17 +2945,17 @@ uint WavePrefixCountBits(bool value);
 __generic<T : __BuiltinType>
 __target_intrinsic(hlsl)
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveMatchScalar($0)")
+__target_intrinsic(cuda, "_waveMatchScalar(_getActiveMask(), $0)")
 uint4 WaveMatch(T value);
 __generic<T : __BuiltinType, let N : int>
 __target_intrinsic(hlsl)
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveMatchMultiple($0)")
+__target_intrinsic(cuda, "_waveMatchMultiple(_getActiveMask(), $0)")
 uint4 WaveMatch(vector<T,N> value);
 __generic<T : __BuiltinType, let N : int, let M : int>
 __target_intrinsic(hlsl)
 __cuda_sm_version(7.0)
-__target_intrinsic(cuda, "_waveMatchMultiple($0)")
+__target_intrinsic(cuda, "_waveMatchMultiple(_getActiveMask(), $0)")
 uint4 WaveMatch(matrix<T,N,M> value);
 
 __target_intrinsic(hlsl)
@@ -2822,18 +2967,18 @@ __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 //__target_intrinsic(glsl, "subgroupExclusiveAnd($0)")
-__target_intrinsic(cuda, "_wavePrefixAnd($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixAnd(_getMultiPrefixMask(($1).x), $0)")
 T WaveMultiPrefixBitAnd(T expr, uint4 mask);
 __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveAnd($0)")
-__target_intrinsic(cuda, "_wavePrefixAndMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixAndMultiple(_getMultiPrefixMask(($1).x), $0)")
 __generic<T : __BuiltinArithmeticType, let N : int>
 vector<T,N> WaveMultiPrefixBitAnd(vector<T,N> expr, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixAndMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixAndMultiple(_getMultiPrefixMask(($1).x), $0)")
 matrix<T,N,M> WaveMultiPrefixBitAnd(matrix<T,N,M> expr, uint4 mask);
 
 __generic<T : __BuiltinArithmeticType>
@@ -2841,18 +2986,18 @@ __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 //__target_intrinsic(glsl, "subgroupExclusiveOr($0)")
-__target_intrinsic(cuda, "_wavePrefixOr($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixOr(, _getMultiPrefixMask(($1).x), $0)")
 T WaveMultiPrefixBitOr(T expr, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 //__target_intrinsic(glsl, "subgroupExclusiveOr($0)")
-__target_intrinsic(cuda, "_wavePrefixOrMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixOrMultiple(_getMultiPrefixMask(($1).x), $0)")
 vector<T,N> WaveMultiPrefixBitOr(vector<T,N> expr, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixOrMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixOrMultiple(_getMultiPrefixMask(($1).x), $0)")
 matrix<T,N,M> WaveMultiPrefixBitOr(matrix<T,N,M> expr, uint4 mask);
 
 __generic<T : __BuiltinArithmeticType>
@@ -2860,44 +3005,44 @@ __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveXor($0)")
-__target_intrinsic(cuda, "_wavePrefixXor($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixXor(_getMultiPrefixMask(($1).x), $0)")
 T WaveMultiPrefixBitXor(T expr, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __target_intrinsic(hlsl)
 __glsl_extension(GL_KHR_shader_subgroup_arithmetic)
 __spirv_version(1.3)
 __target_intrinsic(glsl, "subgroupExclusiveXor($0)")
-__target_intrinsic(cuda, "_wavePrefixXorMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixXorMultiple(_getMultiPrefixMask(($1).x), $0)")
 vector<T,N> WaveMultiPrefixBitXor(vector<T,N> expr, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixXorMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixXorMultiple(_getMultiPrefixMask(($1).x), $0)")
 matrix<T,N,M> WaveMultiPrefixBitXor(matrix<T,N,M> expr, uint4 mask);
 
 __generic<T : __BuiltinArithmeticType>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixProduct($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixProduct(_getMultiPrefixMask(($1).x), $0)")
 T WaveMultiPrefixProduct(T value, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int> 
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixProductMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixProductMultiple(_getMultiPrefixMask(($1).x), $0)")
 vector<T,N> WaveMultiPrefixProduct(vector<T,N> value, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixProductMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixProductMultiple(_getMultiPrefixMask(($1).x), $0)")
 matrix<T,N,M> WaveMultiPrefixProduct(matrix<T,N,M> value, uint4 mask);
 
 __generic<T : __BuiltinArithmeticType>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixSum($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixSum(_getMultiPrefixMask(($1).x), $0)")
 T WaveMultiPrefixSum(T value, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixSumMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixSumMultiple(_getMultiPrefixMask(($1).x), $0 )")
 vector<T,N> WaveMultiPrefixSum(vector<T,N> value, uint4 mask);
 __generic<T : __BuiltinArithmeticType, let N : int, let M : int>
 __target_intrinsic(hlsl)
-__target_intrinsic(cuda, "_wavePrefixSumMultiple($0, _getMultiPrefixMask(($1).x))")
+__target_intrinsic(cuda, "_wavePrefixSumMultiple(_getMultiPrefixMask(($1).x), $0)")
 matrix<T,N,M> WaveMultiPrefixSum(matrix<T,N,M> value, uint4 mask);
 
 // `typedef`s to help with the fact that HLSL has been sorta-kinda case insensitive at various points
diff --git a/tests/hlsl-intrinsic/wave-mask/wave.slang b/tests/hlsl-intrinsic/wave-mask/wave.slang
new file mode 100644
index 000000000..6b641906d
--- /dev/null
+++ b/tests/hlsl-intrinsic/wave-mask/wave.slang
@@ -0,0 +1,64 @@
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-cpu -compute 
+//DISABLE_TEST(compute):COMPARE_COMPUTE_EX:-slang -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-slang -compute -dx12 -use-dxil -profile cs_6_0
+//TEST(compute, vulkan):COMPARE_COMPUTE_EX:-vk -compute
+//TEST(compute):COMPARE_COMPUTE_EX:-cuda -compute
+
+//TEST_INPUT:ubuffer(data=[0 0 0 0  0 0 0 0  0 0 0 0  0 0 0 0], stride=4):out,name outputBuffer
+RWStructuredBuffer<int> outputBuffer;
+
+//TEST_INPUT:ubuffer(data=[3 10 2 -1  4 53 4 6  1 2 3 4  7 5 3 1], stride=4):name inputBuffer
+RWStructuredBuffer<int> inputBuffer;
+
+groupshared int sharedMem[32];
+
+int exclusivePrefixSum(WaveMask mask, int index, int waveLaneId, int originalValue, int elementCount)
+{
+    WaveMask localMask = WaveMaskBallot(mask, waveLaneId < elementCount);
+    
+    sharedMem[index] = 0;
+    
+    if(waveLaneId < elementCount)
+    {
+        int temp = 0;
+        int val = originalValue;        
+        
+        for(int  i = 1; i < elementCount; i += i)
+        {
+            int temp = WaveMaskShuffle(localMask, val, waveLaneId - i);            
+            if(waveLaneId >= i)
+            {
+                val += temp;
+            }
+        }        
+        
+        // Make it an exclusive prefix sum
+        val -= originalValue;
+        
+        // Write to shared memory
+        sharedMem[index] = val;
+        
+        // Syncronizes on the mask, and ensures memory fence for shared data write
+        WaveMaskSharedSync(localMask); 
+        return val;
+    }    
+    
+    return 0;    
+}
+
+[numthreads(32, 1, 1)]
+void computeMain(uint3 dispatchThreadID : SV_DispatchThreadID)
+{
+    int index = int(dispatchThreadID.x);
+    const int waveLaneId = WaveGetLaneIndex();
+    
+    const int value = inputBuffer[index];
+    const int elementCount = 9;
+    
+    exclusivePrefixSum(WaveGetActiveMask(), index, waveLaneId, value, elementCount);
+    
+    // It returns the result, but we are going to read from shared memory, to check that aspect worked
+    int prefixValue = sharedMem[index];
+    
+    outputBuffer[index] = prefixValue;
+}
\ No newline at end of file