[llvm] [SDAG] Prefer forming sign_extend for zext nneg per target preference (PR #70725)

[Philip Reames via llvm-commits]

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

Builds on #67982 which recently introduced the nneg flag on a zext instruction.  Note that this change is the first point where the flag is being used for an optimization, and thus may expose latent miscompiles. We've recently taught both CVP and InstCombine to infer the flag when forming zext, but nothing else is using the flag just yet.

>From 369b36d79415fc91156fcf2bcd9deb93ea65521b Mon Sep 17 00:00:00 2001
From: Philip Reames <preames at rivosinc.com>
Date: Mon, 30 Oct 2023 13:51:38 -0700
Subject: [PATCH] [SDAG] Prefer forming sign_extend for zext nneg per target
 preference

Builds on #67982 which recently introduced the nneg flag on a zext
instruction.  Note that this change is the first point where the flag
is being used for an optimization, and thus may expose latent miscompiles.
We've recently taught both CVP and InstCombine to infer the flag when
forming zext, but nothing else is using the flag just yet.
---
 .../CodeGen/SelectionDAG/SelectionDAGBuilder.cpp | 16 ++++++++++++++--
 llvm/test/CodeGen/RISCV/sext-zext-trunc.ll       |  3 +--
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 0e6129aaf521922..c518b1f95e9023e 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3524,8 +3524,20 @@ void SelectionDAGBuilder::visitZExt(const User &I) {
   // ZExt cannot be a no-op cast because sizeof(src) < sizeof(dest).
   // ZExt also can't be a cast to bool for same reason. So, nothing much to do
   SDValue N = getValue(I.getOperand(0));
-  EVT DestVT = DAG.getTargetLoweringInfo().getValueType(DAG.getDataLayout(),
-                                                        I.getType());
+  auto &TLI = DAG.getTargetLoweringInfo();
+  EVT DestVT = TLI.getValueType(DAG.getDataLayout(), I.getType());
+
+  // Since we don't yet have a representation of zext nneg in SDAG or MI,
+  // eagerly use the information to canonicalize towards sign_extend if
+  // that is the target's preference.  TODO: Add nneg support to the
+  // SDAG and MI representations.
+  if (auto *PNI = dyn_cast<PossiblyNonNegInst>(&I);
+      PNI && PNI->hasNonNeg() &&
+      TLI.isSExtCheaperThanZExt(N.getValueType(), DestVT)) {
+    setValue(&I, DAG.getNode(ISD::SIGN_EXTEND, getCurSDLoc(), DestVT, N));
+    return;
+  }
+
   setValue(&I, DAG.getNode(ISD::ZERO_EXTEND, getCurSDLoc(), DestVT, N));
 }
 
diff --git a/llvm/test/CodeGen/RISCV/sext-zext-trunc.ll b/llvm/test/CodeGen/RISCV/sext-zext-trunc.ll
index 7297bfaf0c62ec7..20d73acddea01b0 100644
--- a/llvm/test/CodeGen/RISCV/sext-zext-trunc.ll
+++ b/llvm/test/CodeGen/RISCV/sext-zext-trunc.ll
@@ -501,8 +501,7 @@ define i64 @zext_nneg_i32_to_i64(i32 %a) nounwind {
 ;
 ; RV64-LABEL: zext_nneg_i32_to_i64:
 ; RV64:       # %bb.0:
-; RV64-NEXT:    slli a0, a0, 32
-; RV64-NEXT:    srli a0, a0, 32
+; RV64-NEXT:    sext.w a0, a0
 ; RV64-NEXT:    ret
   %1 = zext nneg i32 %a to i64
   ret i64 %1