[SDAG] Prefer forming sign_extend for zext nneg per target preference

[preames]

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.

[llvmbot]

@llvm/pr-subscribers-llvm-selectiondag

Author: Philip Reames (preames)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/70725.diff

2 Files Affected:

• (modified) llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp (+14-2)
• (modified) llvm/test/CodeGen/RISCV/sext-zext-trunc.ll (+1-2)

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

[nikic]

LGTM

This seems like a reasonable starting point. Something I'm a bit concerned about is that this may break DAGCombiner patterns that look for ZERO_EXTEND. Though given that RISCVCodeGenPrepare already does this replacement aggressively, I guess it's not a big issue in practice.