diff --git a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
--- a/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FunctionLoweringInfo.cpp
@@ -58,9 +58,14 @@
 }
 
 static ISD::NodeType getPreferredExtendForValue(const Value *V) {
-  // For the users of the source value being used for compare instruction, if
-  // the number of signed predicate is greater than unsigned predicate, we
-  // prefer to use SIGN_EXTEND.
+  // For the users of the source value being used for compare instruction or
+  // PHI node:
+  //
+  // 1. if PHI node is only used for return instruction which is required
+  //    to signext or zeroext, we set default ExtendKind to SIGN_EXTEND or
+  //    ZERO_EXTEND.
+  // 2. if the number of signed predicate is greater than unsigned
+  //    predicate, we prefer to use SIGN_EXTEND.
   //
   // With this optimization, we would be able to reduce some redundant sign or
   // zero extension instruction, and eventually more machine CSE opportunities
@@ -71,6 +76,19 @@
     if (const auto *CI = dyn_cast<CmpInst>(U)) {
       NumOfSigned += CI->isSigned();
       NumOfUnsigned += CI->isUnsigned();
+    } else if (const auto *PN = dyn_cast<PHINode>(U)) {
+      if (PN->hasOneUser()) {
+        if (const auto RI = dyn_cast<ReturnInst>(U->user_back())) {
+          const Function *F = RI->getParent()->getParent();
+          if (F->getAttributes().hasRetAttr(Attribute::SExt)) {
+            ExtendKind = ISD::SIGN_EXTEND;
+            NumOfSigned += 1;
+          } else if (F->getAttributes().hasRetAttr(Attribute::ZExt)) {
+            ExtendKind = ISD::ZERO_EXTEND;
+            NumOfUnsigned += 1;
+          }
+        }
+      }
     }
   }
   if (NumOfSigned > NumOfUnsigned)
diff --git a/llvm/test/CodeGen/RISCV/prefer-extend.ll b/llvm/test/CodeGen/RISCV/prefer-extend.ll
new file mode 100644
--- /dev/null
+++ b/llvm/test/CodeGen/RISCV/prefer-extend.ll
@@ -0,0 +1,53 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv32 < %s | FileCheck -check-prefixes=RV32S %s
+; RUN: llc -mtriple=riscv64 < %s | FileCheck -check-prefixes=RV64S %s
+
+define dso_local signext i16 @get_pivot(i16 signext %left, i16 signext %right) {
+; RV32S-LABEL: get_pivot:
+; RV32S:       # %bb.0: # %entry
+; RV32S-NEXT:    sub a2, a1, a0
+; RV32S-NEXT:    slli a2, a2, 16
+; RV32S-NEXT:    srai a2, a2, 16
+; RV32S-NEXT:    blt a2, a0, .LBB0_3
+; RV32S-NEXT:  # %bb.1: # %if.end1
+; RV32S-NEXT:    mv a0, a1
+; RV32S-NEXT:    blt a1, a2, .LBB0_3
+; RV32S-NEXT:  # %bb.2: # %if.end2
+; RV32S-NEXT:    mv a0, a2
+; RV32S-NEXT:  .LBB0_3: # %cleanup
+; RV32S-NEXT:    ret
+;
+; RV64S-LABEL: get_pivot:
+; RV64S:       # %bb.0: # %entry
+; RV64S-NEXT:    sub a2, a1, a0
+; RV64S-NEXT:    slli a2, a2, 48
+; RV64S-NEXT:    srai a2, a2, 48
+; RV64S-NEXT:    blt a2, a0, .LBB0_3
+; RV64S-NEXT:  # %bb.1: # %if.end1
+; RV64S-NEXT:    mv a0, a1
+; RV64S-NEXT:    blt a1, a2, .LBB0_3
+; RV64S-NEXT:  # %bb.2: # %if.end2
+; RV64S-NEXT:    mv a0, a2
+; RV64S-NEXT:  .LBB0_3: # %cleanup
+; RV64S-NEXT:    ret
+entry:
+  %conv = sext i16 %left to i32
+  %conv1 = sext i16 %right to i32
+  %sub = sub nsw i32 %conv1, %conv
+  %conv2 = trunc i32 %sub to i16
+  %conv3 = sext i16 %conv2 to i32
+  %cmp = icmp slt i32 %conv3, %conv
+  br i1 %cmp, label %cleanup, label %if.end1
+
+if.end1:                                          ; preds = %entry
+  %cmp9 = icmp sgt i32 %conv3, %conv1
+  br i1 %cmp9, label %cleanup, label %if.end2
+
+if.end2:                                          ; preds = %if.end1
+  br label %cleanup
+
+cleanup:                                          ; preds = %if.end1, %if.end2, %entry
+  %retval.0 = phi i16 [ %left, %entry ], [ %right, %if.end1 ], [ %conv2, %if.end2 ]
+  ret i16 %retval.0
+}
+