WaveShuffleIndexSinking Pass shuffleGroup by basic block

The WaveShuffleIndexSinking Pass doesn't consider program control flow,
potentially producing invalid IR if users aren't dominated. It merges
instructions globally, risking undefined instructions if blocks are
entered out of order, leading to incorrect calculations. Now,
shuffleGroups are formed by channel and within the same basic block.

Author: MiloszSkobejko

IGC/Compiler/Optimizer/WaveShuffleIndexSinking.cpp

@@ -367,7 +367,7 @@ namespace IGC
         static bool isHoistable( BinaryOperator* inst );
         static bool isHoistableOverAnchor( BinaryOperator* instToHoist, BinaryOperator* anchorInst );
         Function& F;
-        DenseMap<Value*, SmallVector<ShuffleGroup, 4>> ShuffleGroupMap;
+        DenseMap<std::pair<BasicBlock*, Value*>, SmallVector<ShuffleGroup, 4>> ShuffleGroupMap;
         DenseSet<WaveShuffleIndexIntrinsic*> Visited;
     };
 
@@ -451,7 +451,7 @@ bool WaveShuffleIndexSinkingImpl::splitWaveShuffleIndexes()
 bool WaveShuffleIndexSinkingImpl::mergeWaveShuffleIndexes()
 {
     // Map from Source to (Map from Lane to list of duplicate instructions)
-    DenseMap<Value*, DenseMap<ConstantInt*, SmallVector<WaveShuffleIndexIntrinsic*>>> mergeMap;
+    DenseMap<std::pair<BasicBlock*, Value*>, DenseMap<ConstantInt*, SmallVector<WaveShuffleIndexIntrinsic*>>> mergeMap;
     for( auto& BB : F )
     {
         for( auto& I : BB )
@@ -460,7 +460,7 @@ bool WaveShuffleIndexSinkingImpl::mergeWaveShuffleIndexes()
             {
                 if( auto* constantChannel = dyn_cast<ConstantInt>( waveShuffleInst->getChannel() ) )
                 {
-                    mergeMap[ waveShuffleInst->getSrc() ][ constantChannel ].push_back( waveShuffleInst );
+                    mergeMap[ {&BB, waveShuffleInst->getSrc()} ][ constantChannel ].push_back( waveShuffleInst );
                 }
             }
         }
@@ -505,11 +505,13 @@ void WaveShuffleIndexSinkingImpl::gatherShuffleGroups()
                     // Save compute and do not re-process/ create a new ShuffleGroup
                     continue;
                 }
-                if( ShuffleGroupMap.count( waveShuffleInst->getSrc() ) )
+
+                std::pair<BasicBlock*, Value*> bbShuffleGroup = { &BB, waveShuffleInst->getSrc() };
+                if ( ShuffleGroupMap.count( bbShuffleGroup ) )
                 {
                     // Found existing group(s) with the same source, try to match with one of the groups
                     bool match = false;
-                    for( auto& shuffleGroup : ShuffleGroupMap[ waveShuffleInst->getSrc() ] )
+                    for (auto& shuffleGroup : ShuffleGroupMap[ bbShuffleGroup ] )
                     {
                         if( shuffleGroup.match( waveShuffleInst ) )
                         {
@@ -521,13 +523,13 @@ void WaveShuffleIndexSinkingImpl::gatherShuffleGroups()
                     // create new ShuffleGroup since no suitable match was found
                     if( !match )
                     {
-                        ShuffleGroupMap[ waveShuffleInst->getSrc() ].emplace_back( waveShuffleInst );
+                        ShuffleGroupMap[ {&BB, waveShuffleInst->getSrc()} ].emplace_back( waveShuffleInst );
                     }
                 }
                 else
                 {
                     // create new ShuffleGroup for broadcast operations
-                    ShuffleGroupMap[ waveShuffleInst->getSrc() ].emplace_back( waveShuffleInst );
+                    ShuffleGroupMap[ {&BB, waveShuffleInst->getSrc()} ].emplace_back( waveShuffleInst );
                 }
             }
         }

IGC/Compiler/tests/WaveShuffleIndexSinking/constant-src-different-bb.ll

@@ -0,0 +1,46 @@
+;=========================== begin_copyright_notice ============================
+;
+; Copyright (C) 2025 Intel Corporation
+;
+; SPDX-License-Identifier: MIT
+;
+;============================ end_copyright_notice =============================
+; RUN: igc_opt --typed-pointers -igc-wave-shuffle-index-sinking -S < %s | FileCheck %s
+; ------------------------------------------------
+; WaveShuffleIndexSinking
+;
+; This test checks if WaveShuffleIndex instructions are not merged when they are identical, but
+; in separate basic blocks. This behavior helps to avoid potentially invalid LLVM IR being
+; produced by not dominating all users.
+; ------------------------------------------------
+
+define void @test_compare_cases(i32* %dst0, i32* %dst1, i1 %condition) {
+entry:
+  br i1 %condition, label %bb0, label %bb1
+
+bb0:
+; CHECK: [[SHUFFLE0:%.*]] = call i8 @llvm.genx.GenISA.WaveShuffleIndex.i8(i8 1, i32 0, i32 0)
+  %simdShuffle0 = call i8 @llvm.genx.GenISA.WaveShuffleIndex.i8(i8 1, i32 0, i32 0)
+  %and0 = and i8 %simdShuffle0, 1
+  %cmp0 = icmp eq i8 %and0, 0
+  %zxt0 = zext i1 %cmp0 to i32
+  store i32 %zxt0, i32* %dst0
+  br label %exit
+
+bb1:
+; CHECK: [[SHUFFLE1:%.*]] = call i8 @llvm.genx.GenISA.WaveShuffleIndex.i8(i8 1, i32 0, i32 0)
+  %simdShuffle1 = call i8 @llvm.genx.GenISA.WaveShuffleIndex.i8(i8 1, i32 0, i32 0)
+  %and1 = and i8 %simdShuffle1, 1
+  %cmp1 = icmp eq i8 %and1, 0
+  %zxt1 = zext i1 %cmp1 to i32
+  store i32 %zxt1, i32* %dst1
+  br label %exit
+
+exit:
+  ret void
+}
+
+; Function Attrs: convergent nounwind readnone
+declare i8 @llvm.genx.GenISA.WaveShuffleIndex.i8(i8, i32, i32) #0
+
+attributes #0 = { convergent nounwind readnone }