[SandboxVec][BottomUpVec] Fix packing when PHIs are present (llvm#124206)

vporpo · web-flow · commit d2234ca16310 · 2025-01-23T16:29:01.000-08:00
Before this patch we might have emitted pack instructions in between PHI
nodes. This patch fixes it by fixing the insert point of the new packs.
diff --git a/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h b/llvm/include/llvm/Transforms/Vectorize/SandboxVectorizer/VecUtils.h
@@ -137,6 +137,19 @@ class VecUtils {
     }
     return LowestI;
   }
+
+  /// If \p I is not a PHI it returns it. Else it walks down the instruction
+  /// chain looking for the last PHI and returns it. \Returns nullptr if \p I is
+  /// nullptr.
+  static Instruction *getLastPHIOrSelf(Instruction *I) {
+    Instruction *LastI = I;
+    while (I != nullptr && isa<PHINode>(I)) {
+      LastI = I;
+      I = I->getNextNode();
+    }
+    return LastI;
+  }
+
   /// If all values in \p Bndl are of the same scalar type then return it,
   /// otherwise return nullptr.
   static Type *tryGetCommonScalarType(ArrayRef<Value *> Bndl) {
diff --git a/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp b/llvm/lib/Transforms/Vectorize/SandboxVectorizer/Passes/BottomUpVec.cpp
@@ -47,11 +47,13 @@ static SmallVector<Value *, 4> getOperand(ArrayRef<Value *> Bndl,
 /// of BB if no instruction found in \p Vals.
 static BasicBlock::iterator getInsertPointAfterInstrs(ArrayRef<Value *> Vals,
                                                       BasicBlock *BB) {
-  auto *BotI = VecUtils::getLowest(Vals);
+  auto *BotI = VecUtils::getLastPHIOrSelf(VecUtils::getLowest(Vals));
   if (BotI == nullptr)
-    // We are using BB->begin() as the fallback insert point if `ToPack` did
-    // not contain instructions.
-    return BB->begin();
+    // We are using BB->begin() (or after PHIs) as the fallback insert point.
+    return BB->empty()
+               ? BB->begin()
+               : std::next(
+                     VecUtils::getLastPHIOrSelf(&*BB->begin())->getIterator());
   return std::next(BotI->getIterator());
 }
 
diff --git a/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll b/llvm/test/Transforms/SandboxVectorizer/bottomup_basic.ll
@@ -269,3 +269,31 @@ define void @diamondMultiInput(ptr %ptr, ptr %ptrX) {
   store float %sub1, ptr %ptr1
   ret void
 }
+
+define void @diamondWithConstantVector(ptr %ptr) {
+; CHECK-LABEL: define void @diamondWithConstantVector(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:    [[GEPA0:%.*]] = getelementptr i32, ptr [[PTR]], i64 0
+; CHECK-NEXT:    [[GEPB0:%.*]] = getelementptr i32, ptr [[PTR]], i64 10
+; CHECK-NEXT:    store <2 x i32> zeroinitializer, ptr [[GEPA0]], align 4
+; CHECK-NEXT:    store <2 x i32> zeroinitializer, ptr [[GEPB0]], align 4
+; CHECK-NEXT:    ret void
+;
+  %gepA0 = getelementptr i32, ptr %ptr, i64 0
+  %gepA1 = getelementptr i32, ptr %ptr, i64 1
+
+  %gepB0 = getelementptr i32, ptr %ptr, i64 10
+  %gepB1 = getelementptr i32, ptr %ptr, i64 11
+
+  %zext0 = zext i16 0 to i32
+  %zext1 = zext i16 0 to i32
+
+  store i32 %zext0, ptr %gepA0
+  store i32 %zext1, ptr %gepA1
+
+  %orB0 = or i32 0, %zext0
+  %orB1 = or i32 0, %zext1
+  store i32 %orB0, ptr %gepB0
+  store i32 %orB1, ptr %gepB1
+  ret void
+}
diff --git a/llvm/test/Transforms/SandboxVectorizer/pack.ll b/llvm/test/Transforms/SandboxVectorizer/pack.ll
@@ -14,3 +14,77 @@ define void @pack_constants(ptr %ptr) {
   store i8 1, ptr %ptr1
   ret void
 }
+
+; Make sure we don't generate bad IR when packing PHIs.
+; NOTE: This test may become obsolete once we start vectorizing PHIs.
+define void @packPHIs(ptr %ptr) {
+; CHECK-LABEL: define void @packPHIs(
+; CHECK-SAME: ptr [[PTR:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[PHI0:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[PHI1:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[PHI2:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[PHI3:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[PACK:%.*]] = insertelement <2 x i8> poison, i8 [[PHI0]], i32 0
+; CHECK-NEXT:    [[PACK1:%.*]] = insertelement <2 x i8> [[PACK]], i8 [[PHI1]], i32 1
+; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 0
+; CHECK-NEXT:    store <2 x i8> [[PACK1]], ptr [[GEP0]], align 1
+; CHECK-NEXT:    br label %[[LOOP]]
+; CHECK:       [[EXIT:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  br label %loop
+
+loop:
+  %phi0 = phi i8 [0, %entry], [1, %loop]
+  %phi1 = phi i8 [0, %entry], [1, %loop]
+  %phi2 = phi i8 [0, %entry], [1, %loop]
+  %phi3 = phi i8 [0, %entry], [1, %loop]
+  %gep0 = getelementptr i8, ptr %ptr, i64 0
+  %gep1 = getelementptr i8, ptr %ptr, i64 1
+  store i8 %phi0, ptr %gep0
+  store i8 %phi1, ptr %gep1
+  br label %loop
+
+exit:
+  ret void
+}
+
+define void @packFromOtherBB(ptr %ptr, i8 %val) {
+; CHECK-LABEL: define void @packFromOtherBB(
+; CHECK-SAME: ptr [[PTR:%.*]], i8 [[VAL:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[ADD0:%.*]] = add i8 [[VAL]], 0
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i8 [[VAL]], 1
+; CHECK-NEXT:    [[PACK:%.*]] = insertelement <2 x i8> poison, i8 [[ADD0]], i32 0
+; CHECK-NEXT:    [[PACK1:%.*]] = insertelement <2 x i8> [[PACK]], i8 [[MUL1]], i32 1
+; CHECK-NEXT:    br label %[[LOOP:.*]]
+; CHECK:       [[LOOP]]:
+; CHECK-NEXT:    [[PHI0:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[PHI1:%.*]] = phi i8 [ 0, %[[ENTRY]] ], [ 1, %[[LOOP]] ]
+; CHECK-NEXT:    [[GEP0:%.*]] = getelementptr i8, ptr [[PTR]], i64 0
+; CHECK-NEXT:    store <2 x i8> [[PACK1]], ptr [[GEP0]], align 1
+; CHECK-NEXT:    br label %[[LOOP]]
+; CHECK:       [[EXIT:.*:]]
+; CHECK-NEXT:    ret void
+;
+entry:
+  %add0 = add i8 %val, 0
+  %mul1 = mul i8 %val, 1
+  br label %loop
+
+loop:
+  %phi0 = phi i8 [0, %entry], [1, %loop]
+  %phi1 = phi i8 [0, %entry], [1, %loop]
+  %gep0 = getelementptr i8, ptr %ptr, i64 0
+  %gep1 = getelementptr i8, ptr %ptr, i64 1
+  store i8 %add0, ptr %gep0
+  store i8 %mul1, ptr %gep1
+  br label %loop
+
+exit:
+  ret void
+}
diff --git a/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp b/llvm/unittests/Transforms/Vectorize/SandboxVectorizer/VecUtilsTest.cpp
@@ -481,6 +481,36 @@ define void @foo(i8 %v) {
   EXPECT_EQ(sandboxir::VecUtils::getLowest(DiffBBs), nullptr);
 }
 
+TEST_F(VecUtilsTest, GetLastPHIOrSelf) {
+  parseIR(R"IR(
+define void @foo(i8 %v) {
+entry:
+  br label %bb1
+
+bb1:
+  %phi1 = phi i8 [0, %entry], [1, %bb1]
+  %phi2 = phi i8 [0, %entry], [1, %bb1]
+  br label %bb1
+
+bb2:
+  ret void
+}
+)IR");
+  Function &LLVMF = *M->getFunction("foo");
+
+  sandboxir::Context Ctx(C);
+  auto &F = *Ctx.createFunction(&LLVMF);
+  auto &BB = getBasicBlockByName(F, "bb1");
+  auto It = BB.begin();
+  auto *PHI1 = cast<sandboxir::PHINode>(&*It++);
+  auto *PHI2 = cast<sandboxir::PHINode>(&*It++);
+  auto *Br = cast<sandboxir::BranchInst>(&*It++);
+  EXPECT_EQ(sandboxir::VecUtils::getLastPHIOrSelf(PHI1), PHI2);
+  EXPECT_EQ(sandboxir::VecUtils::getLastPHIOrSelf(PHI2), PHI2);
+  EXPECT_EQ(sandboxir::VecUtils::getLastPHIOrSelf(Br), Br);
+  EXPECT_EQ(sandboxir::VecUtils::getLastPHIOrSelf(nullptr), nullptr);
+}
+
 TEST_F(VecUtilsTest, GetCommonScalarType) {
   parseIR(R"IR(
 define void @foo(i8 %v, ptr %ptr) {
