[llvm] [DAG] Canonicalize zero_extend to sign_extend based on target preference (PR #70671)

[Philip Reames via llvm-commits]

https://github.com/preames created https://github.com/llvm/llvm-project/pull/70671

We already had code to stop the canonicalization in the other direction, let's add the inverse combine.  It turned out we'd missed a least one case which unconditionally created zero_extend, so fix that too to avoid a combine loop.

Note that this will most likely be entirely reworked in the near future based on the new nneg flag.  This change was triggered by me investigating what that would look like, and asking why the current structure was incomplete.

>From d19c0c56fb97ed9c248e9f6182419a078c1f5761 Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Mon, 30 Oct 2023 08:30:45 -0700
Subject: [PATCH] [DAG] Canonicalize zero_extend to sign_extend based on target
 preference

We already had code to stop the canonicalization in the other direction, let's
add the inverse combine.  It turned out we'd missed a least one case which
unconditionally created zero_extend, so fix that too to avoid a combine loop.

Note that this will most likely be entirely reworked in the near future
based on the new nneg flag.  This change was triggered by me investigating
what that would look like, and asking why the current structure was incomplete.
---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp    | 11 ++++++++++-
 llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp |  2 +-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index ca5bd4952866886..3f186a71256d084 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -13481,7 +13481,8 @@ SDValue DAGCombiner::visitSIGN_EXTEND(SDNode *N) {
   if (SDValue V = foldSextSetcc(N))
     return V;
 
-  // fold (sext x) -> (zext x) if the sign bit is known zero.
+  // fold (sext x) -> (zext x) if the sign bit is known zero, and that
+  // is what the target prefers.
   if (!TLI.isSExtCheaperThanZExt(N0.getValueType(), VT) &&
       (!LegalOperations || TLI.isOperationLegal(ISD::ZERO_EXTEND, VT)) &&
       DAG.SignBitIsZero(N0))
@@ -13825,6 +13826,14 @@ SDValue DAGCombiner::visitZERO_EXTEND(SDNode *N) {
       return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, SCC);
   }
 
+  // fold (zext x) -> (sext x) if the sign bit is known zero, and
+  // that is what the target prefers.
+  if (TLI.isSExtCheaperThanZExt(N0.getValueType(), VT) &&
+      (!LegalOperations || TLI.isOperationLegal(ISD::SIGN_EXTEND, VT)) &&
+      DAG.SignBitIsZero(N0))
+    return DAG.getNode(ISD::SIGN_EXTEND, DL, VT, N0);
+
+
   // (zext (shl (zext x), cst)) -> (shl (zext x), cst)
   if ((N0.getOpcode() == ISD::SHL || N0.getOpcode() == ISD::SRL) &&
       !TLI.isZExtFree(N0, VT)) {
diff --git a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
index 8b4f3159499122a..ee6f87374823702 100644
--- a/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp
@@ -2419,7 +2419,7 @@ bool TargetLowering::SimplifyDemandedBits(
     Known = Known.sext(BitWidth);
 
     // If the sign bit is known zero, convert this to a zero extend.
-    if (Known.isNonNegative()) {
+    if (Known.isNonNegative() && !isSExtCheaperThanZExt(SrcVT, VT)) {
       unsigned Opc =
           IsVecInReg ? ISD::ZERO_EXTEND_VECTOR_INREG : ISD::ZERO_EXTEND;
       if (!TLO.LegalOperations() || isOperationLegal(Opc, VT))