From 88452508f35f0140be1dbb03516929ed48342236 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Amaury=20S=C3=A9chet?= <deadalnix@gmail.com>
Date: Wed, 5 Jul 2023 17:45:01 +0000
Subject: [PATCH] [DAG] Improve carry reconstruction in combineCarryDiamond.

The gain is usually suffiscient to go the extra mile and reconstruct a carry in some cases.

Reviewed By: RKSimon

Differential Revision: https://reviews.llvm.org/D154533
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp | 16 ++++++++++------
 llvm/test/CodeGen/X86/addcarry.ll             |  9 +++------
 2 files changed, 13 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index 7673ac5..63b7506 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -3017,7 +3017,8 @@ SDValue DAGCombiner::visitADDSAT(SDNode *N) {
   return SDValue();
 }
 
-static SDValue getAsCarry(const TargetLowering &TLI, SDValue V) {
+static SDValue getAsCarry(const TargetLowering &TLI, SDValue V,
+                          bool ForceCarryReconstruction = false) {
   bool Masked = false;
 
   // First, peel away TRUNCATE/ZERO_EXTEND/AND nodes due to legalization.
@@ -3028,11 +3029,17 @@ static SDValue getAsCarry(const TargetLowering &TLI, SDValue V) {
     }
 
     if (V.getOpcode() == ISD::AND && isOneConstant(V.getOperand(1))) {
+      if (ForceCarryReconstruction)
+        return V;
+
       Masked = true;
       V = V.getOperand(0);
       continue;
     }
 
+    if (ForceCarryReconstruction && V.getValueType() == MVT::i1)
+      return V;
+
     break;
   }
 
@@ -3542,11 +3549,8 @@ static SDValue combineCarryDiamond(SelectionDAG &DAG, const TargetLowering &TLI,
     return SDValue();
 
   // Verify that the carry/borrow in is plausibly a carry/borrow bit.
-  // TODO: make getAsCarry() aware of how partial carries are merged.
-  if (CarryIn.getOpcode() != ISD::ZERO_EXTEND)
-    return SDValue();
-  CarryIn = CarryIn.getOperand(0);
-  if (CarryIn.getValueType() != MVT::i1)
+  CarryIn = getAsCarry(TLI, CarryIn, true);
+  if (!CarryIn)
     return SDValue();
 
   SDLoc DL(N);
diff --git a/llvm/test/CodeGen/X86/addcarry.ll b/llvm/test/CodeGen/X86/addcarry.ll
index 6b9883c..af8f921 100644
--- a/llvm/test/CodeGen/X86/addcarry.ll
+++ b/llvm/test/CodeGen/X86/addcarry.ll
@@ -708,13 +708,10 @@ define { i64, i1 } @addcarry_carry_not_i1(i64 %a, i64 %b, i8 %carryin) nounwind
 define { i64, i1 } @addcarry_carry_and_1(i64 %a, i64 %b, i64 %carryin) nounwind {
 ; CHECK-LABEL: addcarry_carry_and_1:
 ; CHECK:       # %bb.0:
-; CHECK-NEXT:    movq %rdx, %rax
-; CHECK-NEXT:    addq %rsi, %rdi
-; CHECK-NEXT:    setb %cl
-; CHECK-NEXT:    andl $1, %eax
-; CHECK-NEXT:    addq %rdi, %rax
+; CHECK-NEXT:    movq %rdi, %rax
+; CHECK-NEXT:    btl $0, %edx
+; CHECK-NEXT:    adcq %rsi, %rax
 ; CHECK-NEXT:    setb %dl
-; CHECK-NEXT:    orb %cl, %dl
 ; CHECK-NEXT:    retq
   %t1 = call { i64, i1 } @llvm.uadd.with.overflow.i64(i64 %a, i64 %b)
   %partial = extractvalue { i64, i1 } %t1, 0
-- 
2.7.4