[RISCV] Prevent constant hoisting for (and (shl X, C), mask<<C)

[craig.topper]

Summary

If the immediate is a shifted mask, we will use a pair of shifts
and never materialize the immediate. Consider the immediate free.

[luismarques]

LGTM.

llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll

85

Nit: having a comment here with the rationale from the summary ("If the immediate is a shifted mask, we will use a pair of shifts and never materialize the immediate") would be convenient, since without the assembly code that might not be obvious.

[reames]

LGTM as well w/the same comment on adding comments. I'd also encourage a comment on the helper routine which describes the pattern you're matching and why that makes the immediate free.

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

if (Imm == 0)
     return TTI::TCC_Free;
 
   // Otherwise, we check how many instructions it will take to materialise.
   const DataLayout &DL = getDataLayout();
   return RISCVMatInt::getIntMatCost(Imm, DL.getTypeSizeInBits(Ty),
                                     getST()->getFeatureBits());
 }
 
+// Look for patterns of shift followed by AND that can be turned into a pair of
+// shifts. We won't need to materialize an immediate for the AND so these can
+// be considered free.
+static bool canUseShiftPair(Instruction *Inst, const APInt &Imm) {
+  uint64_t Mask = Imm.getZExtValue();
+  auto *BO = dyn_cast<BinaryOperator>(Inst->getOperand(0));
+  if (!BO || !BO->hasOneUse())
+    return false;
+
+  if (BO->getOpcode() != Instruction::Shl)
+    return false;
+
+  if (!isa<ConstantInt>(BO->getOperand(1)))
+    return false;
+
+  unsigned ShAmt = cast<ConstantInt>(BO->getOperand(1))->getZExtValue();
+  // (and (shl x, c2), c1) will be matched to (srli (slli x, c2+c3), c3) if c1
+  // is a mask shifted by c2 bits with c3 leading zeros.
+  if (isShiftedMask_64(Mask)) {
+    unsigned Trailing = countTrailingZeros(Mask);
+    if (ShAmt == Trailing)
+      return true;
+  }
+
+  return false;
+}
+
 InstructionCost RISCVTTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
                                                 const APInt &Imm, Type *Ty,
                                                 TTI::TargetCostKind CostKind,
                                                 Instruction *Inst) {
   assert(Ty->isIntegerTy() &&
          "getIntImmCost can only estimate cost of materialising integers");
 
   // We have a Zero register, so 0 is always free.
case Instruction::GetElementPtr:
     return TTI::TCC_Free;
   case Instruction::And:
     // zext.h
     if (Imm == UINT64_C(0xffff) && ST->hasStdExtZbb())
       return TTI::TCC_Free;
     // zext.w
     if (Imm == UINT64_C(0xffffffff) && ST->hasStdExtZba())
       return TTI::TCC_Free;
+    if (Inst && Idx == 1 && Imm.getBitWidth() <= ST->getXLen() &&
+        canUseShiftPair(Inst, Imm))
+      return TTI::TCC_Free;
     [[fallthrough]];
   case Instruction::Add:
   case Instruction::Or:
   case Instruction::Xor:
     Takes12BitImm = true;
     break;
   case Instruction::Mul:
     // Negated power of 2 is a shift and a negate.

llvm/test/Transforms/ConstantHoisting/RISCV/immediates.ll

 ; Check that we don't hoist mul with negated power of 2.
 define i64 @test9(i64 %a) nounwind {
 ; CHECK-LABEL: test9
 ; CHECK: mul i64 %a, -4294967296
   %1 = mul i64 %a, -4294967296
   %2 = mul i64 %1, -4294967296
   ret i64 %2
 }
+
+define i32 @test10(i32 %a, i32 %b) nounwind {
+; CHECK-LABEL: @test10(
+; CHECK: shl i32 %a, 8
+; CHECK: and i32 %1, 65280
+; CHECK: shl i32 %b, 8
+; CHECK: and i32 %3, 65280
+  %1 = shl i32 %a, 8
+  %2 = and i32 %1, 65280
+  %3 = shl i32 %b, 8
+  %4 = and i32 %3, 65280
+  %5 = mul i32 %2, %4
+  ret i32 %5
+}