[llvm] [AArch64] Drop poison-generating flags in `genSubAdd2SubSub` combiner (PR #90028)

[Antonio Frighetto via llvm-commits]

https://github.com/antoniofrighetto created https://github.com/llvm/llvm-project/pull/90028

A miscompilation issue has been addressed with improved handling.

Fixes: https://github.com/llvm/llvm-project/issues/88950.

>From 5743cbc323c09a1d6ed69fe638c9c782404b89aa Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me at antoniofrighetto.com>
Date: Thu, 25 Apr 2024 09:49:24 +0200
Subject: [PATCH] [AArch64] Drop poison-generating flags in `genSubAdd2SubSub`
 combiner

A miscompilation issue has been addressed with improved handling.

Fixes: https://github.com/llvm/llvm-project/issues/88950.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  | 12 +++++++--
 .../AArch64/machine-combiner-subadd2.mir      | 27 +++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 7bf06e71a03059..8653f621ba545f 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -6924,14 +6924,22 @@ genSubAdd2SubSub(MachineFunction &MF, MachineRegisterInfo &MRI,
     assert((Opcode == AArch64::SUBWrr || Opcode == AArch64::SUBXrr) &&
            "Unexpected instruction opcode.");
 
+  uint32_t RootFlags = Root.getFlags();
+  if (RootFlags & MachineInstr::NoSWrap)
+    Root.clearFlag(MachineInstr::NoSWrap);
+  if (RootFlags & MachineInstr::NoUWrap)
+    Root.clearFlag(MachineInstr::NoUWrap);
+
   MachineInstrBuilder MIB1 =
       BuildMI(MF, MIMetadata(Root), TII->get(Opcode), NewVR)
           .addReg(RegA, getKillRegState(RegAIsKill))
-          .addReg(RegB, getKillRegState(RegBIsKill));
+          .addReg(RegB, getKillRegState(RegBIsKill))
+          .setMIFlags(Root.getFlags());
   MachineInstrBuilder MIB2 =
       BuildMI(MF, MIMetadata(Root), TII->get(Opcode), ResultReg)
           .addReg(NewVR, getKillRegState(true))
-          .addReg(RegC, getKillRegState(RegCIsKill));
+          .addReg(RegC, getKillRegState(RegCIsKill))
+          .setMIFlags(Root.getFlags());
 
   InstrIdxForVirtReg.insert(std::make_pair(NewVR, 0));
   InsInstrs.push_back(MIB1);
diff --git a/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir b/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
index d1770bb25fae49..0b09e8a4b5cd38 100644
--- a/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
+++ b/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
@@ -237,3 +237,30 @@ body:              |
     RET_ReallyLR implicit $w0
 
 ...
+---
+# Drop nowrap flags in SUB
+
+# CHECK-LABEL: name: test8
+# CHECK:       %7:gpr64 = SUBXrr %1, %0
+# CHECK-NEXT:  %4:gpr64common = SUBXrr killed %7, killed %2
+
+name:            test8
+registers:
+  - { id: 0, class: gpr64 }
+  - { id: 1, class: gpr64 }
+  - { id: 2, class: gpr64common }
+  - { id: 3, class: gpr64 }
+  - { id: 4, class: gpr64common }
+  - { id: 5, class: gpr64 }
+body:              |
+  bb.0:
+    %1:gpr64 = COPY $x1
+    %0:gpr64 = COPY $x0
+    %2:gpr64common = ORRXri %0:gpr64, 4096
+    %3:gpr64 = ADDXrr killed %2:gpr64common, %0:gpr64
+    %4:gpr64common = nsw SUBSXrr %1:gpr64, killed %3:gpr64, implicit-def dead $nzcv
+    %5:gpr64 = SUBSXri %4:gpr64common, 0, 0, implicit-def $nzcv
+    $x0 = COPY %5:gpr64
+    RET_ReallyLR implicit $x0
+
+...