Fix #18

Author: planetis-m

src/computesim/core.nim

@@ -1,5 +1,6 @@
 # (c) 2024 Antonis Geralis
 import threading/barrier, vectors
+from std/math import isPowerOfTwo
 
 const
   SubgroupSize* {.intdefine.} = 8
@@ -23,6 +24,8 @@ const
       )
     masks
 
+static: assert isPowerOfTwo(SubgroupSize), "SubgroupSize must be a power of two"
+
 type
   ValueType* = enum
     Bool, Int, Uint, Float, Double # Scalar types

src/computesim/subgroupops.nim

@@ -216,7 +216,7 @@ defineSubgroupOp(execShuffle):
   var shuffledVals {.noinit.}: array[SubgroupSize, RawValue]
   # First gather all shuffled values into array
   for threadId in threadsInGroup(group):
-    let srcThreadId = commands[threadId].dirty
+    let srcThreadId = commands[threadId].dirty and SubgroupSize - 1
     # Check if source thread is valid within the group
     var found = false
     for validId in threadsInGroup(group):
@@ -225,7 +225,7 @@ defineSubgroupOp(execShuffle):
         break
     # If source thread is valid, take its value
     # Otherwise use this thread's own value
-    shuffledVals[threadId] = commands[if found: srcThreadId else: threadId].val
+    shuffledVals[threadId] = if found: commands[srcThreadId].val else: default(RawValue)
 
   let valueType = commands[firstThreadId].t
   # Then construct results
@@ -245,7 +245,7 @@ defineSubgroupOp(execShuffleXor):
   var shuffledVals {.noinit.}: array[SubgroupSize, RawValue]
   # First gather all shuffled values into array
   for threadId in threadsInGroup(group):
-    let srcThreadId = threadId xor commands[threadId].dirty
+    let srcThreadId = threadId xor commands[threadId].dirty and SubgroupSize - 1
     # Check if source thread is valid within the group
     var found = false
     for validId in threadsInGroup(group):
@@ -254,7 +254,7 @@ defineSubgroupOp(execShuffleXor):
         break
     # If source thread is valid, take its value
     # Otherwise use this thread's own value
-    shuffledVals[threadId] = commands[if found: srcThreadId else: threadId].val
+    shuffledVals[threadId] = if found: commands[srcThreadId].val else: default(RawValue)
 
   let valueType = commands[firstThreadId].t
   # Then construct results
@@ -273,13 +273,13 @@ defineSubgroupOp(execShuffleXor):
 defineSubgroupOp(execShuffleDown):
   var shuffledVals {.noinit.}: array[SubgroupSize, RawValue]
   for threadId in threadsInGroup(group):
-    let srcThreadId = threadId + commands[threadId].dirty
+    let srcThreadId = (threadId + commands[threadId].dirty) and SubgroupSize - 1
     var found = false
     for validId in threadsInGroup(group):
       if validId == srcThreadId:
         found = true
         break
-    shuffledVals[threadId] = commands[if found: srcThreadId else: threadId].val
+    shuffledVals[threadId] = if found: commands[srcThreadId].val else: default(RawValue)
 
   let valueType = commands[firstThreadId].t
   for threadId in threadsInGroup(group):
@@ -298,16 +298,13 @@ defineSubgroupOp(execShuffleUp):
   var shuffledVals {.noinit.}: array[SubgroupSize, RawValue]
   for threadId in threadsInGroup(group):
     # Convert to signed for safe subtraction
-    let srcThreadId = if threadId >= commands[threadId].dirty:
-      threadId - commands[threadId].dirty
-    else:
-      threadId
+    let srcThreadId = (threadId - commands[threadId].dirty) and SubgroupSize - 1
     var found = false
     for validId in threadsInGroup(group):
       if validId == srcThreadId:
         found = true
         break
-    shuffledVals[threadId] = commands[if found: srcThreadId else: threadId].val
+    shuffledVals[threadId] = if found: commands[srcThreadId].val else: default(RawValue)
 
   let valueType = commands[firstThreadId].t
   for threadId in threadsInGroup(group):

tests/config.nims

@@ -7,7 +7,7 @@
 --define:"ThreadPoolSize=10"
 # --define:"SubgroupSize=32"
 --define:debugSubgroup
-switch("define", "debugSubgroupID:1")
+# switch("define", "debugSubgroupID:1")
 
 when not defined(windows):
   --debugger:"native"

tests/tshuffle.nim

@@ -0,0 +1,53 @@
+import std/math, computesim
+
+proc calculate(output: ptr seq[int32]; numElements: uint32) {.computeShader.} =
+  let tid = gl_GlobalInvocationID.x
+  var value: int32 = 0
+  # Loop from 1 to 2
+  for i in 1'i32..2:
+    if (tid.int32 + i) mod 3 == 0: # skip iteration
+      continue
+    value = subgroupShuffle(tid.int32 + 1, tid.uint32 + i.uint32)
+  # Store the result in the SSBO
+  output[tid] = value
+
+const
+  NumElements = 64'u32
+  WorkGroupSize = 32'u32 # Force underutilization of hardware subgroups
+
+proc main() =
+  # Set up compute dimensions
+  let numWorkGroups = uvec3(ceilDiv(NumElements, WorkGroupSize), 1, 1)
+  let workGroupSize = uvec3(WorkGroupSize, 1, 1)
+
+  # Initialize buffer
+  let output = newSeq[int32](NumElements)
+
+  # Run reduction on CPU
+  runComputeOnCpu(
+    numWorkGroups = numWorkGroups,
+    workGroupSize = workGroupSize,
+    compute = calculate,
+    ssbo = addr output,
+    args = NumElements
+  )
+
+  assert output == @[3'i32, 0, 0, 6, 0, 0, 1, 1, 0, 12, 0, 0, 15, 0, 9, 10,
+                     0, 0, 21, 0, 0, 24, 0, 18, 27, 0, 0, 30, 0, 0, 25, 25,
+                     0, 36, 0, 0, 39, 0, 33, 34, 0, 0, 45, 0, 0, 48, 0, 42,
+                     51, 0, 0, 54, 0, 0, 49, 49, 0, 60, 0, 0, 63, 0, 57, 58]
+
+# Debug Output:
+# - SubgroupID 0
+# [Shuffle #1] inputs {t0: 1, t1: 2, t3: 4, t4: 5, t6: 7, t7: 8} | shuffled: [2, 0, 5, 0, 8, 1]
+# [Shuffle #1] inputs {t0: 1, t2: 3, t3: 4, t5: 6, t6: 7} | shuffled: [3, 0, 6, 0, 1]
+# - SubgroupID 1
+# [Shuffle #1] inputs {t1: 10, t2: 11, t4: 13, t5: 14, t7: 16} | shuffled: [11, 0, 14, 0, 0]
+# [Shuffle #1] inputs {t0: 9, t1: 10, t3: 12, t4: 13, t6: 15, t7: 16} | shuffled: [0, 12, 0, 15, 9, 10]
+# - SubgroupID 2
+# [Shuffle #1] inputs {t0: 17, t2: 19, t3: 20, t5: 22, t6: 23} | shuffled: [0, 20, 0, 23, 0]
+# [Shuffle #1] inputs {t1: 18, t2: 19, t4: 21, t5: 22, t7: 24} | shuffled: [0, 21, 0, 24, 18]
+# Output Buffer:
+# [3, 0, 0, 6, 0, 0, 1, 1, 0, 12, 0, 0, 15, 0, 9, 10, 0, 0, 21, 0, 0, 24, 0, 18, 27, 0, 0, ...
+# - Matches GLSL shader output.
+main()