[llvm] [X86] Respect blockaddress offsets when performing X86 LEA fixups (PR #71641)

[Maurice Heumann via llvm-commits]

https://github.com/momo5502 created https://github.com/llvm/llvm-project/pull/71641

The X86FixupLEAs pass drops blockaddress offsets, when splitting up slow 3-ops LEAs, as can be seen in this example:

https://godbolt.org/z/bEsc3Poje

Before running the pass, the first instruction in bb.0 is a LEA with ebp, ebx and a blockaddress.
After the transformation, the blockaddress is missing.

The reason this happens is because the 3-ops LEA is being splitup into a 2-ops LEA + an add instruction.
However, as hasLEAOffset does not take blockaddresses into consideration, the add is not emitted and thus leading to the offset being dropped.

Taking blockaddresses into consideration fixes this issue and results in the add instruction being emitted.

>From 45db92ab9442eac29b37dd0a3acae9a827ee208a Mon Sep 17 00:00:00 2001
From: Maurice Heumann <maurice.heumann at wibu.com>
Date: Wed, 8 Nov 2023 09:42:34 +0100
Subject: [PATCH] [X86] Respect blockaddress offsets when performing X86 LEA
 fixups

The X86FixupLEAs pass drops blockaddress offsets, when splitting up
slow 3-ops LEAs, as can be seen in this example:

https://godbolt.org/z/bEsc3Poje

Before running the pass, the first instruction in bb.0 is a
LEA with ebp, ebx and a blockaddress.
After the transformation, the blockaddress is missing.

The reason this happens is because the 3-ops LEA is being splitup into
a 2-ops LEA + an add instruction.
However, as hasLEAOffset does not take blockaddresses into consideration,
the add is not emitted and thus leading to the offset being dropped.

Taking blockaddresses into consideration fixes this issue and results
in the add instruction being emitted.
---
 llvm/lib/Target/X86/X86FixupLEAs.cpp          |  3 +-
 .../CodeGen/X86/lea-fixup-blockaddress.mir    | 30 +++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/X86/lea-fixup-blockaddress.mir

diff --git a/llvm/lib/Target/X86/X86FixupLEAs.cpp b/llvm/lib/Target/X86/X86FixupLEAs.cpp
index ceafbf177352a87..beeebf42dfe81ab 100644
--- a/llvm/lib/Target/X86/X86FixupLEAs.cpp
+++ b/llvm/lib/Target/X86/X86FixupLEAs.cpp
@@ -341,7 +341,8 @@ static inline bool hasInefficientLEABaseReg(const MachineOperand &Base,
 }
 
 static inline bool hasLEAOffset(const MachineOperand &Offset) {
-  return (Offset.isImm() && Offset.getImm() != 0) || Offset.isGlobal();
+  return (Offset.isImm() && Offset.getImm() != 0) || Offset.isGlobal() ||
+         Offset.isBlockAddress();
 }
 
 static inline unsigned getADDrrFromLEA(unsigned LEAOpcode) {
diff --git a/llvm/test/CodeGen/X86/lea-fixup-blockaddress.mir b/llvm/test/CodeGen/X86/lea-fixup-blockaddress.mir
new file mode 100644
index 000000000000000..14134e268613bad
--- /dev/null
+++ b/llvm/test/CodeGen/X86/lea-fixup-blockaddress.mir
@@ -0,0 +1,30 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 3
+# RUN: llc -mtriple=i386-unknown-linux-gnu -mattr=slow-3ops-lea -run-pass x86-fixup-LEAs -o -  %s | FileCheck %s
+
+--- |
+  define i32 @square(i32 %0) local_unnamed_addr {
+    %blub = getelementptr i8, ptr blockaddress(@square, %2), i32 %0
+    indirectbr ptr %blub, [label %2]
+
+  2:
+    ret i32 0
+  }
+
+---
+name:            square
+body:             |
+  ; CHECK-LABEL: name: square
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   renamable $eax = LEA32r renamable $ebx, 1, renamable $ebp, 0, $noreg
+  ; CHECK-NEXT:   $eax = ADD32ri $eax, target-flags(x86-gotoff) blockaddress(@square, %ir-block.1), implicit-def $eflags
+  ; CHECK-NEXT:   JMP32r killed renamable $eax
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1 (%ir-block.1, ir-block-address-taken %ir-block.1):
+  ; CHECK-NEXT:   RET 0
+  bb.0:
+    renamable $eax = LEA32r renamable $ebp, 1, renamable $ebx, target-flags(x86-gotoff) blockaddress(@square, %ir-block.1), $noreg
+    JMP32r killed renamable $eax
+
+  bb.1 (%ir-block.1, ir-block-address-taken %ir-block.1):
+    RET 0
+...