[llvm] [SimplifyCFG] Remove limitation on sinking of load/store of alloca (PR #104788)

[Nikita Popov via llvm-commits]

https://github.com/nikic created https://github.com/llvm/llvm-project/pull/104788

This is a followup to https://github.com/llvm/llvm-project/pull/104579 to remove the limitation on sinking loads/stores of allocas entirely, even if this would introduce a phi node.

Nowadays, SROA supports speculating load/store over select/phi. Additionally, SimplifyCFG with sinking only runs at the end of the function simplification pipeline, after SROA. I checked that the two tests modified here still successfully SROA after the SimplifyCFG transform.

We should, however, keep the limitation on lifetime intrinsics. SROA does not have speculation support for these, and I've also found that the way these are handled in the backend is very problematic (https://github.com/llvm/llvm-project/issues/104776), so I think we should leave them alone.

>From cf1ff60895f7d3b1705acf429c1b8d7167ae5312 Mon Sep 17 00:00:00 2001
From: Nikita Popov <npopov at redhat.com>
Date: Mon, 19 Aug 2024 15:13:36 +0200
Subject: [PATCH] [SimplifyCFG] Remove limitation on sinking of load/store of
 alloca

This is a followup to https://github.com/llvm/llvm-project/pull/104579
to remove the limitation on sinking loads/stores of allocas entirely,
even if this would introduce a phi node.

Nowdays, SROA support speculating load/store over select/phi.
Additionally, SimplifyCFG with sinking only runs at the end of
the function simplification pipeline, after SROA. I checked that
the two tests modified here still successfully SROA after the
SimplifyCFG transform.

We should, however, keep the limitation on lifetime intrinsics.
SROA does not have speculation support for these, and I've also
found that the way these are handled in the backend is very
problematic (https://github.com/llvm/llvm-project/issues/104776),
so I think we should leave them alone.
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 19 ++++-------------
 .../SimplifyCFG/X86/sink-common-code.ll       | 21 ++++++-------------
 2 files changed, 10 insertions(+), 30 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index ebdf760bda7f1a..3f680de98470ca 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -2031,21 +2031,10 @@ static bool canSinkInstructions(
       return I->getOperand(OI) == I0->getOperand(OI);
     };
     if (!all_of(Insts, SameAsI0)) {
-      // Because SROA historically couldn't handle speculating stores of
-      // selects, we try not to sink loads, stores or lifetime markers of
-      // allocas when we'd have to create a PHI for the address operand.
-      // TODO: SROA supports speculation for loads and stores now -- remove
-      // this hack?
-      if (isa<StoreInst>(I0) && OI == 1 &&
-          any_of(Insts, [](const Instruction *I) {
-            return isa<AllocaInst>(I->getOperand(1)->stripPointerCasts());
-          }))
-        return false;
-      if (isa<LoadInst>(I0) && OI == 0 &&
-          any_of(Insts, [](const Instruction *I) {
-            return isa<AllocaInst>(I->getOperand(0)->stripPointerCasts());
-          }))
-        return false;
+      // SROA can't speculate lifetime markers of selects/phis, and the
+      // backend may handle such lifetimes incorrectly as well (#104776).
+      // Don't sink lifetimes if it would introduce a phi on the pointer
+      // argument.
       if (isLifeTimeMarker(I0) && OI == 1 &&
           any_of(Insts, [](const Instruction *I) {
             return isa<AllocaInst>(I->getOperand(1)->stripPointerCasts());
diff --git a/llvm/test/Transforms/SimplifyCFG/X86/sink-common-code.ll b/llvm/test/Transforms/SimplifyCFG/X86/sink-common-code.ll
index cd26d949836e68..39b1bec164b9e5 100644
--- a/llvm/test/Transforms/SimplifyCFG/X86/sink-common-code.ll
+++ b/llvm/test/Transforms/SimplifyCFG/X86/sink-common-code.ll
@@ -801,14 +801,8 @@ define i32 @test_pr30188(i1 zeroext %flag, i32 %x) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[Y:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    [[Z:%.*]] = alloca i32, align 4
-; CHECK-NEXT:    br i1 [[FLAG:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
-; CHECK:       if.then:
-; CHECK-NEXT:    store i32 [[X:%.*]], ptr [[Y]], align 4
-; CHECK-NEXT:    br label [[IF_END:%.*]]
-; CHECK:       if.else:
-; CHECK-NEXT:    store i32 [[X]], ptr [[Z]], align 4
-; CHECK-NEXT:    br label [[IF_END]]
-; CHECK:       if.end:
+; CHECK-NEXT:    [[Y_Z:%.*]] = select i1 [[FLAG:%.*]], ptr [[Y]], ptr [[Z]]
+; CHECK-NEXT:    store i32 [[X:%.*]], ptr [[Y_Z]], align 4
 ; CHECK-NEXT:    ret i32 1
 ;
 entry:
@@ -834,17 +828,14 @@ define i32 @test_pr30188a(i1 zeroext %flag, i32 %x) {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[Y:%.*]] = alloca i32, align 4
 ; CHECK-NEXT:    [[Z:%.*]] = alloca i32, align 4
-; CHECK-NEXT:    br i1 [[FLAG:%.*]], label [[IF_THEN:%.*]], label [[IF_ELSE:%.*]]
+; CHECK-NEXT:    br i1 [[FLAG:%.*]], label [[IF_THEN:%.*]], label [[IF_END:%.*]]
 ; CHECK:       if.then:
 ; CHECK-NEXT:    call void @g()
-; CHECK-NEXT:    [[ONE:%.*]] = load i32, ptr [[Y]], align 4
-; CHECK-NEXT:    br label [[IF_END:%.*]]
-; CHECK:       if.else:
-; CHECK-NEXT:    [[THREE:%.*]] = load i32, ptr [[Z]], align 4
 ; CHECK-NEXT:    br label [[IF_END]]
 ; CHECK:       if.end:
-; CHECK-NEXT:    [[THREE_SINK:%.*]] = phi i32 [ [[THREE]], [[IF_ELSE]] ], [ [[ONE]], [[IF_THEN]] ]
-; CHECK-NEXT:    [[FOUR:%.*]] = add i32 [[THREE_SINK]], 2
+; CHECK-NEXT:    [[Z_SINK:%.*]] = phi ptr [ [[Y]], [[IF_THEN]] ], [ [[Z]], [[ENTRY:%.*]] ]
+; CHECK-NEXT:    [[THREE:%.*]] = load i32, ptr [[Z_SINK]], align 4
+; CHECK-NEXT:    [[FOUR:%.*]] = add i32 [[THREE]], 2
 ; CHECK-NEXT:    store i32 [[FOUR]], ptr [[Y]], align 4
 ; CHECK-NEXT:    ret i32 1
 ;