[llvm] [DAGCombiner][RISCV] Prefer to sext i32 non-negative values (PR #65984)

[Yingwei Zheng via llvm-commits]

https://github.com/dtcxzyw created https://github.com/llvm/llvm-project/pull/65984:

By default, `DAGCombiner` folds `sext x` to `zext x` when `x` is non-negative. It will generate redundant `zext` inst seq on riscv64 (typically `slli (srli x, 32), 32`).
godbolt: https://godbolt.org/z/osf6adP1o
This patch applies the transform iff `zext` is **cheaper** than `sext`.


>From c04ac9119a22544cf247271adcd8ce048bfb24c8 Mon Sep 17 00:00:00 2001
From: Yingwei Zheng <dtcxzyw2333 at gmail.com>
Date: Tue, 12 Sep 2023 01:06:30 +0800
Subject: [PATCH] [DAGCombiner][RISCV] Prefer to sext i32 non-negative values

---
 llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp |  3 ++-
 llvm/test/CodeGen/RISCV/aext-to-sext.ll       | 13 ++++++++++++-
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
index d917e8c00c4f92a..f574c746142f959 100644
--- a/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp
@@ -13271,7 +13271,8 @@ SDValue DAGCombiner::visitSIGN_EXTEND(SDNode *N) {
     return V;
 
   // fold (sext x) -> (zext x) if the sign bit is known zero.
-  if ((!LegalOperations || TLI.isOperationLegal(ISD::ZERO_EXTEND, VT)) &&
+  if (!TLI.isSExtCheaperThanZExt(N0.getValueType(), VT) &&
+      (!LegalOperations || TLI.isOperationLegal(ISD::ZERO_EXTEND, VT)) &&
       DAG.SignBitIsZero(N0))
     return DAG.getNode(ISD::ZERO_EXTEND, DL, VT, N0);
 
diff --git a/llvm/test/CodeGen/RISCV/aext-to-sext.ll b/llvm/test/CodeGen/RISCV/aext-to-sext.ll
index 6322f84bad0baae..09803012001c26c 100644
--- a/llvm/test/CodeGen/RISCV/aext-to-sext.ll
+++ b/llvm/test/CodeGen/RISCV/aext-to-sext.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=riscv64 -verify-machineinstrs < %s \
+; RUN: llc -mtriple=riscv64 -mattr=+m -verify-machineinstrs < %s \
 ; RUN:   | FileCheck %s -check-prefix=RV64I
 
 ; Make sure we don't generate an addi in the loop in
@@ -100,3 +100,14 @@ merge:
   %d = zext i32 %b to i64
   ret i64 %d
 }
+
+; We prefer to sign extend i32 non-negative values. The default behavior in
+; DAGCombiner is zero extend. We have a target hook to override it.
+define signext i32 @square(i32 signext %num) {
+; RV64I-LABEL: square:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    mulw a0, a0, a0
+; RV64I-NEXT:    ret
+  %mul = mul nsw i32 %num, %num
+  ret i32 %mul
+}