[InstCombine] Fix incorrect insert point when folding zext of deinterleave

[mshockwave]

Fix invalid module crash in #195330 (comment)

The problem was caused by incorrect IRBuilder insertion point. The insertion point used when generating new instructions might not dominate all the de-interleave fields users. This patch fixes this issue by explicitly setting the insertion point to either llvm.vector.deinterleave2 or the earliest shufflevector instruction (in the de-interleaving group) within the same basicblock.

[llvmorg-github-actions]

@llvm/pr-subscribers-llvm-transforms

Author: Min-Yih Hsu (mshockwave)

Changes

Fix invalid module crash in #195330 (comment)

The problem was caused by incorrect IRBuilder insertion point. The insertion point used when generating new instructions might not dominate all the de-interleave fields users. This patch fixes this issue by explicitly setting the insertion point to either llvm.vector.deinterleave2 or the earliest shufflevector instruction (in the de-interleaving group) within the same basicblock.

---

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

2 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp (+14-4)
• (modified) llvm/test/Transforms/InstCombine/fold-zext-of-deinterleave.ll (+26)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
index 3a7fbbcb468da..5abad37e173c4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -3332,11 +3332,12 @@ InstCombinerImpl::foldExtractionOfVectorDeinterleave(ZExtInst &RootZExt) {
   Value *DIV;
 
   using namespace PatternMatch;
-  Value *SVI = nullptr, *DI = nullptr;
+  Value *SV = nullptr, *DI = nullptr;
+  Instruction *SVI = nullptr;
   if (!match(&RootZExt,
              m_ZExt(m_CombineOr(
                  m_ExtractValue(m_Value(DI, m_Deinterleave2(m_Value(DIV)))),
-                 m_Value(SVI, m_Shuffle(m_Value(), m_Value()))))))
+                 m_Value(SV, m_Shuffle(m_Value(), m_Value()))))))
     return nullptr;
 
   auto isDeinterleaveShuffle =
@@ -3361,9 +3362,10 @@ InstCombinerImpl::foldExtractionOfVectorDeinterleave(ZExtInst &RootZExt) {
 
   // Validate either the shufflevector or the vector.deinterleave2 and obtain
   // the value they're de-interleaving.
-  if (SVI) {
+  if (SV) {
+    SVI = cast<Instruction>(SV);
     // We will find other shufflevectors later.
-    DIV = isDeinterleaveShuffle(cast<Instruction>(SVI)).first;
+    DIV = isDeinterleaveShuffle(SVI).first;
     if (!DIV)
       return nullptr;
   } else {
@@ -3396,6 +3398,11 @@ InstCombinerImpl::foldExtractionOfVectorDeinterleave(ZExtInst &RootZExt) {
       if (V != DIV)
         continue;
       assert(Index < 2);
+      // Find the earliest field extraction instruction.
+      if (FieldI->getParent() != SVI->getParent())
+        continue;
+      if (FieldI != SVI && FieldI->comesBefore(SVI))
+        SVI = FieldI;
       Fields.push_back({FieldI, Index});
     }
   } else {
@@ -3425,6 +3432,9 @@ InstCombinerImpl::foldExtractionOfVectorDeinterleave(ZExtInst &RootZExt) {
     }
   }
 
+  // This will insert replacement instructions before all the fields users.
+  Builder.SetInsertPoint(DI ? cast<Instruction>(DI) : SVI);
+
   // Double the element size but half the vector length.
   auto *BitcastedTy = VectorType::getExtendedElementVectorType(InputVecTy);
   BitcastedTy = VectorType::getHalfElementsVectorType(BitcastedTy);
diff --git a/llvm/test/Transforms/InstCombine/fold-zext-of-deinterleave.ll b/llvm/test/Transforms/InstCombine/fold-zext-of-deinterleave.ll
index a5345cee69929..796d466c13a0c 100644
--- a/llvm/test/Transforms/InstCombine/fold-zext-of-deinterleave.ll
+++ b/llvm/test/Transforms/InstCombine/fold-zext-of-deinterleave.ll
@@ -125,6 +125,32 @@ define <vscale x 8 x i64> @to_bitcast_deinterleave(ptr %p) {
   ret <vscale x 8 x i64> %r
 }
 
+; Regression test for a crash: https://github.com/llvm/llvm-project/pull/195330#issuecomment-4635245035
+define <16 x i16> @zext_shufflevector_correct_insertpoint(<16 x i8> %load) {
+; CHECK-LABEL: define <16 x i16> @zext_shufflevector_correct_insertpoint(
+; CHECK-SAME: <16 x i8> [[LOAD:%.*]]) {
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze <16 x i8> [[LOAD]]
+; CHECK-NEXT:    [[TMP2:%.*]] = bitcast <16 x i8> [[TMP1]] to <8 x i16>
+; CHECK-NEXT:    [[ZEXT:%.*]] = and <8 x i16> [[TMP2]], splat (i16 255)
+; CHECK-NEXT:    [[ZEXT9:%.*]] = lshr <8 x i16> [[TMP2]], splat (i16 8)
+; CHECK-NEXT:    [[ADD:%.*]] = add nuw nsw <8 x i16> [[ZEXT]], splat (i16 1)
+; CHECK-NEXT:    [[TMP3:%.*]] = or <8 x i16> [[ADD]], [[ZEXT9]]
+; CHECK-NEXT:    [[OR:%.*]] = shufflevector <8 x i16> [[TMP3]], <8 x i16> zeroinitializer, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+; CHECK-NEXT:    ret <16 x i16> [[OR]]
+;
+  %shufflevector4 = shufflevector <16 x i8> %load, <16 x i8> zeroinitializer, <8 x i32> <i32 1, i32 3, i32 5, i32 7, i32 9, i32 11, i32 13, i32 15>
+  %shufflevector = shufflevector <16 x i8> %load, <16 x i8> zeroinitializer, <8 x i32> <i32 0, i32 2, i32 4, i32 6, i32 8, i32 10, i32 12, i32 14>
+  %zext = zext <8 x i8> %shufflevector to <8 x i16>
+  %shufflevector7 = shufflevector <8 x i8> %shufflevector4, <8 x i8> zeroinitializer, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison, i32 poison>
+  %shufflevector8 = shufflevector <16 x i8> %shufflevector7, <16 x i8> zeroinitializer, <8 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7>
+  %add = add <8 x i16> %zext, splat (i16 1)
+  %shufflevector6 = shufflevector <8 x i16> %add, <8 x i16> zeroinitializer, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %zext9 = zext <8 x i8> %shufflevector8 to <8 x i16>
+  %shufflevector10 = shufflevector <8 x i16> %zext9, <8 x i16> zeroinitializer, <16 x i32> <i32 0, i32 1, i32 2, i32 3, i32 4, i32 5, i32 6, i32 7, i32 8, i32 9, i32 10, i32 11, i32 12, i32 13, i32 14, i32 15>
+  %or = or <16 x i16> %shufflevector6, %shufflevector10
+  ret <16 x i16> %or
+}
+
 define <8 x i64> @negative_zext_shufflevector_not_zext(<16 x i32> %v) {
 ; CHECK-LABEL: define <8 x i64> @negative_zext_shufflevector_not_zext(
 ; CHECK-SAME: <16 x i32> [[V:%.*]]) {

[topperc]

LGTM