[SLP]Fix mask processing for reused gathered scalars

Need to sync the mask between cost and actual emission to avoid bugs in
mask calculation

Fixes #122324

Author: alexey-bataev

llvm/lib/Transforms/Vectorize/SLPVectorizer.cpp

@@ -10973,7 +10973,19 @@ class BoUpSLP::ShuffleCostEstimator : public BaseShuffleAnalysis {
       }
     }
 
-    ::addMask(CommonMask, ExtMask, /*ExtendingManyInputs=*/true);
+    if (!ExtMask.empty()) {
+      if (CommonMask.empty()) {
+        CommonMask.assign(ExtMask.begin(), ExtMask.end());
+      } else {
+        SmallVector<int> NewMask(ExtMask.size(), PoisonMaskElem);
+        for (int I = 0, Sz = ExtMask.size(); I < Sz; ++I) {
+          if (ExtMask[I] == PoisonMaskElem)
+            continue;
+          NewMask[I] = CommonMask[ExtMask[I]];
+        }
+        CommonMask.swap(NewMask);
+      }
+    }
     if (CommonMask.empty()) {
       assert(InVectors.size() == 1 && "Expected only one vector with no mask");
       return Cost;

llvm/test/Transforms/SLPVectorizer/X86/shuffle-mask-emission.ll

@@ -0,0 +1,43 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -passes=slp-vectorizer -mtriple=x86_64-unknown-linux-gnu < %s | FileCheck %s
+
+define i1 @test() {
+; CHECK-LABEL: define i1 @test() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[H_PROMOTED118_I_FR:%.*]] = freeze i32 1
+; CHECK-NEXT:    [[TMP0:%.*]] = insertelement <4 x i32> <i32 0, i32 0, i32 poison, i32 0>, i32 [[H_PROMOTED118_I_FR]], i32 2
+; CHECK-NEXT:    [[TMP1:%.*]] = xor <4 x i32> zeroinitializer, [[TMP0]]
+; CHECK-NEXT:    [[TMP2:%.*]] = add <4 x i32> zeroinitializer, [[TMP0]]
+; CHECK-NEXT:    [[TMP3:%.*]] = shufflevector <4 x i32> [[TMP1]], <4 x i32> [[TMP2]], <4 x i32> <i32 0, i32 1, i32 6, i32 7>
+; CHECK-NEXT:    [[TMP4:%.*]] = shufflevector <4 x i32> [[TMP0]], <4 x i32> [[TMP3]], <4 x i32> <i32 2, i32 2, i32 7, i32 2>
+; CHECK-NEXT:    [[TMP5:%.*]] = add <4 x i32> [[TMP3]], [[TMP4]]
+; CHECK-NEXT:    [[TMP6:%.*]] = and <4 x i32> [[TMP5]], <i32 0, i32 1, i32 1, i32 1>
+; CHECK-NEXT:    [[TMP7:%.*]] = icmp eq <4 x i32> [[TMP6]], <i32 1, i32 0, i32 0, i32 0>
+; CHECK-NEXT:    [[TMP8:%.*]] = call i1 @llvm.vector.reduce.or.v4i1(<4 x i1> [[TMP7]])
+; CHECK-NEXT:    [[OP_RDX:%.*]] = or i1 [[TMP8]], false
+; CHECK-NEXT:    ret i1 [[OP_RDX]]
+;
+entry:
+  %h.promoted118.i.fr = freeze i32 1
+  %invariant.op.i51 = add i32 %h.promoted118.i.fr, 0
+  %conv25.i = xor i32 0, 0
+  %add.i.i = add i32 %conv25.i, %h.promoted118.i.fr
+  %sext.i.mask = and i32 %add.i.i, 0
+  %cmp27.i = icmp eq i32 %sext.i.mask, 1
+  %0 = or i1 %cmp27.i, false
+  %conv25.i.1 = add i32 0, 0
+  %add.i.i.1 = add i32 %conv25.i.1, %h.promoted118.i.fr
+  %sext.i.1.mask = and i32 %add.i.i.1, 1
+  %cmp27.i.1 = icmp eq i32 %sext.i.1.mask, 0
+  %conv25.1.i.1 = xor i32 0, 0
+  %add.i.1.i.1 = add i32 %conv25.1.i.1, %h.promoted118.i.fr
+  %sext.1.i.1.mask = and i32 %add.i.1.i.1, 1
+  %cmp27.1.i.1 = icmp eq i32 %sext.1.i.1.mask, 0
+  %add.i.2.reass.i.1 = add i32 %invariant.op.i51, %conv25.i.1
+  %sext.2.i.1.mask = and i32 %add.i.2.reass.i.1, 1
+  %cmp27.2.i.1 = icmp eq i32 %sext.2.i.1.mask, 0
+  %1 = or i1 %cmp27.1.i.1, %cmp27.2.i.1
+  %2 = or i1 %cmp27.i.1, %1
+  %3 = or i1 %0, %2
+  ret i1 %3
+}