Added support to fold Shuffle and Broadcast operations.

m4rs-mt · Mar 3, 2022 · 338689c · 338689c
1 parent deaa4ff
commit 338689c
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 27 deletions.
diff --git a/Src/ILGPU.Tests/WarpOperations.cs b/Src/ILGPU.Tests/WarpOperations.cs
@@ -79,10 +79,14 @@ public void WarpBarrier(int length)
             Verify(buffer.View, expected);
         }
 
-        internal static void WarpBroadcastKernel(ArrayView1D<int, Stride1D.Dense> data)
+        internal static void WarpBroadcastKernel(
+            ArrayView1D<int, Stride1D.Dense> data,
+            ArrayView1D<int, Stride1D.Dense> data2,
+            int c)
         {
             var idx = Grid.GlobalIndex.X;
             data[idx] = Warp.Broadcast(Group.IdxX, Warp.WarpSize - 1);
+            data2[idx] = Warp.Broadcast(c, Warp.WarpSize - 2);
         }
 
         [Theory]
@@ -94,21 +98,28 @@ public void WarpBroadcast(int length)
         {
             var warpSize = Accelerator.WarpSize;
             using var buffer = Accelerator.Allocate1D<int>(length * warpSize);
+            using var buffer2 = Accelerator.Allocate1D<int>(length * warpSize);
             var extent = new KernelConfig(
                 length,
                 warpSize);
-            Execute(extent, buffer.View);
+            Execute(extent, buffer.View, buffer2.View, length);
 
             var expected = Enumerable.Repeat(warpSize - 1, length * warpSize).ToArray();
             Verify(buffer.View, expected);
+
+            var expected2 = Enumerable.Repeat(length, length * warpSize).ToArray();
+            Verify(buffer2.View, expected2);
         }
 
         internal static void WarpShuffleKernel(
             Index1D index,
-            ArrayView1D<int, Stride1D.Dense> data)
+            ArrayView1D<int, Stride1D.Dense> data,
+            ArrayView1D<int, Stride1D.Dense> data2,
+            int c)
         {
             var targetIdx = Warp.WarpSize - 1;
             data[index] = Warp.Shuffle(Warp.LaneIdx, targetIdx);
+            data2[index] = Warp.Shuffle(c, targetIdx);
         }
 
         [Theory]
@@ -120,11 +131,15 @@ public void WarpShuffle(int warpMultiplier)
         {
             var length = Accelerator.WarpSize * warpMultiplier;
             using var dataBuffer = Accelerator.Allocate1D<int>(length);
-            Execute(length, dataBuffer.View);
+            using var dataBuffer2 = Accelerator.Allocate1D<int>(length);
+            Execute(length, dataBuffer.View, dataBuffer2.View, length);
 
             var expected = Enumerable.Repeat(
                 Accelerator.WarpSize - 1, length).ToArray();
             Verify(dataBuffer.View, expected);
+
+            var expected2 = Enumerable.Repeat(length, length).ToArray();
+            Verify(dataBuffer2.View, expected2);
         }
 
         internal static void WarpShuffleDownKernel(

diff --git a/Src/ILGPU/IR/Construction/Threads.cs b/Src/ILGPU/IR/Construction/Threads.cs
@@ -48,6 +48,37 @@ public MemoryValue CreateBarrier(
                 GetInitializer(location),
                 kind));
 
+        /// <summary>
+        /// Returns true if the given variable is a constant with respect to a broadcast
+        /// or shuffle value operating on the warp or the group level.
+        /// </summary>
+        /// <param name="variable">The variable to test.</param>
+        /// <returns>
+        /// True, if the given variable is a constant in the parent broadcast or shuffle
+        /// context.
+        /// </returns>
+        private static bool IsShuffleOrBroadcastConstant(Value variable) =>
+            variable switch
+            {
+                // Entry-point parameters can be considered uniform
+                Parameter param => param.Method.HasFlags(MethodFlags.EntryPoint),
+                PrimitiveValue _ => true,
+                WarpSizeValue _ => true,
+                GroupDimensionValue _ => true,
+                GridDimensionValue _ => true,
+                GridIndexValue _ => true,
+                UnaryArithmeticValue unary =>
+                    IsShuffleOrBroadcastConstant(unary.Value),
+                BinaryArithmeticValue binary =>
+                    IsShuffleOrBroadcastConstant(binary.Left) &&
+                    IsShuffleOrBroadcastConstant(binary.Right),
+                TernaryArithmeticValue ternary =>
+                    IsShuffleOrBroadcastConstant(ternary.First) &&
+                    IsShuffleOrBroadcastConstant(ternary.Second) &&
+                    IsShuffleOrBroadcastConstant(ternary.Third),
+                _ => false
+            };
+
         /// <summary>
         /// Creates a new broadcast operation.
         /// </summary>
@@ -63,14 +94,16 @@ public ValueReference CreateBroadcast(
             Value variable,
             Value origin,
             BroadcastKind kind) =>
-            Append(new Broadcast(
-                GetInitializer(location),
-                variable,
-                origin,
-                kind));
+            IsShuffleOrBroadcastConstant(variable)
+                ? variable
+                : Append(new Broadcast(
+                    GetInitializer(location),
+                    variable,
+                    origin,
+                    kind));
 
         /// <summary>
-        /// Creates a new shuffle operation.
+        /// Creates a new shuffle operation involving all lanes of a warp.
         /// </summary>
         /// <param name="location">The current location.</param>
         /// <param name="variable">The variable.</param>
@@ -82,11 +115,13 @@ public ValueReference CreateShuffle(
             Value variable,
             Value origin,
             ShuffleKind kind) =>
-            Append(new WarpShuffle(
-                GetInitializer(location),
-                variable,
-                origin,
-                kind));
+            IsShuffleOrBroadcastConstant(variable)
+                ? variable
+                : Append(new WarpShuffle(
+                    GetInitializer(location),
+                    variable,
+                    origin,
+                    kind));
 
         /// <summary>
         /// Creates a new sub-warp shuffle operation that operates
@@ -104,17 +139,19 @@ public ValueReference CreateShuffle(
             Value origin,
             Value width,
             ShuffleKind kind) =>
-            width is WarpSizeValue
-            ? CreateShuffle(
-                location,
-                variable,
-                origin,
-                kind)
-            : Append(new SubWarpShuffle(
-                GetInitializer(location),
-                variable,
-                origin,
-                width,
-                kind));
+            IsShuffleOrBroadcastConstant(variable)
+                ? (ValueReference)variable
+                : width is WarpSizeValue
+                    ? CreateShuffle(
+                        location,
+                        variable,
+                        origin,
+                        kind)
+                    : Append(new SubWarpShuffle(
+                        GetInitializer(location),
+                        variable,
+                        origin,
+                        width,
+                        kind));
     }
 }