[HLK] Modify Wave match test logic to support modifications in different lanes and vector position

tools/clang/unittests/HLSLExec/LongVectors.cpp

@@ -1614,41 +1614,88 @@ template <typename T> T waveMultiPrefixProduct(T A, UINT) {
 
 template <typename T> struct Op<OpType::WaveMatch, T, 1> : StrictValidation {};
 
+static constexpr UINT ComputeWaveMask(UINT NumWaves) {
+  return (NumWaves < 64) ? (1ULL << NumWaves) - 1 : ~0ULL;
+}
+
+// Helper struct to build the expected result for WaveMatch tests.
+struct WaveMatchResultBuilder {
+
+private:
+  uint64_t LowWaveMask;
+  uint64_t HighWaveMask;
+  uint64_t ActiveLanesLow;
+  uint64_t ActiveLanesHigh;
+
+public:
+  WaveMatchResultBuilder(UINT NumWaves)
+      : ActiveLanesLow(0), ActiveLanesHigh(0) {
+    VERIFY_IS_TRUE(NumWaves <= 128);
+    const UINT LowWaves = std::min(64U, NumWaves);
+    const UINT HighWaves = NumWaves - LowWaves;
+    LowWaveMask = ComputeWaveMask(LowWaves);
+    HighWaveMask = ComputeWaveMask(HighWaves);
+  }
+
+  void SetLane(UINT LaneID) {
+    if (LaneID < 64)
+      ActiveLanesLow |= (1ULL << LaneID) & LowWaveMask;
+    else
+      ActiveLanesHigh |= (1ULL << (LaneID - 64)) & HighWaveMask;
+  }
+
+  void ClearLane(UINT LaneID) {
+    if (LaneID < 64)
+      ActiveLanesLow &= ~(1ULL << LaneID) & LowWaveMask;
+    else
+      ActiveLanesHigh &= ~(1ULL << (LaneID - 64)) & HighWaveMask;
+  }
+
+  void InvertLanes() {
+    ActiveLanesLow = ~ActiveLanesLow & LowWaveMask;
+    ActiveLanesHigh = ~ActiveLanesHigh & HighWaveMask;
+  }
+
+  void ComputeExpected(UINT *Dest) {
+    Dest[0] = static_cast<UINT>(ActiveLanesLow);
+    Dest[1] = static_cast<UINT>(ActiveLanesLow >> 32);
+    Dest[2] = static_cast<UINT>(ActiveLanesHigh);
+    Dest[3] = static_cast<UINT>(ActiveLanesHigh >> 32);
+  }
+};
+
 template <typename T> struct ExpectedBuilder<OpType::WaveMatch, T> {
   static std::vector<UINT> buildExpected(Op<OpType::WaveMatch, T, 1> &,
                                          const InputSets<T> &,
                                          const UINT WaveSize) {
-    // For this test, the shader arranges it so that lane 0 is different from
-    // all the other lanes. Besides that all other lines write their result of
-    // WaveMatch as well.
+    // For this test, the shader arranges it so that lanes 0, WAVE_SIZE/2 and
+    // WAVE_SIZE-1 are different from all the other lanes, also those
+    // lanes modify the vector at positions 0, WAVE_SIZE/2 and WAVE_SIZE-1
+    // respectively, if the input vector has enough elements. Besides that all
+    // other lanes write their result of WaveMatch as well.
 
     std::vector<UINT> Expected;
     Expected.assign(WaveSize * 4, 0);
 
-    const UINT LowWaves = std::min(64U, WaveSize);
-    const UINT HighWaves = WaveSize - LowWaves;
-
-    const uint64_t LowWaveMask =
-        (LowWaves < 64) ? (1ULL << LowWaves) - 1 : ~0ULL;
+    const UINT MidLaneID = WaveSize / 2;
+    const UINT LastLaneID = WaveSize - 1;
 
-    const uint64_t HighWaveMask =
-        (HighWaves < 64) ? (1ULL << HighWaves) - 1 : ~0ULL;
+    WaveMatchResultBuilder UnchangedLanes(WaveSize);
+    UnchangedLanes.InvertLanes();
+    UnchangedLanes.ClearLane(0);
+    UnchangedLanes.ClearLane(MidLaneID);
+    UnchangedLanes.ClearLane(LastLaneID);
 
-    const uint64_t LowExpected = ~1ULL & LowWaveMask;
-    const uint64_t HighExpected = ~0ULL & HighWaveMask;
+    for (UINT LaneID = 0; LaneID < WaveSize; ++LaneID) {
+      const UINT Index = LaneID * 4;
 
-    Expected[0] = 1;
-    Expected[1] = 0;
-    Expected[2] = 0;
-    Expected[3] = 0;
-
-    // all lanes other than the first one have the same result
-    for (UINT I = 1; I < WaveSize; ++I) {
-      const UINT Index = I * 4;
-      Expected[Index] = static_cast<UINT>(LowExpected);
-      Expected[Index + 1] = static_cast<UINT>(LowExpected >> 32);
-      Expected[Index + 2] = static_cast<UINT>(HighExpected);
-      Expected[Index + 3] = static_cast<UINT>(HighExpected >> 32);
+      if (LaneID == 0 || LaneID == MidLaneID || LaneID == LastLaneID) {
+        WaveMatchResultBuilder ChangedLanes(WaveSize);
+        ChangedLanes.SetLane(LaneID);
+        ChangedLanes.ComputeExpected(&Expected[Index]);
+        continue;
+      }
+      UnchangedLanes.ComputeExpected(&Expected[Index]);
     }
 
     return Expected;

tools/clang/unittests/HLSLExec/ShaderOpArith.xml

@@ -4408,15 +4408,24 @@ void MSMain(uint GID : SV_GroupIndex,
         #ifdef FUNC_WAVE_MATCH
         void TestWaveMatch(vector<TYPE, NUM> Vector)
         {
-            if(WaveGetLaneIndex() == 0)
+            uint LaneIndex = WaveGetLaneIndex();
+            bool ShouldModify = ( LaneIndex == 0 || 
+                                  LaneIndex == (WAVE_SIZE / 2) || 
+                                  LaneIndex == (WAVE_SIZE - 1));
+
+            if(ShouldModify && LaneIndex < NUM)
             {
-              if(Vector[0] == (TYPE)0)
-                Vector[0] = (TYPE) 1;
-              else if(Vector[0] == (TYPE)1)
-                Vector[0] = (TYPE) 0;
+              if(Vector[LaneIndex] == (TYPE) 0)
+                Vector[LaneIndex] = (TYPE) 1;
+              else if(Vector[LaneIndex] == (TYPE) 1)
+                Vector[LaneIndex] = (TYPE) 0;
               else
-                Vector[0] = (TYPE) 1;
+                Vector[LaneIndex] = (TYPE) 1;
             }
+
+            // Making sure all lanes finish updating their vectors.
+            AllMemoryBarrierWithGroupSync();
+
             uint4 result = WaveMatch(Vector);
             uint index = WaveGetLaneIndex();