llvm-project/llvm/lib/Target/ARM/MVETailPredication.cpp
David Green eb64450afc [ARM] Convert active.lane.masks to vctp with non-zero starts
This attempts to expand the logic in the MVETailPredication pass to convert
active lane masks that the vectorizer produces to vctp instructions that the
backend can later turn into tail predicated loops. Especially for addrecs with
non-zero starts that can be created from epilog vectorization. There is some
adjustment to the logic to handle this, moving some of the code to check the
addrec earlier so that we can get the start value. This start value is then
incorporated into the logic of checkin the new vctp is valid, and there is a
newly added check that it is known to be a multiple of the VF as we expect.

Differential Revision: https://reviews.llvm.org/D146517
2023-03-29 14:17:10 +01:00

468 lines
19 KiB
C++

//===- 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);
/// Rematerialize the iteration count in exit blocks, which enables
/// ARMLowOverheadLoops to better optimise away loop update statements inside
/// hardware-loops.
void RematerializeIterCount();
};
} // 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)