[llvm] [Codegen, BasicBlockSections] Avoid cloning blocks which have their machine block address taken. (PR #94296)

[Rahman Lavaee via llvm-commits]

https://github.com/rlavaee created https://github.com/llvm/llvm-project/pull/94296

These blocks usually show up in the form of branches within inline assembly. Since it's hard to rewire them, we fully omit paths with such blocks from path cloning.

>From 3aa2bf51fd348fecee3b2323c55481cfd9a799ed Mon Sep 17 00:00:00 2001
From: Rahman Lavaee <rahmanl at google.com>
Date: Tue, 4 Jun 2024 00:05:38 +0000
Subject: [PATCH] [Codegen, BasicBlockSections] Avoid cloning blocks which have
 their machine block address taken.

These blocks usually show up in the form of branches within inline
assembly. Since it's hard to rewire them, we fully omit paths with such
blocks from path cloning.
---
 llvm/lib/CodeGen/BasicBlockPathCloning.cpp    | 10 +++++++
 .../basic-block-sections-cloning-invalid.ll   | 27 ++++++++++++++-----
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
index 901542e8507bd..19f824850607c 100644
--- a/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
+++ b/llvm/lib/CodeGen/BasicBlockPathCloning.cpp
@@ -119,6 +119,16 @@ bool IsValidCloning(const MachineFunction &MF,
           return false;
         }
       }
+      if (PathBB->isMachineBlockAddressTaken()) {
+        // Avoid cloning blocks which have their address taken since we can't
+        // rewire branches to those blocks as easily (e.g., branches within
+        // inline assembly).
+        WithColor::warning()
+            << "block #" << BBID
+            << " has its machine block address taken in function "
+            << MF.getName() << "\n";
+        return false;
+      }
     }
 
     if (I != ClonePath.size() - 1 && !PathBB->empty() &&
diff --git a/llvm/test/CodeGen/X86/basic-block-sections-cloning-invalid.ll b/llvm/test/CodeGen/X86/basic-block-sections-cloning-invalid.ll
index 521ec43ef050a..c316ef9f8f260 100644
--- a/llvm/test/CodeGen/X86/basic-block-sections-cloning-invalid.ll
+++ b/llvm/test/CodeGen/X86/basic-block-sections-cloning-invalid.ll
@@ -13,8 +13,8 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'v1' > %t2
 ; RUN: echo 'f foo' >> %t2
 ; RUN: echo 'p 0 2 3' >> %t2
-; RUN: echo 'p 0 1 3' >> %t2
-; RUN: echo 'c 0 1.1 3.2 2.1 3.1 1' >> %t2
+; RUN: echo 'p 0 1 2 3' >> %t2
+; RUN: echo 'c 0 1.1 2.2 3.2 2.1 3.1 1' >> %t2
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t2 2> %t2.err | FileCheck %s --check-prefixes=PATH
 ; RUN: FileCheck %s --check-prefixes=WARN1 < %t2.err
 ; RUN: echo 'v1' > %t3
@@ -23,6 +23,14 @@ declare void @effect(i32 zeroext)
 ; RUN: echo 'c 0 100.1 1' >> %t3
 ; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t3 2> %t3.err | FileCheck %s
 ; RUN: FileCheck %s --check-prefixes=WARN2 < %t3.err
+; RUN: echo 'v1' > %t4
+; RUN: echo 'f foo' >> %t4
+; RUN: echo 'p 1 6' >> %t4
+; RUN: echo 'c 0 1 6.1' >> %t4
+; RUN: llc < %s -mtriple=x86_64-pc-linux -O0 -function-sections -basic-block-sections=%t4 2> %t4.err | FileCheck %s
+; RUN: FileCheck %s --check-prefixes=WARN3 < %t4.err
+
+
 
 define void @foo(i1 %a, i1 %b, i1 %c, i1 %d) {
 b0:
@@ -31,23 +39,29 @@ b0:
 
 b1:                                           ; preds = %b0
   call void @effect(i32 1)
-  br i1 %b, label %b2, label %b3
+  br i1 %b, label %b2, label %b6
 
 b2:                                             ; preds = %b1
   call void @effect(i32 2)
   br label %b3
 
-b3:                                            ; preds = %b0, %b1, %b2
+b3:                                            ; preds = %b0, %b2
   call void @effect(i32 3)
   br i1 %c, label %b4, label %b5
 
 b4:                                             ; preds = %b3
   call void @effect(i32 4)
-  br i1 %d, label %b5, label %cold
+  callbr void asm sideeffect "je ${0:l}", "!i,~{dirflag},~{fpsr},~{flags}"()
+    to label %b5 [label %b6]
 
 b5:                                            ; preds = %b3, %b4
   call void @effect(i32 5)
   ret void
+
+b6:                                            ; preds = %b1, %b4
+  call void @effect(i32 6)
+  ret void
+
 cold:
   call void @effect(i32 6)                     ; preds = %b4
   ret void
@@ -59,7 +73,7 @@ cold:
 
 ; CHECK:   je .LBB0_3
 ; PATH:  # %bb.7:      # %b1
-; PATH:  # %bb.8:      # %b3
+; PATH:  # %bb.8:      # %b2
 ; PATH:    jne .LBB0_4
 ; CHECK: # %bb.1:      # %b1
 ; CHECK:   jne foo.cold
@@ -69,4 +83,5 @@ cold:
 ;; Check the warnings
 ; WARN1: warning: block #2 is not a successor of block #0 in function foo
 ; WARN2: warning: no block with id 100 in function foo
+; WARN3: warning: block #6 has its machine block address taken in function foo