[llvm] [DAGCombiner][AArch64] Make combineCarryDiamond avoid creating UADDO_CARRY with carry in larger than setcc result type. (PR #89121)

[Craig Topper via llvm-commits]

https://github.com/topperc updated https://github.com/llvm/llvm-project/pull/89121

>From 25ecd0b66973f51f7db3c7680372323c2cfbc10b Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Wed, 17 Apr 2024 11:35:45 -0700
Subject: [PATCH] [DAGCombiner][AArch64] Make combineCarryDiamond avoid
 creating UADDO_CARRY with carry in larger than setcc result type.

In the attach test case we were creating a UADDO_CARRY with i1
carry out and i41 carry in. i41 exceeds is larger than the setcc
result type for AArch64 which is i32. i41 needs to be promoted to
i64 since it is larger than i32. The type legalizer tried to use
promoteTargetBoolean, but that can only promote from a type smaller
than setcc result type.

The easiest fix here is to force the carryin type to match the
carryout type at the type of creation. This should ensure the node
won't exceeed setcc result type as long as the output type doesn't.

I think we should explore requiring the types to match for this node.

Fixes #88966
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp |  2 ++
 llvm/test/CodeGen/AArch64/pr88966.ll          | 29 +++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 llvm/test/CodeGen/AArch64/pr88966.ll

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index c36b1cc9039c26..be724f41992c2e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -3520,6 +3520,8 @@ static SDValue combineCarryDiamond(SelectionDAG &DAG, const TargetLowering &TLI,
     return SDValue();
 
   SDLoc DL(N);
+  CarryIn = DAG.getBoolExtOrTrunc(CarryIn, DL, Carry1->getValueType(1),
+                                  Carry1->getValueType(0));
   SDValue Merged =
       DAG.getNode(NewOp, DL, Carry1->getVTList(), Carry0.getOperand(0),
                   Carry0.getOperand(1), CarryIn);
diff --git a/llvm/test/CodeGen/AArch64/pr88966.ll b/llvm/test/CodeGen/AArch64/pr88966.ll
new file mode 100644
index 00000000000000..0743773d171d72
--- /dev/null
+++ b/llvm/test/CodeGen/AArch64/pr88966.ll
@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc < %s -mtriple=aarch64 | FileCheck %s
+
+define i32 @f(ptr %arg, i41 %arg1, ptr %arg2) {
+; CHECK-LABEL: f:
+; CHECK:       // %bb.0: // %bb
+; CHECK-NEXT:    and w9, w1, #0x1
+; CHECK-NEXT:    mov w10, #1 // =0x1
+; CHECK-NEXT:    mov x8, x0
+; CHECK-NEXT:    cmp w9, #1
+; CHECK-NEXT:    mov w0, wzr
+; CHECK-NEXT:    adc x9, xzr, x10
+; CHECK-NEXT:    str x9, [x2]
+; CHECK-NEXT:    str xzr, [x8]
+; CHECK-NEXT:    ret
+bb:
+  %new0 = and i41 %arg1, 1
+  %last = trunc i41 %new0 to i1
+  %i = add i64 0, 1
+  %i3 = zext i1 %last to i64
+  %i4 = add i64 %i, %i3
+  %i5 = icmp ult i64 %i, 0
+  %i6 = icmp ult i64 %i4, %i
+  %i7 = and i1 %i5, %i6
+  %i8 = zext i1 %i7 to i64
+  store i64 %i4, ptr %arg2, align 8
+  store i64 %i8, ptr %arg, align 8
+  ret i32 0
+}