[llvm] [RISCV] Don't look for sext in RISCVCodeGenPrepare::visitAnd. (PR #78798)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/78798

We want to know the upper 33 bits of the And Input are zero. SExt
only guarantees they are the same.
    
We originally checked for SExt or ZExt when we were using
isImpliedByDomCondition because a ZExt may have been changed to SExt
before we visited the And.
    
We are no longer using isImpliedByDomCondition so we can only look
for zext with the nneg flag.
    
While here, switch to PatternMatch to simplify the code.
    
Fixes #78783

>From 02622f65286a2371f49efc55718f69a7d9ad8abb Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Fri, 19 Jan 2024 14:07:49 -0800
Subject: [PATCH 1/2] [RISCV] Add test case for #78783. NFC

---
 llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
index 0e4e04af4a3fe7e..f58b01c7395af6d 100644
--- a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
+++ b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
@@ -91,3 +91,15 @@ for.body:                                         ; preds = %for.body, %for.body
   %niter.ncmp.1 = icmp eq i64 %niter.next.1, %unroll_iter
   br i1 %niter.ncmp.1, label %for.cond.cleanup.loopexit.unr-lcssa, label %for.body
 }
+
+; TODO: We should not change 4294967295 to -1 here.
+define i64 @bug(i32 %x) {
+; CHECK-LABEL: @bug(
+; CHECK-NEXT:    [[A:%.*]] = sext i32 [[X:%.*]] to i64
+; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], -1
+; CHECK-NEXT:    ret i64 [[B]]
+;
+  %a = sext i32 %x to i64
+  %b = and i64 %a, 4294967295
+  ret i64 %b
+}

>From 7e50f6690c9d33adb1ad926e003d9bbdc81901d8 Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Fri, 19 Jan 2024 14:21:09 -0800
Subject: [PATCH 2/2] [RISCV] Don't look for sext in
 RISCVCodeGenPrepare::visitAnd.

We want to know the upper 33 bits of the And Input are zero. SExt
only guarantees they are the same.

We originally checked for SExt or ZExt when we were using
isImpliedByDomCondition because a ZExt may have been changed to SExt
before we visited the And.

We are no longer using isImpliedByDomCondition so we can only look
for zext with the nneg flag.

While here, switch to PatternMatch to simplify the code.

Fixes #78783
---
 llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp  | 18 ++++++------------
 .../test/CodeGen/RISCV/riscv-codegenprepare.ll |  4 ++--
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
index 6434532afd40981..53fcc527e615dd4 100644
--- a/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
+++ b/llvm/lib/Target/RISCV/RISCVCodeGenPrepare.cpp
@@ -21,6 +21,7 @@
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/Intrinsics.h"
+#include "llvm/IR/PatternMatch.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 
@@ -67,20 +68,13 @@ bool RISCVCodeGenPrepare::visitAnd(BinaryOperator &BO) {
   if (!BO.getType()->isIntegerTy(64))
     return false;
 
-  auto canBeSignExtend = [](Instruction *I) {
-    if (isa<SExtInst>(I))
-      return true;
-    if (isa<ZExtInst>(I))
-      return I->hasNonNeg();
-    return false;
-  };
+  using namespace PatternMatch;
 
-  // Left hand side should be a sext or zext nneg.
-  Instruction *LHS = dyn_cast<Instruction>(BO.getOperand(0));
-  if (!LHS || !canBeSignExtend(LHS))
+  // Left hand side should be a zext nneg.
+  Value *LHSSrc;
+  if (!match(BO.getOperand(0), m_NNegZExt(m_Value(LHSSrc))))
     return false;
 
-  Value *LHSSrc = LHS->getOperand(0);
   if (!LHSSrc->getType()->isIntegerTy(32))
     return false;
 
@@ -100,7 +94,7 @@ bool RISCVCodeGenPrepare::visitAnd(BinaryOperator &BO) {
 
   // Sign extend the constant and replace the And operand.
   C = SignExtend64<32>(C);
-  BO.setOperand(1, ConstantInt::get(LHS->getType(), C));
+  BO.setOperand(1, ConstantInt::get(RHS->getType(), C));
 
   return true;
 }
diff --git a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
index f58b01c7395af6d..2179a0d26cf9826 100644
--- a/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
+++ b/llvm/test/CodeGen/RISCV/riscv-codegenprepare.ll
@@ -92,11 +92,11 @@ for.body:                                         ; preds = %for.body, %for.body
   br i1 %niter.ncmp.1, label %for.cond.cleanup.loopexit.unr-lcssa, label %for.body
 }
 
-; TODO: We should not change 4294967295 to -1 here.
+; Make sure we do not change 4294967295 to -1 here.
 define i64 @bug(i32 %x) {
 ; CHECK-LABEL: @bug(
 ; CHECK-NEXT:    [[A:%.*]] = sext i32 [[X:%.*]] to i64
-; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], -1
+; CHECK-NEXT:    [[B:%.*]] = and i64 [[A]], 4294967295
 ; CHECK-NEXT:    ret i64 [[B]]
 ;
   %a = sext i32 %x to i64