[llvm-branch-commits] [VectorCombine] Fix miscompile caused by collisions in foldShuffleToIdentity (PR #188668)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: wanglei (wangleiat)

<details>
<summary>Changes</summary>

`foldShuffleToIdentity` previously used a `SmallPtrSet<Use *, 4>` keyed
only on the first lane's `Use *` to track equivalent sequences. This
caused incorrect folding when different shuffle sequences shared the
same initial `Use *`.

This patch fixes the miscompile by tracking the full `InstLane` sequence
via `SmallVector<SmallVector<InstLane>>`, ensuring exact sequence
matching.

It also corrects an implicit `Use *` to `Value *` conversion by changing
`dyn_cast<Constant>(FrontU)` to `dyn_cast<Constant>(FrontU->get())`.

Fixes #<!-- -->180338


---
Full diff: https://github.com/llvm/llvm-project/pull/188668.diff


2 Files Affected:

- (modified) llvm/lib/Transforms/Vectorize/VectorCombine.cpp (+20-16) 
- (modified) llvm/test/Transforms/VectorCombine/LoongArch/shuffle-identity-miscompile.ll (+4-2) 


``````````diff
diff --git a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
index 136f4fe238b34..42680c1fe272d 100644
--- a/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
+++ b/llvm/lib/Transforms/Vectorize/VectorCombine.cpp
@@ -3509,22 +3509,26 @@ static bool isFreeConcat(ArrayRef<InstLane> Item, TTI::TargetCostKind CostKind,
   return true;
 }
 
-static Value *generateNewInstTree(ArrayRef<InstLane> Item, FixedVectorType *Ty,
-                                  const SmallPtrSet<Use *, 4> &IdentityLeafs,
-                                  const SmallPtrSet<Use *, 4> &SplatLeafs,
-                                  const SmallPtrSet<Use *, 4> &ConcatLeafs,
-                                  IRBuilderBase &Builder,
-                                  const TargetTransformInfo *TTI) {
+static Value *
+generateNewInstTree(ArrayRef<InstLane> Item, FixedVectorType *Ty,
+                    const SmallVectorImpl<SmallVector<InstLane>> &IdentityLeafs,
+                    const SmallVectorImpl<SmallVector<InstLane>> &SplatLeafs,
+                    const SmallVectorImpl<SmallVector<InstLane>> &ConcatLeafs,
+                    IRBuilderBase &Builder, const TargetTransformInfo *TTI) {
   auto [FrontU, FrontLane] = Item.front();
 
-  if (IdentityLeafs.contains(FrontU)) {
+  auto Contains = [](const SmallVectorImpl<SmallVector<InstLane>> &Set,
+                     ArrayRef<InstLane> Item) {
+    return any_of(Set, [&](ArrayRef<InstLane> E) { return E == Item; });
+  };
+
+  if (Contains(IdentityLeafs, Item))
     return FrontU->get();
-  }
-  if (SplatLeafs.contains(FrontU)) {
+  if (Contains(SplatLeafs, Item)) {
     SmallVector<int, 16> Mask(Ty->getNumElements(), FrontLane);
     return Builder.CreateShuffleVector(FrontU->get(), Mask);
   }
-  if (ConcatLeafs.contains(FrontU)) {
+  if (Contains(ConcatLeafs, Item)) {
     unsigned NumElts =
         cast<FixedVectorType>(FrontU->get()->getType())->getNumElements();
     SmallVector<Value *> Values(Item.size() / NumElts, nullptr);
@@ -3613,7 +3617,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
 
   SmallVector<SmallVector<InstLane>> Worklist;
   Worklist.push_back(Start);
-  SmallPtrSet<Use *, 4> IdentityLeafs, SplatLeafs, ConcatLeafs;
+  SmallVector<SmallVector<InstLane>> IdentityLeafs, SplatLeafs, ConcatLeafs;
   unsigned NumVisited = 0;
 
   while (!Worklist.empty()) {
@@ -3642,11 +3646,11 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
           return !E.value().first || (IsEquiv(E.value().first->get(), FrontV) &&
                                       E.value().second == (int)E.index());
         })) {
-      IdentityLeafs.insert(FrontU);
+      IdentityLeafs.push_back(std::move(Item));
       continue;
     }
     // Look for constants, for the moment only supporting constant splats.
-    if (auto *C = dyn_cast<Constant>(FrontU);
+    if (auto *C = dyn_cast<Constant>(FrontU->get());
         C && C->getSplatValue() &&
         all_of(drop_begin(Item), [Item](InstLane &IL) {
           Value *FrontV = Item.front().first->get();
@@ -3655,7 +3659,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
                         cast<Constant>(U->get())->getSplatValue() ==
                             cast<Constant>(FrontV)->getSplatValue());
         })) {
-      SplatLeafs.insert(FrontU);
+      SplatLeafs.push_back(std::move(Item));
       continue;
     }
     // Look for a splat value.
@@ -3664,7 +3668,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
           auto [U, Lane] = IL;
           return !U || (U->get() == FrontU->get() && Lane == FrontLane);
         })) {
-      SplatLeafs.insert(FrontU);
+      SplatLeafs.push_back(std::move(Item));
       continue;
     }
 
@@ -3749,7 +3753,7 @@ bool VectorCombine::foldShuffleToIdentity(Instruction &I) {
     }
 
     if (isFreeConcat(Item, CostKind, TTI)) {
-      ConcatLeafs.insert(FrontU);
+      ConcatLeafs.push_back(std::move(Item));
       continue;
     }
 
diff --git a/llvm/test/Transforms/VectorCombine/LoongArch/shuffle-identity-miscompile.ll b/llvm/test/Transforms/VectorCombine/LoongArch/shuffle-identity-miscompile.ll
index 94815d8a604bc..1582dc4c59708 100644
--- a/llvm/test/Transforms/VectorCombine/LoongArch/shuffle-identity-miscompile.ll
+++ b/llvm/test/Transforms/VectorCombine/LoongArch/shuffle-identity-miscompile.ll
@@ -7,8 +7,10 @@ define i32 @shuffle_different_lanes() local_unnamed_addr {
 ; CHECK-NEXT:  [[ENTRY:.*:]]
 ; CHECK-NEXT:    [[TMP0:%.*]] = call <2 x i64> @llvm.loongarch.lsx.vadd.d(<2 x i64> <i64 8067293510808050, i64 0>, <2 x i64> zeroinitializer)
 ; CHECK-NEXT:    [[TMP1:%.*]] = trunc <2 x i64> [[TMP0]] to <2 x i16>
-; CHECK-NEXT:    [[TMP2:%.*]] = trunc <2 x i64> [[TMP0]] to <2 x i16>
-; CHECK-NEXT:    [[TMP3:%.*]] = bitcast <2 x i16> zeroinitializer to i32
+; CHECK-NEXT:    [[TMP2:%.*]] = shufflevector <2 x i64> [[TMP0]], <2 x i64> poison, <2 x i32> <i32 1, i32 1>
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc <2 x i64> [[TMP2]] to <2 x i16>
+; CHECK-NEXT:    [[SHUFFLE_I:%.*]] = sub <2 x i16> [[TMP1]], [[TMP4]]
+; CHECK-NEXT:    [[TMP3:%.*]] = bitcast <2 x i16> [[SHUFFLE_I]] to i32
 ; CHECK-NEXT:    ret i32 [[TMP3]]
 ;
 entry:

``````````

</details>


https://github.com/llvm/llvm-project/pull/188668