[RISCV] Teach getIntImmCostInst about (X & -(1 << C1) & 0xffffffff) == C2 << C1

[topperc]

We can rewrite this to (srai(w)/srli X, C1) == C2 so the AND immediate is
free. This transform is done by performSETCCCombine in RISCVISelLowering.cpp.

This fixes the opaque constant case mentioned in #157416.

[llvmbot]

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

Author: Craig Topper (topperc)

Changes

We can rewrite this to (sraiw X, C1) == C2 so the AND immediate is
free. This transform is done by performSETCCCombine in RISCVISelLowering.cpp.

This fixes the opaque constant case mentioned in #157416.

---

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

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp (+39)
• (modified) llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll (+52)

diff --git a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
index a06faa414a2ef..82892dfd28bd1 100644
--- a/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
+++ b/llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp
@@ -166,6 +166,42 @@ static bool canUseShiftPair(Instruction *Inst, const APInt &Imm) {
   return false;
 }
 
+// If this is a 64-bit AND with a mask of the form -(1 << C) in the lower 32
+// bits and the only user is an equality comparison, we might be able to use a
+// sraiw instead. This avoids the need to materialize the AND constant.
+static bool canUseSRAIWCmp(Instruction *Inst, const APInt &Imm) {
+  if (!Inst->hasOneUse())
+    return false;
+
+  // Look for equality comparison.
+  auto *Cmp = dyn_cast<ICmpInst>(*Inst->user_begin());
+  if (!Cmp || !Cmp->isEquality())
+    return false;
+
+  // Right hand side of comparison should be a constant.
+  auto *C = dyn_cast<ConstantInt>(Cmp->getOperand(1));
+  if (!C)
+    return false;
+
+  uint64_t Mask = Imm.getZExtValue();
+
+  // Mask should be of the form -(1 << C) in the lower 32 bits.
+  if (!isShiftedMask_64(Mask) || !isUInt<32>(Mask) ||
+      !isPowerOf2_32(-uint32_t(Mask)))
+    return false;
+
+  // Comparison constant should be a subset of Mask.
+  uint64_t CmpC = C->getZExtValue();
+  if ((CmpC & Mask) != CmpC)
+    return false;
+
+  // We'll need to sign extend the comparison constant and shift it right. Make
+  // sure the new constant can use addi/xori+seqz/snez.
+  unsigned ShiftBits = llvm::countr_zero(Mask);
+  int64_t NewCmpC = SignExtend64<32>(CmpC) >> ShiftBits;
+  return NewCmpC >= -2048 && NewCmpC <= 2048;
+}
+
 InstructionCost RISCVTTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
                                                 const APInt &Imm, Type *Ty,
                                                 TTI::TargetCostKind CostKind,
@@ -223,6 +259,9 @@ InstructionCost RISCVTTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
     if (Inst && Idx == 1 && Imm.getBitWidth() <= ST->getXLen() &&
         canUseShiftPair(Inst, Imm))
       return TTI::TCC_Free;
+    if (Inst && Idx == 1 && Imm.getBitWidth() == 64 &&
+        canUseSRAIWCmp(Inst, Imm))
+      return TTI::TCC_Free;
     Takes12BitImm = true;
     break;
   case Instruction::Add:
diff --git a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
index 55b0d1f0bf7be..2f95133380436 100644
--- a/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
+++ b/llvm/test/CodeGen/RISCV/and-negpow2-cmp.ll
@@ -155,3 +155,55 @@ define i1 @test9(i64 %x) {
   %b = icmp eq i64 %a, u0x08000000
   ret i1 %b
 }
+
+; Make sure the and constant doesn't get converted to an opaque constant by
+; ConstantHoisting. If it's an opaque constant, we'll have addi -16 and addi 15.
+define i64 @test10(i64 %0) #0 {
+; RV32-LABEL: test10:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    addi a0, a0, -1
+; RV32-NEXT:    andi a0, a0, -16
+; RV32-NEXT:    snez a0, a0
+; RV32-NEXT:    li a1, 0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test10:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    addi a0, a0, -1
+; RV64-NEXT:    sraiw a0, a0, 4
+; RV64-NEXT:    snez a0, a0
+; RV64-NEXT:    ret
+entry:
+  %1 = add nuw nsw i64 %0, u0xffffffff
+  %2 = and i64 %1, u0xfffffff0
+  %3 = icmp ne i64 %2, 0
+  %4 = zext i1 %3 to i64
+  ret i64 %4
+}
+
+; Make sure the and constant doesn't get converted to an opaque constant by
+; ConstantHoisting. If it's an opaque constant, we'll have addi -16 and addi 15.
+define i64 @test11(i64 %0) #0 {
+; RV32-LABEL: test11:
+; RV32:       # %bb.0: # %entry
+; RV32-NEXT:    addi a0, a0, -1
+; RV32-NEXT:    srai a0, a0, 4
+; RV32-NEXT:    addi a0, a0, 1621
+; RV32-NEXT:    seqz a0, a0
+; RV32-NEXT:    li a1, 0
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: test11:
+; RV64:       # %bb.0: # %entry
+; RV64-NEXT:    addi a0, a0, -1
+; RV64-NEXT:    sraiw a0, a0, 4
+; RV64-NEXT:    addi a0, a0, 1621
+; RV64-NEXT:    seqz a0, a0
+; RV64-NEXT:    ret
+entry:
+  %1 = add nuw nsw i64 %0, u0xffffffff
+  %2 = and i64 %1, u0xfffffff0
+  %3 = icmp eq i64 %2, u0xffff9ab0
+  %4 = zext i1 %3 to i64
+  ret i64 %4
+}

[pfusik]

This is not limited to SRAIW (also works on RV32), so please rename the function and update the comment and the commit message to avoid confusion.

[pfusik]

This is good and has a further potential for combining into addiw a0, a0, -3.

[pfusik]

Nit: mismatched parenthesis in the subject line.

[pfusik]

Since you renamed the function, you need to update its name at the call site.
Otherwise LGTM.