[RISCV] Rewrite RISCVCodeGenPrepare using zext nneg [nfc-ish]

[preames]

This stacks on #70725. Once we have lowering for zext nneg, we can rewrite all of the existing RISCVCodeGenPrepare login in terms of zext nneg instead of sext. The change isn't NFC from the perspective of the individual pass, but should be from the perspective of codegen as a whole.

As noted in the TODO, one piece can be moved to instcombine, but I'll leave that to a separate commit.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Philip Reames (preames)

Changes

This stacks on #70725. Once we have lowering for zext nneg, we can rewrite all of the existing RISCVCodeGenPrepare login in terms of zext nneg instead of sext. The change isn't NFC from the perspective of the individual pass, but should be from the perspective of codegen as a whole.

As noted in the TODO, one piece can be moved to instcombine, but I'll leave that to a separate commit.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp (+22-28)
• (modified) llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll (+2-2)

diff --git a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
index 2fcd9a40588a733..7bc7e3924ca7026 100644
--- a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
+++ b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
@@ -62,38 +62,32 @@ bool RISCVCodeGenPrepare::visitZExtInst(ZExtInst &ZExt) {
   if (!ST->is64Bit())
     return false;
 
+  if (ZExt.hasNonNeg())
+    return false;
+
   Value *Src = ZExt.getOperand(0);
 
   // We only care about ZExt from i32 to i64.
   if (!ZExt.getType()->isIntegerTy(64) || !Src->getType()->isIntegerTy(32))
     return false;
 
-  // Look for an opportunity to replace (i64 (zext (i32 X))) with a sext if we
-  // can determine that the sign bit of X is zero via a dominating condition.
-  // This often occurs with widened induction variables.
+  // Look for an opportunity to infer nneg on a zext if we can determine that
+  // the sign bit of X is zero via a dominating condition. This often occurs
+  // with widened induction variables.
   if (isImpliedByDomCondition(ICmpInst::ICMP_SGE, Src,
                               Constant::getNullValue(Src->getType()), &ZExt,
                               *DL).value_or(false)) {
-    auto *SExt = new SExtInst(Src, ZExt.getType(), "", &ZExt);
-    SExt->takeName(&ZExt);
-    SExt->setDebugLoc(ZExt.getDebugLoc());
-
-    ZExt.replaceAllUsesWith(SExt);
-    ZExt.eraseFromParent();
+    ZExt.setNonNeg(true);
     ++NumZExtToSExt;
     return true;
   }
 
-  // Convert (zext (abs(i32 X, i1 1))) -> (sext (abs(i32 X, i1 1))). If abs of
+  // Convert (zext (abs(i32 X, i1 1))) -> (zext nneg (abs(i32 X, i1 1))). If abs of
   // INT_MIN is poison, the sign bit is zero.
+  // TODO: Move this to instcombine now that we have zext nneg in IR.
   using namespace PatternMatch;
   if (match(Src, m_Intrinsic<Intrinsic::abs>(m_Value(), m_One()))) {
-    auto *SExt = new SExtInst(Src, ZExt.getType(), "", &ZExt);
-    SExt->takeName(&ZExt);
-    SExt->setDebugLoc(ZExt.getDebugLoc());
-
-    ZExt.replaceAllUsesWith(SExt);
-    ZExt.eraseFromParent();
+    ZExt.setNonNeg(true);
     ++NumZExtToSExt;
     return true;
   }
@@ -102,9 +96,8 @@ bool RISCVCodeGenPrepare::visitZExtInst(ZExtInst &ZExt) {
 }
 
 // Try to optimize (i64 (and (zext/sext (i32 X), C1))) if C1 has bit 31 set,
-// but bits 63:32 are zero. If we can prove that bit 31 of X is 0, we can fill
-// the upper 32 bits with ones. A separate transform will turn (zext X) into
-// (sext X) for the same condition.
+// but bits 63:32 are zero. If we know that bit 31 of X is 0, we can fill
+// the upper 32 bits with ones.
 bool RISCVCodeGenPrepare::visitAnd(BinaryOperator &BO) {
   if (!ST->is64Bit())
     return false;
@@ -112,9 +105,17 @@ bool RISCVCodeGenPrepare::visitAnd(BinaryOperator &BO) {
   if (!BO.getType()->isIntegerTy(64))
     return false;
 
-  // Left hand side should be sext or zext.
+  auto canBeSignExtend = [](Instruction *I) {
+    if (isa<SExtInst>(I))
+      return true;
+    if (isa<ZExtInst>(I))
+      return I->hasNonNeg();
+    return false;
+  };
+
+  // Left hand side should be a sext or zext nneg.
   Instruction *LHS = dyn_cast<Instruction>(BO.getOperand(0));
-  if (!LHS || (!isa<SExtInst>(LHS) && !isa<ZExtInst>(LHS)))
+  if (!LHS || !canBeSignExtend(LHS))
     return false;
 
   Value *LHSSrc = LHS->getOperand(0);
@@ -135,13 +136,6 @@ bool RISCVCodeGenPrepare::visitAnd(BinaryOperator &BO) {
   if (!isUInt<32>(C) || isInt<12>(C) || !isInt<12>(SignExtend64<32>(C)))
     return false;
 
-  // If we can determine the sign bit of the input is 0, we can replace the
-  // And mask constant.
-  if (!isImpliedByDomCondition(ICmpInst::ICMP_SGE, LHSSrc,
-                               Constant::getNullValue(LHSSrc->getType()),
-                               LHS, *DL).value_or(false))
-    return false;
-
   // Sign extend the constant and replace the And operand.
   C = SignExtend64<32>(C);
   BO.setOperand(1, ConstantInt::get(LHS->getType(), C));
diff --git a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
index b7530c80c417a13..b4f0918635650b9 100644
--- a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
+++ b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
@@ -9,7 +9,7 @@ define void @test1(ptr nocapture noundef %a, i32 noundef signext %n) {
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp sgt i32 [[N:%.*]], 0
 ; CHECK-NEXT:    br i1 [[CMP3]], label [[FOR_BODY_PREHEADER:%.*]], label [[FOR_COND_CLEANUP:%.*]]
 ; CHECK:       for.body.preheader:
-; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = sext i32 [[N]] to i64
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext nneg i32 [[N]] to i64
 ; CHECK-NEXT:    br label [[FOR_BODY:%.*]]
 ; CHECK:       for.cond.cleanup.loopexit:
 ; CHECK-NEXT:    br label [[FOR_COND_CLEANUP]]
@@ -60,7 +60,7 @@ define void @test2(ptr nocapture noundef %a, i32 noundef signext %n) {
 ; CHECK-NEXT:    [[CMP3:%.*]] = icmp sgt i32 [[N:%.*]], 0
 ; CHECK-NEXT:    br i1 [[CMP3]], label [[FOR_BODY_PREHEADER:%.*]], label [[FOR_COND_CLEANUP:%.*]]
 ; CHECK:       for.body.preheader:
-; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = sext i32 [[N]] to i64
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext nneg i32 [[N]] to i64
 ; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[WIDE_TRIP_COUNT]], 1
 ; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i32 [[N]], 1
 ; CHECK-NEXT:    br i1 [[TMP0]], label [[FOR_COND_CLEANUP_LOOPEXIT_UNR_LCSSA:%.*]], label [[FOR_BODY_PREHEADER_NEW:%.*]]

[github-actions]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff 83c560b3bf46e9b5a65f9a41b60e21898e286c9c 64ddf53850c4ef5197157059819217ee6bf974f6 -- llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp

View the diff from clang-format here.

diff --git a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
index 7bc7e3924..ff0272ff7 100644
--- a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
+++ b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
@@ -82,8 +82,8 @@ bool RISCVCodeGenPrepare::visitZExtInst(ZExtInst &ZExt) {
     return true;
   }
 
-  // Convert (zext (abs(i32 X, i1 1))) -> (zext nneg (abs(i32 X, i1 1))). If abs of
-  // INT_MIN is poison, the sign bit is zero.
+  // Convert (zext (abs(i32 X, i1 1))) -> (zext nneg (abs(i32 X, i1 1))). If abs
+  // of INT_MIN is poison, the sign bit is zero.
   // TODO: Move this to instcombine now that we have zext nneg in IR.
   using namespace PatternMatch;
   if (match(Src, m_Intrinsic<Intrinsic::abs>(m_Value(), m_One()))) {

[topperc]

LGTM