//===- MVETailPredication.cpp - MVE Tail Predication ------------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
/// \file
/// Armv8.1m introduced MVE, M-Profile Vector Extension, and low-overhead
/// branches to help accelerate DSP applications. These two extensions,
/// combined with a new form of predication called tail-predication, can be used
/// to provide implicit vector predication within a low-overhead loop.
/// This is implicit because the predicate of active/inactive lanes is
/// calculated by hardware, and thus does not need to be explicitly passed
/// to vector instructions. The instructions responsible for this are the
/// DLSTP and WLSTP instructions, which setup a tail-predicated loop and the
/// the total number of data elements processed by the loop. The loop-end
/// LETP instruction is responsible for decrementing and setting the remaining
/// elements to be processed and generating the mask of active lanes.
///
/// The HardwareLoops pass inserts intrinsics identifying loops that the
/// backend will attempt to convert into a low-overhead loop. The vectorizer is
/// responsible for generating a vectorized loop in which the lanes are
/// predicated upon an get.active.lane.mask intrinsic. This pass looks at these
/// get.active.lane.mask intrinsic and attempts to convert them to VCTP
/// instructions. This will be picked up by the ARM Low-overhead loop pass later
/// in the backend, which performs the final transformation to a DLSTP or WLSTP
/// tail-predicated loop.
//
//===----------------------------------------------------------------------===//

#include "ARM.h"
#include "ARMSubtarget.h"
#include "ARMTargetTransformInfo.h"
#include "llvm/Analysis/LoopInfo.h"
#include "llvm/Analysis/LoopPass.h"
#include "llvm/Analysis/ScalarEvolution.h"
#include "llvm/Analysis/ScalarEvolutionExpressions.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/CodeGen/TargetPassConfig.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsARM.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/InitializePasses.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/LoopUtils.h"
#include "llvm/Transforms/Utils/ScalarEvolutionExpander.h"

using namespace llvm;

#define DEBUG_TYPE "mve-tail-predication"
#define DESC "Transform predicated vector loops to use MVE tail predication"

cl::opt<TailPredication::Mode> EnableTailPredication(
   "tail-predication", cl::desc("MVE tail-predication pass options"),
   cl::init(TailPredication::Enabled),
   cl::values(clEnumValN(TailPredication::Disabled, "disabled",
                         "Don't tail-predicate loops"),
              clEnumValN(TailPredication::EnabledNoReductions,
                         "enabled-no-reductions",
                         "Enable tail-predication, but not for reduction loops"),
              clEnumValN(TailPredication::Enabled,
                         "enabled",
                         "Enable tail-predication, including reduction loops"),
              clEnumValN(TailPredication::ForceEnabledNoReductions,
                         "force-enabled-no-reductions",
                         "Enable tail-predication, but not for reduction loops, "
                         "and force this which might be unsafe"),
              clEnumValN(TailPredication::ForceEnabled,
                         "force-enabled",
                         "Enable tail-predication, including reduction loops, "
                         "and force this which might be unsafe")));


namespace {

class MVETailPredication : public LoopPass {
  SmallVector<IntrinsicInst*, 4> MaskedInsts;
  Loop *L = nullptr;
  ScalarEvolution *SE = nullptr;
  TargetTransformInfo *TTI = nullptr;
  const ARMSubtarget *ST = nullptr;

public:
  static char ID;

  MVETailPredication() : LoopPass(ID) { }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.addRequired<ScalarEvolutionWrapperPass>();
    AU.addRequired<LoopInfoWrapperPass>();
    AU.addRequired<TargetPassConfig>();
    AU.addRequired<TargetTransformInfoWrapperPass>();
    AU.addPreserved<LoopInfoWrapperPass>();
    AU.setPreservesCFG();
  }

  bool runOnLoop(Loop *L, LPPassManager&) override;

private:
  /// Perform the relevant checks on the loop and convert active lane masks if
  /// possible.
  bool TryConvertActiveLaneMask(Value *TripCount);

  /// Perform several checks on the arguments of @llvm.get.active.lane.mask
  /// intrinsic. E.g., check that the loop induction variable and the element
  /// count are of the form we expect, and also perform overflow checks for
  /// the new expressions that are created.
  const SCEV *IsSafeActiveMask(IntrinsicInst *ActiveLaneMask, Value *TripCount);

