[llvm] [LICM] Use OverflowTracking to preserve NUW/NSW when reassociating. (PR #140404)

[Florian Hahn via llvm-commits]

https://github.com/fhahn updated https://github.com/llvm/llvm-project/pull/140404

>From f6b033254c45d1724a85f8ac4c5f8666be4e4aab Mon Sep 17 00:00:00 2001
From: Florian Hahn <flo at fhahn.com>
Date: Sat, 17 May 2025 18:19:16 +0100
Subject: [PATCH] [LICM] Use OverflowTracking to preserve NUW/NSW when
 reassociating.

This enables preserving NSW when both adds have NSW and NUW.

For now, set AllKnownNonNegative/AllKnownNonZero to false when using in
LICM.

https://alive2.llvm.org/ce/z/uu79Xc

Depends on https://github.com/llvm/llvm-project/pull/140403 (included in
PR)
---
 llvm/lib/Transforms/Scalar/LICM.cpp      | 19 +++++++++++--------
 llvm/test/Transforms/LICM/hoist-binop.ll |  4 ++--
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/LICM.cpp b/llvm/lib/Transforms/Scalar/LICM.cpp
index 7d89a13fa3bab..7c7e0dcff0886 100644
--- a/llvm/lib/Transforms/Scalar/LICM.cpp
+++ b/llvm/lib/Transforms/Scalar/LICM.cpp
@@ -2864,14 +2864,7 @@ static bool hoistBOAssociation(Instruction &I, Loop &L,
   auto *NewBO = BinaryOperator::Create(
       Opcode, LV, Inv, BO->getName() + ".reass", BO->getIterator());
 
-  // Copy NUW for ADDs if both instructions have it.
-  if (Opcode == Instruction::Add && BO->hasNoUnsignedWrap() &&
-      BO0->hasNoUnsignedWrap()) {
-    // If `Inv` was not constant-folded, a new Instruction has been created.
-    if (auto *I = dyn_cast<Instruction>(Inv))
-      I->setHasNoUnsignedWrap(true);
-    NewBO->setHasNoUnsignedWrap(true);
-  } else if (Opcode == Instruction::FAdd || Opcode == Instruction::FMul) {
+  if (Opcode == Instruction::FAdd || Opcode == Instruction::FMul) {
     // Intersect FMF flags for FADD and FMUL.
     FastMathFlags Intersect = BO->getFastMathFlags() & BO0->getFastMathFlags();
     if (auto *I = dyn_cast<Instruction>(Inv))
@@ -2884,6 +2877,16 @@ static bool hoistBOAssociation(Instruction &I, Loop &L,
     if (auto *I = dyn_cast<PossiblyDisjointInst>(Inv))
       I->setIsDisjoint(Disjoint);
     cast<PossiblyDisjointInst>(NewBO)->setIsDisjoint(Disjoint);
+  } else {
+    OverflowTracking Flags;
+    Flags.AllKnownNonNegative = false;
+    Flags.AllKnownNonZero = false;
+    Flags.mergeFlags(*BO);
+    Flags.mergeFlags(*BO0);
+    // If `Inv` was not constant-folded, a new Instruction has been created.
+    if (auto *I = dyn_cast<Instruction>(Inv))
+      Flags.applyFlags(*I);
+    Flags.applyFlags(*NewBO);
   }
 
   BO->replaceAllUsesWith(NewBO);
diff --git a/llvm/test/Transforms/LICM/hoist-binop.ll b/llvm/test/Transforms/LICM/hoist-binop.ll
index 33161090a8ccf..1b1347776fb9e 100644
--- a/llvm/test/Transforms/LICM/hoist-binop.ll
+++ b/llvm/test/Transforms/LICM/hoist-binop.ll
@@ -371,13 +371,13 @@ loop:
 define void @add_nuw_nsw(i64 %c1, i64 %c2) {
 ; CHECK-LABEL: @add_nuw_nsw(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = add nuw i64 [[C1:%.*]], [[C2:%.*]]
+; CHECK-NEXT:    [[INVARIANT_OP:%.*]] = add nuw nsw i64 [[C1:%.*]], [[C2:%.*]]
 ; CHECK-NEXT:    br label [[LOOP:%.*]]
 ; CHECK:       loop:
 ; CHECK-NEXT:    [[INDEX:%.*]] = phi i64 [ 0, [[ENTRY:%.*]] ], [ [[INDEX_NEXT_REASS:%.*]], [[LOOP]] ]
 ; CHECK-NEXT:    [[STEP_ADD:%.*]] = add nuw nsw i64 [[INDEX]], [[C1]]
 ; CHECK-NEXT:    call void @use(i64 [[STEP_ADD]])
-; CHECK-NEXT:    [[INDEX_NEXT_REASS]] = add nuw i64 [[INDEX]], [[INVARIANT_OP]]
+; CHECK-NEXT:    [[INDEX_NEXT_REASS]] = add nuw nsw i64 [[INDEX]], [[INVARIANT_OP]]
 ; CHECK-NEXT:    br label [[LOOP]]
 ;
 entry: