NVPTXPromoteParamAlign.cpp

Go to the documentation of this file.

//===-- NVPTXPromoteParamAlign.cpp - Promote .param alignment ------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Increase the .param-space alignment of NVPTX arguments and return values so
// their loads and stores can be vectorized. On every defined function:
//
//   1. Give each byval param an explicit ABI `align` for its pointee type
//      (capped at the PTX max). byval already implies this, but the alignment
//      is otherwise invisible to IR alignment analyses.
//   2. For a local function whose every use is a type-compatible direct call,
//      we control all call sites and raise aggregate/byval param and return
//      alignment to at least 16 (for 128-bit vectorization). This is recorded
//      as `stackalign` and mirrored onto the calls.
//   3. Propagate the result onto byval loads at a known offset, since `align`
//      and `stackalign` aren't both picked up by IR alignment analyses.
//
// (2) runs before (1) so byval `align` still matches between caller and callee
// while stackalign is mirrored onto the calls.
//
//===----------------------------------------------------------------------===//
 
#include "NVPTX.h"
#include "NVPTXUtilities.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/IR/Attributes.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/Module.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
#include "llvm/Support/Debug.h"
#include <optional>
#include <queue>
 
#define DEBUG_TYPE "nvptx-promote-param-align"
 
using namespace llvm;
 
namespace {
class NVPTXPromoteParamAlignLegacyPass : public ModulePass {
  bool runOnModule(Module &M) override;
 
public:
  static char ID;
  NVPTXPromoteParamAlignLegacyPass() : ModulePass(ID) {}
  StringRef getPassName() const override {
    return "Promote alignment of parameters and return values (NVPTX)";
  }
};
} // namespace
 
char NVPTXPromoteParamAlignLegacyPass::ID = 0;
 
INITIALIZE_PASS(NVPTXPromoteParamAlignLegacyPass, "nvptx-promote-param-align",
                "Promote alignment of parameters and return values (NVPTX)",
                false, false)
 
// Return true if the attributes that determine an NVPTX .param slot's layout
// match.
static bool layoutAttrsMatch(AttributeSet CalleeAttrs, AttributeSet CallAttrs) {
  if (CalleeAttrs.getByValType() != CallAttrs.getByValType() ||
      CalleeAttrs.getStackAlignment() != CallAttrs.getStackAlignment())
    return false;
 
  // `align` only affects the layout for byval parameters.
  return !CalleeAttrs.getByValType() ||
         CalleeAttrs.getAlignment() == CallAttrs.getAlignment();
}
 
static bool callSiteMatchesCalleeABI(const CallBase &CB, const Function &F) {
  const AttributeList CalleeAttrs = F.getAttributes();
  const AttributeList CallAttrs = CB.getAttributes();
 
  if (!layoutAttrsMatch(CalleeAttrs.getRetAttrs(), CallAttrs.getRetAttrs()))
    return false;
 
  return all_of(seq(F.arg_size()), [&](size_t I) {
    return layoutAttrsMatch(CalleeAttrs.getParamAttrs(I),
                            CallAttrs.getParamAttrs(I));
  });
}
 
// Promotable if the function is local and every use is an ABI-compatible direct
// call, so we control every call site and can raise alignment on both sides.
static bool canPromoteParamAlign(Function &F) {
  if (F.isDeclaration() || !F.hasLocalLinkage())
    return false;
 
  if (F.hasAddressTaken(/*Users=*/nullptr, /*IgnoreCallbackUses=*/false,
                        /*IgnoreAssumeLikeCalls=*/true,
                        /*IgnoreLLVMUsed=*/true))
    return false;
 
  return all_of(F.users(), [&](const User *U) {
    const auto *CB = dyn_cast<CallBase>(U);
    if (!CB || CB->getCalledOperand() != &F)
      return true;
    return CB->getFunctionType() == F.getFunctionType() &&
           callSiteMatchesCalleeABI(*CB, F);
  });
}
 
// Raise the alignment of every load reachable from a byval pointer at a known
// constant offset. Must stay in sync with the param load/store in LowerCall.
static bool propagateAlignmentToLoads(Value *Val, Align NewAlign,
                                      const DataLayout &DL) {
  struct Load {
    LoadInst *Inst;
    uint64_t Offset;
  };
 
  struct LoadContext {
    Value *InitialVal;
    uint64_t Offset;
  };
 
  SmallVector<Load> Loads;
  std::queue<LoadContext> Worklist;
  Worklist.push({Val, 0});
 
  while (!Worklist.empty()) {
    LoadContext Ctx = Worklist.front();
    Worklist.pop();
 
    for (User *CurUser : Ctx.InitialVal->users()) {
      if (auto *I = dyn_cast<LoadInst>(CurUser))
        Loads.push_back({I, Ctx.Offset});
      else if (isa<BitCastInst>(CurUser) || isa<AddrSpaceCastInst>(CurUser))
        Worklist.push({cast<Instruction>(CurUser), Ctx.Offset});
      else if (auto *I = dyn_cast<GetElementPtrInst>(CurUser)) {
        APInt OffsetAccumulated =
            APInt::getZero(DL.getIndexTypeSizeInBits(I->getType()));
 
        if (!I->accumulateConstantOffset(DL, OffsetAccumulated))
          continue;
 
        uint64_t OffsetLimit = -1;
        uint64_t Offset = OffsetAccumulated.getLimitedValue(OffsetLimit);
        assert(Offset != OffsetLimit && "Expect Offset less than UINT64_MAX");
 
        Worklist.push({I, Ctx.Offset + Offset});
      }
    }
  }
 
  bool Changed = false;
  for (Load &CurLoad : Loads) {
    Align NewLoadAlign = commonAlignment(NewAlign, CurLoad.Offset);
    if (NewLoadAlign > CurLoad.Inst->getAlign()) {
      CurLoad.Inst->setAlignment(NewLoadAlign);
      Changed = true;
    }
  }
  return Changed;
}
 
// Bump an alignment up to at least 16 (for 128-bit vectorization), or nullopt
// if it's already large enough.
static MaybeAlign getPromotedParamAlign(Align CurrentAlign) {
  const Align PromotedAlign = std::max(CurrentAlign, Align(16));
  if (PromotedAlign > CurrentAlign)
    return PromotedAlign;
  return std::nullopt;
}
 
static bool promoteParamAlign(Function &F) {
  if (!canPromoteParamAlign(F))
    return false;
 
  LLVMContext &Ctx = F.getContext();
  const DataLayout &DL = F.getDataLayout();
 
  // Promoted (arg index, new alignment) pairs, to mirror onto call sites.
  SmallVector<std::pair<unsigned, Align>, 8> PromotedParams;
  MaybeAlign PromotedRet;
 
  // Promote aggregate and byval parameters.
  for (Argument &Arg : F.args()) {
    const bool IsByVal = Arg.hasByValAttr();
    Type *ArgTy = IsByVal ? Arg.getParamByValType() : Arg.getType();
    if (ArgTy->isEmptyTy() || (!IsByVal && !shouldPassAsArray(ArgTy)))
      continue;
 
    // An explicit stackalign already wins at emission time, nothing to promote.
    if (Arg.getParamStackAlign())
      continue;
    const unsigned ArgNo = Arg.getArgNo();
 
    // `align` only applies to byval (pointer) args, not by-value aggregates.
    Align CurrentAlign = getPTXParamTypeAlign(ArgTy, DL);
    if (IsByVal)
      CurrentAlign = std::max(CurrentAlign, Arg.getParamAlign().valueOrOne());
    const MaybeAlign PromotedAlign = getPromotedParamAlign(CurrentAlign);
    if (!PromotedAlign)
      continue;
 
    LLVM_DEBUG(dbgs() << "Promoting alignment of " << Arg << " to "
                      << PromotedAlign->value() << '\n');
    Arg.addAttr(Attribute::getWithStackAlignment(Ctx, *PromotedAlign));
    PromotedParams.emplace_back(ArgNo, *PromotedAlign);
  }
 
  // Promote an aggregate return value.
  Type *RetTy = F.getReturnType();
  if (shouldPassAsArray(RetTy) && !RetTy->isEmptyTy() &&
      !F.getAttributes().getRetStackAlignment()) {
    const MaybeAlign PromotedAlign =
        getPromotedParamAlign(getPTXParamTypeAlign(RetTy, DL));
    if (PromotedAlign) {
      F.addRetAttr(Attribute::getWithStackAlignment(Ctx, *PromotedAlign));
      PromotedRet = *PromotedAlign;
    }
  }
 
  if (PromotedParams.empty() && !PromotedRet)
    return false;
 
  // Mirror the promotion onto every direct call site so both sides agree on the
  // .param layout. canPromoteParamAlign already verified they're
  // ABI-compatible.
  for (User *U : F.users()) {
    auto *CB = dyn_cast<CallBase>(U);
    if (!CB || CB->getCalledOperand() != &F)
      continue;
 
    for (const auto &[ArgNo, PromotedAlign] : PromotedParams)
      CB->addParamAttr(ArgNo,
                       Attribute::getWithStackAlignment(Ctx, PromotedAlign));
    if (PromotedRet)
      CB->addRetAttr(Attribute::getWithStackAlignment(Ctx, *PromotedRet));
 
    assert(callSiteMatchesCalleeABI(*CB, F) &&
           "mirroring must preserve call-site/callee ABI compatibility");
  }
 
  return true;
}
 
// Spell out each byval parameter's ABI alignment as an explicit `align` (step 1
// above). Runs after promoteParamAlign, which needs byval `align` to still
// match between callers and callees.
static bool setByValParamABIAlign(Function &F) {
  if (F.isDeclaration())
    return false;
 
  LLVMContext &Ctx = F.getContext();
  const DataLayout &DL = F.getDataLayout();
  bool Changed = false;
  for (Argument &Arg : F.args()) {
    if (!Arg.hasByValAttr())
      continue;
    Type *ETy = Arg.getParamByValType();
    if (ETy->isEmptyTy())
      continue;
    const Align ABIAlign = getPTXParamTypeAlign(ETy, DL);
    if (Arg.getParamAlign().valueOrOne() >= ABIAlign)
      continue;
    Arg.removeAttr(Attribute::Alignment);
    Arg.addAttr(Attribute::getWithAlignment(Ctx, ABIAlign));
    Changed = true;
  }
  return Changed;
}
 
// Propagate each byval parameter's .param alignment onto its constant-offset
// loads (step 3 above). Runs after promoteParamAlign so the promoted
// `stackalign` is included, and on every function so kernels and external
// functions benefit too.
static bool propagateByValParamLoadAlign(Function &F) {
  if (F.isDeclaration())
    return false;
 
  const DataLayout &DL = F.getDataLayout();
  bool Changed = false;
  for (Argument &Arg : F.args()) {
    if (!Arg.hasByValAttr())
      continue;
    Type *ETy = Arg.getParamByValType();
    if (ETy->isEmptyTy())
      continue;
    const unsigned ParamIdx = Arg.getArgNo() + AttributeList::FirstArgIndex;
    const Align ParamAlign = getDeviceByValParamAlign(&F, ETy, ParamIdx, DL);
    Changed |= propagateAlignmentToLoads(&Arg, ParamAlign, DL);
  }
  return Changed;
}
 
static bool promoteParamAlignModule(Module &M) {
  bool Changed = false;
  for (Function &F : M) {
    // Order matters (see the file header): promote, normalize `align`, then
    // propagate to loads.
    Changed |= promoteParamAlign(F);
    Changed |= setByValParamABIAlign(F);
    Changed |= propagateByValParamLoadAlign(F);
  }
  return Changed;
}
 
bool NVPTXPromoteParamAlignLegacyPass::runOnModule(Module &M) {
  return promoteParamAlignModule(M);
}
 
ModulePass *llvm::createNVPTXPromoteParamAlignPass() {
  return new NVPTXPromoteParamAlignLegacyPass();
}
 
PreservedAnalyses NVPTXPromoteParamAlignPass::run(Module &M,
                                                  ModuleAnalysisManager &AM) {
  return promoteParamAlignModule(M) ? PreservedAnalyses::none()
                                    : PreservedAnalyses::all();
}