  /// Insert the intrinsic to represent the effect of tail predication.
  void InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask, Value *Start);
};

} // end namespace

bool MVETailPredication::runOnLoop(Loop *L, LPPassManager&) {
  if (skipLoop(L) || !EnableTailPredication)
    return false;

  MaskedInsts.clear();
  Function &F = *L->getHeader()->getParent();
  auto &TPC = getAnalysis<TargetPassConfig>();
  auto &TM = TPC.getTM<TargetMachine>();
  ST = &TM.getSubtarget<ARMSubtarget>(F);
  TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
  SE = &getAnalysis<ScalarEvolutionWrapperPass>().getSE();
  this->L = L;

  // The MVE and LOB extensions are combined to enable tail-predication, but
  // there's nothing preventing us from generating VCTP instructions for v8.1m.
  if (!ST->hasMVEIntegerOps() || !ST->hasV8_1MMainlineOps()) {
    LLVM_DEBUG(dbgs() << "ARM TP: Not a v8.1m.main+mve target.\n");
    return false;
  }

  BasicBlock *Preheader = L->getLoopPreheader();
  if (!Preheader)
    return false;

  auto FindLoopIterations = [](BasicBlock *BB) -> IntrinsicInst* {
    for (auto &I : *BB) {
      auto *Call = dyn_cast<IntrinsicInst>(&I);
      if (!Call)
        continue;

      Intrinsic::ID ID = Call->getIntrinsicID();
      if (ID == Intrinsic::start_loop_iterations ||
          ID == Intrinsic::test_start_loop_iterations)
        return cast<IntrinsicInst>(&I);
    }
    return nullptr;
  };

  // Look for the hardware loop intrinsic that sets the iteration count.
  IntrinsicInst *Setup = FindLoopIterations(Preheader);

  // The test.set iteration could live in the pre-preheader.
  if (!Setup) {
    if (!Preheader->getSinglePredecessor())
      return false;
    Setup = FindLoopIterations(Preheader->getSinglePredecessor());
    if (!Setup)
      return false;
  }

  LLVM_DEBUG(dbgs() << "ARM TP: Running on Loop: " << *L << *Setup << "\n");

  bool Changed = TryConvertActiveLaneMask(Setup->getArgOperand(0));

  return Changed;
}

