From 53efc2461fded2f7999ed6c82dfcfd1a55bcaf19 Mon Sep 17 00:00:00 2001
From: Vladimir Radosavljevic <vr@matterlabs.dev>
Date: Tue, 29 Oct 2024 10:23:46 +0100
Subject: [PATCH] [TailDuplication] Add option to enable tail duplication of
 fallthrough BBs

During TailDuplication pass, fallthrough BBs are not duplicated.
This is fine for most cases, but there are cases where we want
to duplicate fallthrough BBs as well (e.g. EvmEmulator). This
patch adds the `tail-dup-fallthrough-bbs` option to enable
tail duplication of fallthrough BBs.

PR: #727.

Signed-off-by: Vladimir Radosavljevic <vr@matterlabs.dev>
---
 llvm/lib/CodeGen/TailDuplicator.cpp                  | 11 ++++++++++-
 llvm/test/CodeGen/EraVM/tail-dup-fallthrough-bbs.mir | 11 +++--------
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/CodeGen/TailDuplicator.cpp b/llvm/lib/CodeGen/TailDuplicator.cpp
index 797e417339a0..350412ed32c8 100644
--- a/llvm/lib/CodeGen/TailDuplicator.cpp
+++ b/llvm/lib/CodeGen/TailDuplicator.cpp
@@ -76,6 +76,13 @@ static cl::opt<bool>
 static cl::opt<unsigned> TailDupLimit("tail-dup-limit", cl::init(~0U),
                                       cl::Hidden);
 
+// EraVM local begin
+static cl::opt<bool>
+    TailDupFallthroughBBs("tail-dup-fallthrough-bbs",
+                          cl::desc("Tail duplicate fallthrough basic blocks"),
+                          cl::init(false), cl::Hidden);
+// EraVM local end
+
 void TailDuplicator::initMF(MachineFunction &MFin, bool PreRegAlloc,
                             const MachineBranchProbabilityInfo *MBPIin,
                             MBFIWrapper *MBFIin,
@@ -559,8 +566,10 @@ bool TailDuplicator::shouldTailDuplicate(bool IsSimple,
   // When doing tail-duplication during layout, the block ordering is in flux,
   // so canFallThrough returns a result based on incorrect information and
   // should just be ignored.
-  if (!LayoutMode && TailBB.canFallThrough())
+  // EraVM local begin
+  if (!TailDupFallthroughBBs && !LayoutMode && TailBB.canFallThrough())
     return false;
+  // EraVM local end
 
   // Don't try to tail-duplicate single-block loops.
   if (TailBB.isSuccessor(&TailBB))
diff --git a/llvm/test/CodeGen/EraVM/tail-dup-fallthrough-bbs.mir b/llvm/test/CodeGen/EraVM/tail-dup-fallthrough-bbs.mir
index 68c3d32bfca6..8c35957e34a7 100644
--- a/llvm/test/CodeGen/EraVM/tail-dup-fallthrough-bbs.mir
+++ b/llvm/test/CodeGen/EraVM/tail-dup-fallthrough-bbs.mir
@@ -1,5 +1,5 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 2
-# RUN: llc -x mir -run-pass tailduplication -verify-machineinstrs -compile-twice=false < %s | FileCheck %s
+# RUN: llc -x mir -run-pass tailduplication -tail-dup-fallthrough-bbs -verify-machineinstrs -compile-twice=false < %s | FileCheck %s
 
 --- |
 
@@ -15,17 +15,12 @@ tracksRegLiveness: true
 body:             |
   ; CHECK-LABEL: name: test
   ; CHECK: bb.0:
-  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
-  ; CHECK-NEXT:   liveins: $r1
-  ; CHECK-NEXT: {{  $}}
-  ; CHECK-NEXT:   J %bb.1
-  ; CHECK-NEXT: {{  $}}
-  ; CHECK-NEXT: bb.1:
   ; CHECK-NEXT:   successors: %bb.3(0x40000000), %bb.2(0x40000000)
   ; CHECK-NEXT:   liveins: $r1
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT:   dead $r0 = SUBxrr_v i256 10, $r1, i256 0, implicit-def $flags
-  ; CHECK-NEXT:   JCl %bb.3, i256 3, implicit killed $flags
+  ; CHECK-NEXT:   JCl %bb.3, 3, implicit $flags
+  ; CHECK-NEXT:   J %bb.2
   ; CHECK-NEXT: {{  $}}
   ; CHECK-NEXT: bb.2:
   ; CHECK-NEXT:   liveins: $r1