[RISCV] Prevent constant hoisting for (ashr (add (shl X, 32), C<<32), 32).

[craig.topper]

Summary

Instead of hoisting the C<<32, keep in place so DAGCombiner can
match the pattern to an addw.

[ributzka]

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

42

Is this DAGCombine RISCV specific or a generic combine?

[craig.topper]

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

42

It's RISC-V specific. See https://reviews.llvm.org/D128869

[craig.topper]

Explicitly check the opcode of Inst in isShiftedADDWOrSUBW. This can get called with Opcode
being Instruction::Add, but the instruction being a GEP.

[ributzka]

RISCV is not my expertise, but LGTM with respect to constant hoisting.

[craig.topper]

Make sure all 3 instructions are in the same BB to be sure the pattern can be matched.

[craig.topper]

Abandoning in favor a pre-consthoist peephole in RISCVCodeGenPrepare. The single basic block restriction is too restrictive. We need to help SelectionDAG out here.

llvm/lib/Target/RISCV/RISCVTargetTransformInfo.cpp

 //===-- RISCVTargetTransformInfo.cpp - RISC-V specific TTI ----------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "RISCVTargetTransformInfo.h"
 #include "MCTargetDesc/RISCVMatInt.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/CodeGen/BasicTTIImpl.h"
 #include "llvm/CodeGen/TargetLowering.h"
+#include "llvm/IR/PatternMatch.h"
 #include <cmath>
 using namespace llvm;
 
 #define DEBUG_TYPE "riscvtti"
 
 static cl::opt<unsigned> RVVRegisterWidthLMUL(
     "riscv-v-register-bit-width-lmul",
     cl::desc(
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());
 }
 
+// Recognize (ashr (add (shl X, 32), C << 32), 32) which DAGCombine can turn
+// into (addw X, C). If we hoist the constant, DAGCombine can't see the pattern.
+// The shl could also be a mul with -(1 << 32). The add can also be a sub with
+// a constant left hand side. We assume InstCombine canonicalized the add/sub.
+static bool isShiftedADDWOrSUBW(Instruction *Inst, const APInt &Imm) {
+  using namespace PatternMatch;
+  if (Inst->getOpcode() != Instruction::Add &&
+      Inst->getOpcode() != Instruction::Sub)
+    return false;
+  // Needs to be a 64-bit immediate with at least 32 trailing zeros.
+  if (Imm.getBitWidth() != 64 || Imm.countTrailingZeros() < 32)
+    return false;
+  unsigned ShlIdx = Inst->getOpcode() == Instruction::Add ? 0 : 1;
+  Value *ShlOp = Inst->getOperand(ShlIdx);
+  // Non-immediate source should be a shl by 32 or a mul by -(1 << 32).
+  if (!match(ShlOp, m_Shl(m_Value(), m_SpecificInt(32))) &&
+      !match(ShlOp, m_Mul(m_Value(), m_SpecificInt(-1ULL << 32))))
+    return false;
+  BasicBlock *BB = Inst->getParent();
+  // Need to be in the same basic block.
+  if (cast<Instruction>(ShlOp)->getParent() != BB)
+    return false;
+  // User should be an ashr by 32.
+  return Inst->hasOneUse() &&
+         match(*Inst->user_begin(), m_AShr(m_Value(), m_SpecificInt(32))) &&
+         cast<Instruction>(*Inst->user_begin())->getParent() == BB;
+}
+
 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.
InstructionCost RISCVTTIImpl::getIntImmCostInst(unsigned Opcode, unsigned Idx,
   case Instruction::AShr:
     Takes12BitImm = true;
     ImmArgIdx = 1;
     break;
   default:
     break;
   }
 
+  if (Inst && ST->is64Bit() && isShiftedADDWOrSUBW(Inst, Imm))
+    return TTI::TCC_Free;
+
   if (Takes12BitImm) {
     // Check immediate is the correct argument...
     if (Instruction::isCommutative(Opcode) || Idx == ImmArgIdx) {
       // ... and fits into the 12-bit immediate.
       if (Imm.getMinSignedBits() <= 64 &&
           getTLI()->isLegalAddImmediate(Imm.getSExtValue())) {
         return TTI::TCC_Free;
       }

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

 ; RUN: opt -mtriple=riscv32-unknown-elf -S -consthoist < %s | FileCheck %s
-; RUN: opt -mtriple=riscv64-unknown-elf -S -consthoist < %s | FileCheck %s
+; RUN: opt -mtriple=riscv64-unknown-elf -S -consthoist < %s | FileCheck %s --check-prefixes=CHECK,RV64
 
 ; Check that we don't hoist immediates with small values.
 define i64 @test1(i64 %a) nounwind {
 ; CHECK-LABEL: test1
 ; CHECK-NOT: %const = bitcast i64 2 to i64
   %1 = mul i64 %a, 2
   %2 = add i64 %1, 2
   ret i64 %2

 ; 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
 }
+
+; The mul is equivalent to (neg (shl X, 32)). The mul+add+ashr can be replaced
+; with subw using a much smaller constant.
+define i64 @test10(i64 %a) nounwind {
+; RV64-LABEL: test10
+; RV64: add i64 %1, 4294967296
+; RV64: add i64 %4, 4294967296
+  %1 = mul i64 %a, -4294967296
+  %2 = add i64 %1, 4294967296
+  %3 = ashr i64 %1, 32
+  %4 = mul i64 %3, -4294967296
+  %5 = add i64 %4, 4294967296
+  %6 = ashr i64 %5, 32
+  ret i64 %6
+}
+
+; The mul is equivalent to (neg (shl X, 32)). The mul+sub+ashr can be replaced
+; with addw using a much smaller constant.
+define i64 @test11(i64 %a) nounwind {
+; RV64-LABEL: test11
+; RV64: sub i64 4294967296, %1
+; RV64: sub i64 4294967296, %4
+  %1 = mul i64 %a, -4294967296
+  %2 = sub i64 4294967296, %1
+  %3 = ashr i64 %1, 32
+  %4 = mul i64 %3, -4294967296
+  %5 = sub i64 4294967296, %4
+  %6 = ashr i64 %5, 32
+  ret i64 %6
+}
+
+; The shl+add+ashr can be replaced with subw using a much smaller constant.
+define i64 @test12(i64 %a) nounwind {
+; RV64-LABEL: test12
+; RV64: add i64 %1, 4294967296
+; RV64: add i64 %4, 4294967296
+  %1 = shl i64 %a, 32
+  %2 = add i64 %1, 4294967296
+  %3 = ashr i64 %1, 32
+  %4 = shl i64 %3, 32
+  %5 = add i64 %4, 4294967296
+  %6 = ashr i64 %5, 32
+  ret i64 %6
+}
+
+; The shl+shl+ashr can be replaced with addw using a much smaller constant.
+define i64 @test13(i64 %a) nounwind {
+; RV64-LABEL: test13
+; RV64: sub i64 4294967296, %1
+; RV64: sub i64 4294967296, %4
+  %1 = shl i64 %a, 32
+  %2 = sub i64 4294967296, %1
+  %3 = ashr i64 %1, 32
+  %4 = shl i64 %3, 32
+  %5 = sub i64 4294967296, %4
+  %6 = ashr i64 %5, 32
+  ret i64 %6
+}