// The active lane intrinsic has this form:
//
//    @llvm.get.active.lane.mask(IV, TC)
//
// Here we perform checks that this intrinsic behaves as expected,
// which means:
//
// 1) Check that the TripCount (TC) belongs to this loop (originally).
// 2) The element count (TC) needs to be sufficiently large that the decrement
//    of element counter doesn't overflow, which means that we need to prove:
//        ceil(ElementCount / VectorWidth) >= TripCount
//    by rounding up ElementCount up:
//        ((ElementCount + (VectorWidth - 1)) / VectorWidth
//    and evaluate if expression isKnownNonNegative:
//        (((ElementCount + (VectorWidth - 1)) / VectorWidth) - TripCount
// 3) The IV must be an induction phi with an increment equal to the
//    vector width.
const SCEV *MVETailPredication::IsSafeActiveMask(IntrinsicInst *ActiveLaneMask,
                                                 Value *TripCount) {
  bool ForceTailPredication =
    EnableTailPredication == TailPredication::ForceEnabledNoReductions ||
    EnableTailPredication == TailPredication::ForceEnabled;

  Value *ElemCount = ActiveLaneMask->getOperand(1);
  bool Changed = false;
  if (!L->makeLoopInvariant(ElemCount, Changed))
    return nullptr;

  auto *EC= SE->getSCEV(ElemCount);
  auto *TC = SE->getSCEV(TripCount);
  int VectorWidth =
      cast<FixedVectorType>(ActiveLaneMask->getType())->getNumElements();
  if (VectorWidth != 2 && VectorWidth != 4 && VectorWidth != 8 &&
      VectorWidth != 16)
    return nullptr;
  ConstantInt *ConstElemCount = nullptr;

  // 1) Smoke tests that the original scalar loop TripCount (TC) belongs to
  // this loop.  The scalar tripcount corresponds the number of elements
  // processed by the loop, so we will refer to that from this point on.
  if (!SE->isLoopInvariant(EC, L)) {
    LLVM_DEBUG(dbgs() << "ARM TP: element count must be loop invariant.\n");
    return nullptr;
  }

  // 2) Find out if IV is an induction phi. Note that we can't use Loop
  // helpers here to get the induction variable, because the hardware loop is
  // no longer in loopsimplify form, and also the hwloop intrinsic uses a
  // different counter. Using SCEV, we check that the induction is of the
  // form i = i + 4, where the increment must be equal to the VectorWidth.
  auto *IV = ActiveLaneMask->getOperand(0);
  auto *IVExpr = SE->getSCEV(IV);
  auto *AddExpr = dyn_cast<SCEVAddRecExpr>(IVExpr);

  if (!AddExpr) {
    LLVM_DEBUG(dbgs() << "ARM TP: induction not an add expr: "; IVExpr->dump());
    return nullptr;
  }
  // Check that this AddRec is associated with this loop.
  if (AddExpr->getLoop() != L) {
    LLVM_DEBUG(dbgs() << "ARM TP: phi not part of this loop\n");
    return nullptr;
  }
  auto *Step = dyn_cast<SCEVConstant>(AddExpr->getOperand(1));
  if (!Step) {
    LLVM_DEBUG(dbgs() << "ARM TP: induction step is not a constant: ";
               AddExpr->getOperand(1)->dump());
    return nullptr;
  }
  auto StepValue = Step->getValue()->getSExtValue();
  if (VectorWidth != StepValue) {
    LLVM_DEBUG(dbgs() << "ARM TP: Step value " << StepValue
                      << " doesn't match vector width " << VectorWidth << "\n");
    return nullptr;
  }

  if ((ConstElemCount = dyn_cast<ConstantInt>(ElemCount))) {
    ConstantInt *TC = dyn_cast<ConstantInt>(TripCount);
    if (!TC) {
      LLVM_DEBUG(dbgs() << "ARM TP: Constant tripcount expected in "
                           "set.loop.iterations\n");
      return nullptr;
    }

    // Calculate 2 tripcount values and check that they are consistent with
    // each other. The TripCount for a predicated vector loop body is
    // ceil(ElementCount/Width), or floor((ElementCount+Width-1)/Width) as we
    // work it out here.
    uint64_t TC1 = TC->getZExtValue();
    uint64_t TC2 =
        (ConstElemCount->getZExtValue() + VectorWidth - 1) / VectorWidth;

    // If the tripcount values are inconsistent, we can't insert the VCTP and
    // trigger tail-predication; keep the intrinsic as a get.active.lane.mask
    // and legalize this.
    if (TC1 != TC2) {
      LLVM_DEBUG(dbgs() << "ARM TP: inconsistent constant tripcount values: "
                 << TC1 << " from set.loop.iterations, and "
                 << TC2 << " from get.active.lane.mask\n");
      return nullptr;
    }
  } else if (!ForceTailPredication) {
    // 3) We need to prove that the sub expression that we create in the
    // tail-predicated loop body, which calculates the remaining elements to be
    // processed, is non-negative, i.e. it doesn't overflow:
    //
    //   ((ElementCount + VectorWidth - 1) / VectorWidth) - TripCount >= 0
    //
    // This is true if:
    //
    //    TripCount == (ElementCount + VectorWidth - 1) / VectorWidth
    //
    // which what we will be using here.
    //
    auto *VW = SE->getSCEV(ConstantInt::get(TripCount->getType(), VectorWidth));
    // ElementCount + (VW-1):
    auto *Start = AddExpr->getStart();
    auto *ECPlusVWMinus1 = SE->getAddExpr(EC,
        SE->getSCEV(ConstantInt::get(TripCount->getType(), VectorWidth - 1)));

    // Ceil = ElementCount + (VW-1) / VW
    auto *Ceil = SE->getUDivExpr(ECPlusVWMinus1, VW);

    // Prevent unused variable warnings with TC
    (void)TC;
    LLVM_DEBUG({
      dbgs() << "ARM TP: Analysing overflow behaviour for:\n";
      dbgs() << "ARM TP: - TripCount = " << *TC << "\n";
      dbgs() << "ARM TP: - ElemCount = " << *EC << "\n";
      dbgs() << "ARM TP: - Start = " << *Start << "\n";
      dbgs() << "ARM TP: - BETC = " << *SE->getBackedgeTakenCount(L) << "\n";
      dbgs() << "ARM TP: - VecWidth =  " << VectorWidth << "\n";
      dbgs() << "ARM TP: - (ElemCount+VW-1) / VW = " << *Ceil << "\n";
    });

    // As an example, almost all the tripcount expressions (produced by the
    // vectoriser) look like this:
    //
    //   TC = ((-4 + (4 * ((3 + %N) /u 4))<nuw> - start) /u 4)
    //
    // and "ElementCount + (VW-1) / VW":
    //
    //   Ceil = ((3 + %N) /u 4)
    //
    // Check for equality of TC and Ceil by calculating SCEV expression
    // TC - Ceil and test it for zero.
    //
    const SCEV *Div = SE->getUDivExpr(
        SE->getAddExpr(SE->getMulExpr(Ceil, VW), SE->getNegativeSCEV(VW),
                       SE->getNegativeSCEV(Start)),
        VW);
    const SCEV *Sub = SE->getMinusSCEV(SE->getBackedgeTakenCount(L), Div);
    LLVM_DEBUG(dbgs() << "ARM TP: - Sub       = "; Sub->dump());

    // Use context sensitive facts about the path to the loop to refine.  This
    // comes up as the backedge taken count can incorporate context sensitive
    // reasoning, and our RHS just above doesn't.
    Sub = SE->applyLoopGuards(Sub, L);
    LLVM_DEBUG(dbgs() << "ARM TP: - (Guarded) = "; Sub->dump());

    if (!Sub->isZero()) {
      LLVM_DEBUG(dbgs() << "ARM TP: possible overflow in sub expression.\n");
      return nullptr;
    }
  }

  // Check that the start value is a multiple of the VectorWidth.
  // TODO: This could do with a method to check if the scev is a multiple of
  // VectorWidth. For the moment we just check for constants, muls and unknowns
  // (which use MaskedValueIsZero and seems to be the most common).
  if (auto *BaseC = dyn_cast<SCEVConstant>(AddExpr->getStart())) {
    if (BaseC->getAPInt().urem(VectorWidth) == 0)
      return SE->getMinusSCEV(EC, BaseC);
  } else if (auto *BaseV = dyn_cast<SCEVUnknown>(AddExpr->getStart())) {
    Type *Ty = BaseV->getType();
    APInt Mask = APInt::getLowBitsSet(Ty->getPrimitiveSizeInBits(),
                                      Log2_64(VectorWidth));
    if (MaskedValueIsZero(BaseV->getValue(), Mask,
                          L->getHeader()->getModule()->getDataLayout()))
      return SE->getMinusSCEV(EC, BaseV);
  } else if (auto *BaseMul = dyn_cast<SCEVMulExpr>(AddExpr->getStart())) {
    if (auto *BaseC = dyn_cast<SCEVConstant>(BaseMul->getOperand(0)))
      if (BaseC->getAPInt().urem(VectorWidth) == 0)
        return SE->getMinusSCEV(EC, BaseC);
    if (auto *BaseC = dyn_cast<SCEVConstant>(BaseMul->getOperand(1)))
      if (BaseC->getAPInt().urem(VectorWidth) == 0)
        return SE->getMinusSCEV(EC, BaseC);
  }

  LLVM_DEBUG(
      dbgs() << "ARM TP: induction base is not know to be a multiple of VF: "
             << *AddExpr->getOperand(0) << "\n");
  return nullptr;
}

