[SDAG] SetCC: remove spurious extensions

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -13924,6 +13924,34 @@ SDValue DAGCombiner::visitSETCC(SDNode *N) {
       }
     }
   }
+
+  // (setcc (zext a), (zext b), setu??) -> (setcc a, b, setu??)
+  // (setcc (sext a), (sext b), sets??) -> (setcc a, b, sets??)
+  if ((ISD::isUnsignedIntSetCC(Cond) && N0.getOpcode() == ISD::ZERO_EXTEND &&
+       N1.getOpcode() == ISD::ZERO_EXTEND) ||
+      (ISD::isSignedIntSetCC(Cond) && N0.getOpcode() == ISD::SIGN_EXTEND &&
+       N1.getOpcode() == ISD::SIGN_EXTEND)) {
+    SDValue LHS = N0.getOperand(0), RHS = N1.getOperand(0);
+    EVT SmallVT =
+        LHS.getScalarValueSizeInBits() > RHS.getScalarValueSizeInBits()
+            ? LHS.getValueType()
+            : RHS.getValueType();
+    if (!LegalOperations ||
+        (SmallVT.isSimple() &&
+         TLI.isCondCodeLegal(Cond, SmallVT.getSimpleVT()))) {
+      LHS = DAG.getExtOrTrunc(ISD::isSignedIntSetCC(Cond), LHS, SDLoc(LHS),
+                              SmallVT);
+      RHS = DAG.getExtOrTrunc(ISD::isSignedIntSetCC(Cond), RHS, SDLoc(RHS),
+                              SmallVT);
+      SDValue NewSetCC =
+          DAG.getSetCC(DL, getSetCCResultType(SmallVT), LHS, RHS, Cond);
+      // Promote to a legal type for setcc, then adjust back to VT (if before
+      // LegalOperations)
+      return DAG.getZExtOrTrunc(
+          TLI.promoteTargetBoolean(DAG, NewSetCC, N0.getValueType()), DL, VT);
+    }
+  }
+
   return SDValue();
 }
 

llvm/test/CodeGen/AArch64/arm64-vcmp.ll

@@ -234,3 +234,29 @@ define <1 x i64> @cmnez_d(<1 x i64> %A) nounwind {
   %mask = sext <1 x i1> %tst to <1 x i64>
   ret <1 x i64> %mask
 }
+
+; Check for the elimination of spurious type extensions
+define <16 x i1> @abdu_cmp(<16 x i8> %a, <16 x i8> %b, <16 x i8> %g) {
+; CHECK-LABEL: abdu_cmp:
+; CHECK:         uabd.16b v0, v0, v1
+; CHECK-NEXT:    cmhi.16b v0, v2, v0
+; CHECK-NEXT:    ret
+  %za = zext <16 x i8> %a to <16 x i32>
+  %zb = zext <16 x i8> %b to <16 x i32>
+  %zg = zext <16 x i8> %g to <16 x i32>
+  %mx = call <16 x i32> @llvm.umax.v16i32(<16 x i32> %za, <16 x i32> %zb)
+  %mn = call <16 x i32> @llvm.umin.v16i32(<16 x i32> %za, <16 x i32> %zb)
+  %abdu = sub <16 x i32> %mx, %mn
+  %cond = icmp ult <16 x i32> %abdu, %zg
+  ret <16 x i1> %cond
+}
+
+define <16 x i1> @sext_cmp(<16 x i8> %a, <16 x i8> %b) {
+; CHECK-LABEL: sext_cmp:
+; CHECK:         cmgt.16b v0, v0, v1
+; CHECK-NEXT:    ret
+  %za = sext <16 x i8> %a to <16 x i32>
+  %zb = sext <16 x i8> %b to <16 x i32>
+  %cond = icmp slt <16 x i32> %za, %zb
+  ret <16 x i1> %cond
+}