[llvm] [BranchFolding] Add an option to tail merge only bbs without successors (PR #101969)

[Vladimir Radosavljevic via llvm-commits]

https://github.com/vladimirradosavljevic updated https://github.com/llvm/llvm-project/pull/101969

>From 39acbc2c90e977bbcf488bd2d7518e8042ffe073 Mon Sep 17 00:00:00 2001
From: Vladimir Radosavljevic <vr at matterlabs.dev>
Date: Mon, 5 Aug 2024 14:26:37 +0200
Subject: [PATCH] [BranchFolding] Add an option to tail merge only bbs without
 successors

This patch adds -tail-merge-only-bbs-without-succ option to
BranchFolding pass. The only difference between this option
and -enable-tail-merge is that this option can reduce code
size in case function has multiple return points that can be
merged.
---
 llvm/lib/CodeGen/BranchFolding.cpp            |  8 +++
 ...nch-folding-tail-merge-bbs-without-succ.ll | 64 +++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 llvm/test/CodeGen/RISCV/branch-folding-tail-merge-bbs-without-succ.ll

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 1dc278586f117..41c947c988077 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -85,6 +85,11 @@ TailMergeSize("tail-merge-size",
               cl::desc("Min number of instructions to consider tail merging"),
               cl::init(3), cl::Hidden);
 
+static cl::opt<bool> TailMergeOnlyBBsWithoutSucc(
+    "tail-merge-only-bbs-without-succ",
+    cl::desc("Tail merge only basic blocks without successors"),
+    cl::init(false), cl::Hidden);
+
 namespace {
 
   /// BranchFolderPass - Wrap branch folder in a machine function pass.
@@ -1035,6 +1040,9 @@ bool BranchFolder::TailMergeBlocks(MachineFunction &MF) {
   if (MergePotentials.size() >= 2)
     MadeChange |= TryTailMergeBlocks(nullptr, nullptr, MinCommonTailLength);
 
+  if (TailMergeOnlyBBsWithoutSucc)
+    return MadeChange;
+
   // Look at blocks (IBB) with multiple predecessors (PBB).
   // We change each predecessor to a canonical form, by
   // (1) temporarily removing any unconditional branch from the predecessor
diff --git a/llvm/test/CodeGen/RISCV/branch-folding-tail-merge-bbs-without-succ.ll b/llvm/test/CodeGen/RISCV/branch-folding-tail-merge-bbs-without-succ.ll
new file mode 100644
index 0000000000000..e3d8ffcf0cebb
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/branch-folding-tail-merge-bbs-without-succ.ll
@@ -0,0 +1,64 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -mtriple=riscv32 -tail-merge-only-bbs-without-succ=true < %s \
+; RUN:   | FileCheck %s --check-prefix=TAIL-MERGE-NO-SUCC
+; RUN: llc -mtriple=riscv32 -enable-tail-merge=false < %s \
+; RUN:   | FileCheck %s --check-prefix=NO-TAIL-MERGE
+
+define i32 @test(ptr %addr,i1 %cond1, i1 %cond2) {
+; TAIL-MERGE-NO-SUCC-LABEL: test:
+; TAIL-MERGE-NO-SUCC:       # %bb.0: # %entry
+; TAIL-MERGE-NO-SUCC-NEXT:    andi a1, a1, 1
+; TAIL-MERGE-NO-SUCC-NEXT:    beqz a1, .LBB0_4
+; TAIL-MERGE-NO-SUCC-NEXT:  # %bb.1: # %bb1
+; TAIL-MERGE-NO-SUCC-NEXT:    li a1, 2
+; TAIL-MERGE-NO-SUCC-NEXT:    andi a2, a2, 1
+; TAIL-MERGE-NO-SUCC-NEXT:    beqz a2, .LBB0_3
+; TAIL-MERGE-NO-SUCC-NEXT:  # %bb.2: # %bb2
+; TAIL-MERGE-NO-SUCC-NEXT:    sw a1, 0(a0)
+; TAIL-MERGE-NO-SUCC-NEXT:    li a0, 0
+; TAIL-MERGE-NO-SUCC-NEXT:    ret
+; TAIL-MERGE-NO-SUCC-NEXT:  .LBB0_3: # %bb3
+; TAIL-MERGE-NO-SUCC-NEXT:    sw a1, 0(a0)
+; TAIL-MERGE-NO-SUCC-NEXT:  .LBB0_4: # %bb5
+; TAIL-MERGE-NO-SUCC-NEXT:    li a0, 0
+; TAIL-MERGE-NO-SUCC-NEXT:    ret
+;
+; NO-TAIL-MERGE-LABEL: test:
+; NO-TAIL-MERGE:       # %bb.0: # %entry
+; NO-TAIL-MERGE-NEXT:    andi a1, a1, 1
+; NO-TAIL-MERGE-NEXT:    beqz a1, .LBB0_3
+; NO-TAIL-MERGE-NEXT:  # %bb.1: # %bb1
+; NO-TAIL-MERGE-NEXT:    li a1, 2
+; NO-TAIL-MERGE-NEXT:    andi a2, a2, 1
+; NO-TAIL-MERGE-NEXT:    beqz a2, .LBB0_4
+; NO-TAIL-MERGE-NEXT:  # %bb.2: # %bb2
+; NO-TAIL-MERGE-NEXT:    sw a1, 0(a0)
+; NO-TAIL-MERGE-NEXT:    li a0, 0
+; NO-TAIL-MERGE-NEXT:    ret
+; NO-TAIL-MERGE-NEXT:  .LBB0_3: # %bb5
+; NO-TAIL-MERGE-NEXT:    li a0, 0
+; NO-TAIL-MERGE-NEXT:    ret
+; NO-TAIL-MERGE-NEXT:  .LBB0_4: # %bb3
+; NO-TAIL-MERGE-NEXT:    sw a1, 0(a0)
+; NO-TAIL-MERGE-NEXT:    li a0, 0
+; NO-TAIL-MERGE-NEXT:    ret
+entry:
+  br i1 %cond1, label %bb1, label %bb5
+
+bb1:
+  br i1 %cond2, label %bb2, label %bb3
+
+bb2:
+  store i32 2, ptr %addr, align 4
+  br label %bb4
+
+bb3:
+  store i32 2, ptr %addr, align 4
+  br label %bb4
+
+bb4:
+  ret i32 0
+
+bb5:
+  ret i32 0
+}