void MVETailPredication::InsertVCTPIntrinsic(IntrinsicInst *ActiveLaneMask,
                                             Value *Start) {
  IRBuilder<> Builder(L->getLoopPreheader()->getTerminator());
  Module *M = L->getHeader()->getModule();
  Type *Ty = IntegerType::get(M->getContext(), 32);
  unsigned VectorWidth =
      cast<FixedVectorType>(ActiveLaneMask->getType())->getNumElements();

  // Insert a phi to count the number of elements processed by the loop.
  Builder.SetInsertPoint(L->getHeader()->getFirstNonPHI());
  PHINode *Processed = Builder.CreatePHI(Ty, 2);
  Processed->addIncoming(Start, L->getLoopPreheader());

  // Replace @llvm.get.active.mask() with the ARM specific VCTP intrinic, and
  // thus represent the effect of tail predication.
  Builder.SetInsertPoint(ActiveLaneMask);
  ConstantInt *Factor = ConstantInt::get(cast<IntegerType>(Ty), VectorWidth);

  Intrinsic::ID VCTPID;
  switch (VectorWidth) {
  default:
    llvm_unreachable("unexpected number of lanes");
  case 2:  VCTPID = Intrinsic::arm_mve_vctp64; break;
  case 4:  VCTPID = Intrinsic::arm_mve_vctp32; break;
  case 8:  VCTPID = Intrinsic::arm_mve_vctp16; break;
  case 16: VCTPID = Intrinsic::arm_mve_vctp8; break;
  }
  Function *VCTP = Intrinsic::getDeclaration(M, VCTPID);
  Value *VCTPCall = Builder.CreateCall(VCTP, Processed);
  ActiveLaneMask->replaceAllUsesWith(VCTPCall);

  // Add the incoming value to the new phi.
  // TODO: This add likely already exists in the loop.
  Value *Remaining = Builder.CreateSub(Processed, Factor);
  Processed->addIncoming(Remaining, L->getLoopLatch());
  LLVM_DEBUG(dbgs() << "ARM TP: Insert processed elements phi: "
             << *Processed << "\n"
             << "ARM TP: Inserted VCTP: " << *VCTPCall << "\n");
}

bool MVETailPredication::TryConvertActiveLaneMask(Value *TripCount) {
  SmallVector<IntrinsicInst *, 4> ActiveLaneMasks;
  for (auto *BB : L->getBlocks())
    for (auto &I : *BB)
      if (auto *Int = dyn_cast<IntrinsicInst>(&I))
        if (Int->getIntrinsicID() == Intrinsic::get_active_lane_mask)
          ActiveLaneMasks.push_back(Int);

  if (ActiveLaneMasks.empty())
    return false;

  LLVM_DEBUG(dbgs() << "ARM TP: Found predicated vector loop.\n");

  for (auto *ActiveLaneMask : ActiveLaneMasks) {
    LLVM_DEBUG(dbgs() << "ARM TP: Found active lane mask: "
                      << *ActiveLaneMask << "\n");

    const SCEV *StartSCEV = IsSafeActiveMask(ActiveLaneMask, TripCount);
    if (!StartSCEV) {
      LLVM_DEBUG(dbgs() << "ARM TP: Not safe to insert VCTP.\n");
      return false;
    }
    LLVM_DEBUG(dbgs() << "ARM TP: Safe to insert VCTP. Start is " << *StartSCEV
                      << "\n");
    SCEVExpander Expander(*SE, L->getHeader()->getModule()->getDataLayout(),
                          "start");
    Instruction *Ins = L->getLoopPreheader()->getTerminator();
    Value *Start = Expander.expandCodeFor(StartSCEV, StartSCEV->getType(), Ins);
    LLVM_DEBUG(dbgs() << "ARM TP: Created start value " << *Start << "\n");
    InsertVCTPIntrinsic(ActiveLaneMask, Start);
  }

  // Remove dead instructions and now dead phis.
  for (auto *II : ActiveLaneMasks)
    RecursivelyDeleteTriviallyDeadInstructions(II);
  for (auto *I : L->blocks())
    DeleteDeadPHIs(I);
  return true;
}

Pass *llvm::createMVETailPredicationPass() {
  return new MVETailPredication();
}

char MVETailPredication::ID = 0;

INITIALIZE_PASS_BEGIN(MVETailPredication, DEBUG_TYPE, DESC, false, false)
INITIALIZE_PASS_END(MVETailPredication, DEBUG_TYPE, DESC, false, false)