[llvm] [NVPTX] Promote internal function alignments in IR pass (PR #208040)
Alex MacLean via llvm-commits
llvm-commits at lists.llvm.org
Tue Jul 7 22:27:54 PDT 2026
https://github.com/AlexMaclean updated https://github.com/llvm/llvm-project/pull/208040
>From 6cbd3afe71114f108340e1c90d9e1ef299a8a526 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Tue, 7 Jul 2026 09:09:30 -0700
Subject: [PATCH 1/2] [NVPTX] Promote interal function alignments in IR pass
---
llvm/lib/Target/NVPTX/CMakeLists.txt | 2 +-
llvm/lib/Target/NVPTX/NVPTX.h | 10 +-
llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp | 22 +-
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 14 +-
llvm/lib/Target/NVPTX/NVPTXPassRegistry.def | 2 +-
.../Target/NVPTX/NVPTXPromoteParamAlign.cpp | 317 ++++++++++++++++++
.../Target/NVPTX/NVPTXSetByValParamAlign.cpp | 162 ---------
llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp | 4 +-
llvm/lib/Target/NVPTX/NVPTXUtilities.cpp | 57 ++--
llvm/lib/Target/NVPTX/NVPTXUtilities.h | 36 +-
.../test/CodeGen/NVPTX/promote-param-align.ll | 217 ++++++++++++
.../CodeGen/NVPTX/set-byval-param-align.ll | 120 -------
.../secondary/llvm/lib/Target/NVPTX/BUILD.gn | 2 +-
13 files changed, 597 insertions(+), 368 deletions(-)
create mode 100644 llvm/lib/Target/NVPTX/NVPTXPromoteParamAlign.cpp
delete mode 100644 llvm/lib/Target/NVPTX/NVPTXSetByValParamAlign.cpp
create mode 100644 llvm/test/CodeGen/NVPTX/promote-param-align.ll
delete mode 100644 llvm/test/CodeGen/NVPTX/set-byval-param-align.ll
diff --git a/llvm/lib/Target/NVPTX/CMakeLists.txt b/llvm/lib/Target/NVPTX/CMakeLists.txt
index d2b5083a5f83f..00f19eee21704 100644
--- a/llvm/lib/Target/NVPTX/CMakeLists.txt
+++ b/llvm/lib/Target/NVPTX/CMakeLists.txt
@@ -35,8 +35,8 @@ set(NVPTXCodeGen_sources
NVPTXMCExpr.cpp
NVPTXMarkKernelPtrsGlobal.cpp
NVPTXPeephole.cpp
- NVPTXSetByValParamAlign.cpp
NVPTXPrologEpilogPass.cpp
+ NVPTXPromoteParamAlign.cpp
NVPTXProxyRegErasure.cpp
NVPTXRegisterInfo.cpp
NVPTXReplaceImageHandles.cpp
diff --git a/llvm/lib/Target/NVPTX/NVPTX.h b/llvm/lib/Target/NVPTX/NVPTX.h
index 57ed9bb43aeea..914295bf76e93 100644
--- a/llvm/lib/Target/NVPTX/NVPTX.h
+++ b/llvm/lib/Target/NVPTX/NVPTX.h
@@ -49,7 +49,7 @@ MachineFunctionPass *createNVPTXPrologEpilogPass();
MachineFunctionPass *createNVPTXReplaceImageHandlesPass();
FunctionPass *createNVPTXImageOptimizerPass();
ModulePass *createNVPTXLowerArgsPass();
-FunctionPass *createNVPTXSetByValParamAlignPass();
+ModulePass *createNVPTXPromoteParamAlignPass();
FunctionPass *createNVPTXLowerAllocaPass();
FunctionPass *createNVPTXLowerUnreachablePass(bool TrapUnreachable,
bool NoTrapAfterNoreturn);
@@ -72,7 +72,7 @@ void initializeNVPTXLowerAggrCopiesPass(PassRegistry &);
void initializeNVPTXLowerAllocaPass(PassRegistry &);
void initializeNVPTXLowerUnreachablePass(PassRegistry &);
void initializeNVPTXLowerArgsLegacyPassPass(PassRegistry &);
-void initializeNVPTXSetByValParamAlignLegacyPassPass(PassRegistry &);
+void initializeNVPTXPromoteParamAlignLegacyPassPass(PassRegistry &);
void initializeNVPTXProxyRegErasurePass(PassRegistry &);
void initializeNVPTXForwardParamsPassPass(PassRegistry &);
void initializeNVPTXAddressFolderPassPass(PassRegistry &);
@@ -111,9 +111,9 @@ struct NVPTXCopyByValArgsPass : OptionalPassInfoMixin<NVPTXCopyByValArgsPass> {
PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
};
-struct NVPTXSetByValParamAlignPass
- : OptionalPassInfoMixin<NVPTXSetByValParamAlignPass> {
- PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
+struct NVPTXPromoteParamAlignPass
+ : OptionalPassInfoMixin<NVPTXPromoteParamAlignPass> {
+ PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
};
struct NVPTXLowerArgsPass : OptionalPassInfoMixin<NVPTXLowerArgsPass> {
diff --git a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
index 2f4964fe33042..a22e2cd05a931 100644
--- a/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXAsmPrinter.cpp
@@ -364,19 +364,9 @@ void NVPTXAsmPrinter::emitCallPrototype(const CallBase &CB,
Type *Ty = CB.getArgOperand(I)->getType();
if (CB.paramHasAttr(I, Attribute::ByVal)) {
- // Indirect calls need strict ABI alignment so we disable optimizations by
- // not providing a function to optimize.
Type *ETy = CB.getParamByValType(I);
- // Mirror the byval alignment computed by SelectionDAGBuilder: prefer an
- // explicit stack/param alignment, otherwise fall back to the byval type
- // alignment.
- MaybeAlign InitialAlign = CB.getParamStackAlign(I);
- if (!InitialAlign)
- InitialAlign = CB.getParamAlign(I);
- Align ByValAlign =
- InitialAlign.value_or(TLI->getByValTypeAlignment(ETy, DL));
- Align ParamByValAlign =
- getDeviceByValParamAlign(/*F=*/nullptr, ETy, ByValAlign, DL);
+ Align ParamByValAlign = getDeviceByValParamAlign(
+ &CB, ETy, I + AttributeList::FirstArgIndex, DL);
O << ".param .align " << ParamByValAlign.value() << " .b8 _["
<< DL.getTypeAllocSize(ETy) << "]";
@@ -1532,12 +1522,10 @@ void NVPTXAsmPrinter::emitFunctionParamList(const Function *F, raw_ostream &O) {
// <a> = optimal alignment for the element type; always multiple of
// PAL.getParamAlignment
// size = typeallocsize of element type
+ const unsigned ParamIdx = Arg.getArgNo() + AttributeList::FirstArgIndex;
const Align OptimalAlign =
- IsKernelFunc
- ? getPTXParamAlign(
- F, ETy, Arg.getArgNo() + AttributeList::FirstArgIndex, DL)
- : getDeviceByValParamAlign(F, ETy,
- Arg.getParamAlign().valueOrOne(), DL);
+ IsKernelFunc ? getPTXParamAlign(F, ETy, ParamIdx, DL)
+ : getDeviceByValParamAlign(F, ETy, ParamIdx, DL);
O << "\t.param .align " << OptimalAlign.value() << " .b8 " << ParamSym
<< "[" << DL.getTypeAllocSize(ETy) << "]";
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index fa0b9fe5f7051..82e4886ab7ebe 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -1370,16 +1370,10 @@ SDValue NVPTXTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
Type *ETy = (IsByVal ? Arg.IndirectType : Arg.Ty);
const Align ArgAlign = [&]() {
- if (IsByVal) {
- // The ByValAlign in the Outs[OIdx].Flags is always set at this point,
- // so we don't need to worry whether it's naturally aligned or not.
- // See TargetLowering::LowerCallTo().
- const Align InitialAlign = ArgOuts[0].Flags.getNonZeroByValAlign();
- return getDeviceByValParamAlign(CB->getCalledFunction(), ETy,
- InitialAlign, DL);
- }
- return getPTXParamAlign(CB, Arg.Ty, ArgI + AttributeList::FirstArgIndex,
- DL);
+ const unsigned ParamIdx = ArgI + AttributeList::FirstArgIndex;
+ if (IsByVal)
+ return getDeviceByValParamAlign(CB, ETy, ParamIdx, DL);
+ return getPTXParamAlign(CB, Arg.Ty, ParamIdx, DL);
}();
const unsigned TySize = DL.getTypeAllocSize(ETy);
diff --git a/llvm/lib/Target/NVPTX/NVPTXPassRegistry.def b/llvm/lib/Target/NVPTX/NVPTXPassRegistry.def
index 70d4e92a32e2f..7fd6265cda106 100644
--- a/llvm/lib/Target/NVPTX/NVPTXPassRegistry.def
+++ b/llvm/lib/Target/NVPTX/NVPTXPassRegistry.def
@@ -19,6 +19,7 @@
MODULE_PASS("generic-to-nvvm", GenericToNVVMPass())
MODULE_PASS("nvptx-lower-args", NVPTXLowerArgsPass(*this))
MODULE_PASS("nvptx-lower-ctor-dtor", NVPTXCtorDtorLoweringPass())
+MODULE_PASS("nvptx-promote-param-align", NVPTXPromoteParamAlignPass())
MODULE_PASS("nvvm-reflect", NVVMReflectPass())
#undef MODULE_PASS
@@ -39,7 +40,6 @@ FUNCTION_ALIAS_ANALYSIS("nvptx-aa", NVPTXAA())
#endif
FUNCTION_PASS("nvvm-intr-range", NVVMIntrRangePass())
FUNCTION_PASS("nvptx-copy-byval-args", NVPTXCopyByValArgsPass())
-FUNCTION_PASS("nvptx-set-byval-param-align", NVPTXSetByValParamAlignPass())
FUNCTION_PASS("nvptx-mark-kernel-ptrs-global", NVPTXMarkKernelPtrsGlobalPass())
FUNCTION_PASS("nvptx-tag-invariant-loads", NVPTXTagInvariantLoadsPass())
FUNCTION_PASS("nvptx-ir-peephole", NVPTXIRPeepholePass())
diff --git a/llvm/lib/Target/NVPTX/NVPTXPromoteParamAlign.cpp b/llvm/lib/Target/NVPTX/NVPTXPromoteParamAlign.cpp
new file mode 100644
index 0000000000000..e1df83d9994b0
--- /dev/null
+++ b/llvm/lib/Target/NVPTX/NVPTXPromoteParamAlign.cpp
@@ -0,0 +1,317 @@
+//===-- 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();
+}
diff --git a/llvm/lib/Target/NVPTX/NVPTXSetByValParamAlign.cpp b/llvm/lib/Target/NVPTX/NVPTXSetByValParamAlign.cpp
deleted file mode 100644
index bd5cbfb4ac6b6..0000000000000
--- a/llvm/lib/Target/NVPTX/NVPTXSetByValParamAlign.cpp
+++ /dev/null
@@ -1,162 +0,0 @@
-//===-- NVPTXSetByValParamAlign.cpp - Set byval 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Set explicit alignment on byval parameter attributes in the NVPTX backend.
-// Without this, the alignment is left unspecified and IR-level analyses (e.g.,
-// computeKnownBits via Value::getPointerAlignment) conservatively assume
-// Align(1), since the actual alignment is a target-specific codegen detail not
-// visible at the IR level.
-//
-// The alignment is chosen as follows:
-// - Externally-visible functions: ABI type alignment (capped at 128).
-// - Internal/private functions: max(16, ABI align) to enable 128-bit
-// vectorized param loads. The compiler can _increase_ alignment beyond ABI
-// in this case because it has control over all of the call sites and byval
-// parameters are copies allocated by the caller in .param space.
-//
-// After updating the attribute, the pass propagates the improved alignment to
-// all loads from the byval pointer that use a known constant offset.
-//
-// TODO: Consider removing the load propagation in favor of infer-alignment,
-// which should be able to pick up the improved alignment from the attribute.
-//
-//===----------------------------------------------------------------------===//
-
-#include "NVPTX.h"
-#include "NVPTXUtilities.h"
-#include "llvm/IR/Function.h"
-#include "llvm/IR/Instructions.h"
-#include "llvm/InitializePasses.h"
-#include "llvm/Pass.h"
-#include "llvm/Support/Debug.h"
-#include <queue>
-
-#define DEBUG_TYPE "nvptx-set-byval-param-align"
-
-using namespace llvm;
-
-namespace {
-class NVPTXSetByValParamAlignLegacyPass : public FunctionPass {
- bool runOnFunction(Function &F) override;
-
-public:
- static char ID;
- NVPTXSetByValParamAlignLegacyPass() : FunctionPass(ID) {}
- StringRef getPassName() const override {
- return "Set alignment of byval parameters (NVPTX)";
- }
-};
-} // namespace
-
-char NVPTXSetByValParamAlignLegacyPass::ID = 0;
-
-INITIALIZE_PASS(NVPTXSetByValParamAlignLegacyPass,
- "nvptx-set-byval-param-align",
- "Set alignment of byval parameters (NVPTX)", false, false)
-
-static Align setByValParamAlign(Argument *Arg) {
- Function *F = Arg->getParent();
- Type *ByValType = Arg->getParamByValType();
- const DataLayout &DL = F->getDataLayout();
-
- const Align OptimizedAlign = getPTXPromotedParamTypeAlign(F, ByValType, DL);
- const Align CurrentAlign = Arg->getParamAlign().valueOrOne();
-
- if (CurrentAlign >= OptimizedAlign)
- return CurrentAlign;
-
- LLVM_DEBUG(dbgs() << "Try to use alignment " << OptimizedAlign.value()
- << " instead of " << CurrentAlign.value() << " for " << *Arg
- << '\n');
-
- Arg->removeAttr(Attribute::Alignment);
- Arg->addAttr(Attribute::getWithAlignment(F->getContext(), OptimizedAlign));
-
- return OptimizedAlign;
-}
-
-// Adjust alignment of arguments passed byval in .param address space. We can
-// increase alignment of such arguments in a way that ensures that we can
-// effectively vectorize their loads. We should also traverse all loads from
-// byval pointer and adjust their alignment, if those were using known offset.
-// Such alignment changes must be conformed with parameter store and load in
-// NVPTXTargetLowering::LowerCall.
-static void 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});
- }
- }
- }
-
- for (Load &CurLoad : Loads) {
- Align NewLoadAlign = commonAlignment(NewAlign, CurLoad.Offset);
- Align CurLoadAlign = CurLoad.Inst->getAlign();
- CurLoad.Inst->setAlignment(std::max(NewLoadAlign, CurLoadAlign));
- }
-}
-
-static bool setByValParamAlignment(Function &F) {
- const DataLayout &DL = F.getDataLayout();
- bool Changed = false;
- for (Argument &Arg : F.args()) {
- if (!Arg.hasByValAttr())
- continue;
- const Align NewArgAlign = setByValParamAlign(&Arg);
- propagateAlignmentToLoads(&Arg, NewArgAlign, DL);
- Changed = true;
- }
- return Changed;
-}
-
-bool NVPTXSetByValParamAlignLegacyPass::runOnFunction(Function &F) {
- return setByValParamAlignment(F);
-}
-
-FunctionPass *llvm::createNVPTXSetByValParamAlignPass() {
- return new NVPTXSetByValParamAlignLegacyPass();
-}
-
-PreservedAnalyses
-NVPTXSetByValParamAlignPass::run(Function &F, FunctionAnalysisManager &AM) {
- return setByValParamAlignment(F) ? PreservedAnalyses::none()
- : PreservedAnalyses::all();
-}
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index 2bc2158fccae1..fc954c3397caa 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -111,7 +111,7 @@ extern "C" LLVM_ABI LLVM_EXTERNAL_VISIBILITY void LLVMInitializeNVPTXTarget() {
initializeNVPTXAssignValidGlobalNamesPass(PR);
initializeNVPTXAtomicLowerPass(PR);
initializeNVPTXLowerArgsLegacyPassPass(PR);
- initializeNVPTXSetByValParamAlignLegacyPassPass(PR);
+ initializeNVPTXPromoteParamAlignLegacyPassPass(PR);
initializeNVPTXMarkKernelPtrsGlobalLegacyPassPass(PR);
initializeNVPTXLowerAllocaPass(PR);
initializeNVPTXLowerUnreachablePass(PR);
@@ -367,7 +367,7 @@ void NVPTXPassConfig::addIRPasses() {
// before the address space inference passes.
if (getNVPTXTargetMachine().getDrvInterface() == NVPTX::CUDA)
addPass(createNVPTXMarkKernelPtrsGlobalPass());
- addPass(createNVPTXSetByValParamAlignPass());
+ addPass(createNVPTXPromoteParamAlignPass());
addPass(createNVPTXLowerArgsPass());
if (getOptLevel() != CodeGenOptLevel::None) {
addAddressSpaceInferencePasses();
diff --git a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
index b617c379d5f5e..2b636f0ac8cad 100644
--- a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
@@ -33,31 +33,13 @@ Function *llvm::getMaybeBitcastedCallee(const CallBase *CB) {
return dyn_cast<Function>(CB->getCalledOperand()->stripPointerCasts());
}
-Align llvm::getPTXPromotedParamTypeAlign(const Function *F, Type *ArgTy,
- const DataLayout &DL) {
+Align llvm::getPTXParamTypeAlign(Type *ArgTy, const DataLayout &DL) {
// Capping the alignment to 128 bytes as that is the maximum alignment
// supported by PTX.
- const Align ABITypeAlign = std::min(Align(128), DL.getABITypeAlign(ArgTy));
-
- // If a function has linkage different from internal or private, we
- // must use default ABI alignment as external users rely on it. Same
- // for a function that may be called from a function pointer.
- const bool MayOptimizeAlign =
- F && F->hasLocalLinkage() &&
- !F->hasAddressTaken(/*Users=*/nullptr,
- /*IgnoreCallbackUses=*/false,
- /*IgnoreAssumeLikeCalls=*/true,
- /*IgnoreLLVMUsed=*/true);
- assert(!(MayOptimizeAlign && isKernelFunction(*F)) &&
- "Expect kernels to have non-local linkage");
- const Align OptimizedAlign = MayOptimizeAlign ? Align(16) : Align(1);
- return std::max(OptimizedAlign, ABITypeAlign);
+ return std::min(Align(128), DL.getABITypeAlign(ArgTy));
}
-Align llvm::getDeviceByValParamAlign(const Function *F, Type *ArgTy,
- Align InitialAlign, const DataLayout &DL) {
- const Align OptimizedAlign = getPTXPromotedParamTypeAlign(F, ArgTy, DL);
-
+static Align getByValParamAlignFloor(const Function *F) {
// Old ptx versions have a bug. When PTX code takes address of
// byval parameter with alignment < 4, ptxas generates code to
// spill argument into memory. Alas on sm_50+ ptxas generates
@@ -69,9 +51,28 @@ Align llvm::getDeviceByValParamAlign(const Function *F, Type *ArgTy,
// on non-deprecated ptxas versions.
const bool ShouldForceMinAlign =
ForceMinByValParamAlign && (!F || !isKernelFunction(*F));
- const Align AlignFloor = ShouldForceMinAlign ? Align(4) : Align(1);
+ return ShouldForceMinAlign ? Align(4) : Align(1);
+}
+
+Align llvm::getDeviceByValParamAlign(const Function *F, Type *ArgTy,
+ unsigned AttrIdx, const DataLayout &DL) {
+ return std::max(getPTXParamAlign(F, ArgTy, AttrIdx, DL),
+ getByValParamAlignFloor(F));
+}
- return std::max({InitialAlign, OptimizedAlign, AlignFloor});
+Align getDeviceByValParamAlign(const CallBase *CB, Type *ArgTy,
+ unsigned AttrIdx, const DataLayout &DL) {
+ Align ParamAlign = getPTXParamAlign(CB, ArgTy, AttrIdx, DL);
+
+ // For an indirect call getPTXParamAlign can't see the call's own byval
+ // alignment, so fold it in.
+ if (CB && AttrIdx >= AttributeList::FirstArgIndex)
+ ParamAlign = std::max(
+ ParamAlign,
+ CB->getParamAlign(AttrIdx - AttributeList::FirstArgIndex).valueOrOne());
+
+ return std::max(ParamAlign, getByValParamAlignFloor(
+ CB ? CB->getCalledFunction() : nullptr));
}
Align llvm::getPTXParamAlign(const Function *F, Type *Ty, unsigned AttrIdx,
@@ -80,7 +81,7 @@ Align llvm::getPTXParamAlign(const Function *F, Type *Ty, unsigned AttrIdx,
if (MaybeAlign StackAlign = getStackAlign(*F, AttrIdx))
return StackAlign.value();
- Align TypeAlign = getPTXPromotedParamTypeAlign(F, Ty, DL);
+ Align TypeAlign = getPTXParamTypeAlign(Ty, DL);
if (F && AttrIdx >= AttributeList::FirstArgIndex) {
unsigned ArgNo = AttrIdx - AttributeList::FirstArgIndex;
if (F->getAttributes().hasParamAttr(ArgNo, Attribute::ByVal))
@@ -91,14 +92,14 @@ Align llvm::getPTXParamAlign(const Function *F, Type *Ty, unsigned AttrIdx,
Align llvm::getPTXParamAlign(const CallBase *CB, Type *Ty, unsigned Idx,
const DataLayout &DL) {
- const Function *DirectCallee = CB ? CB->getCalledFunction() : nullptr;
-
- if (!DirectCallee && CB) {
+ if (CB)
if (MaybeAlign StackAlign = getStackAlign(*CB, Idx))
return StackAlign.value();
+ // Otherwise resolve the direct callee and use its parameter alignment.
+ const Function *DirectCallee = CB ? CB->getCalledFunction() : nullptr;
+ if (!DirectCallee && CB)
DirectCallee = getMaybeBitcastedCallee(CB);
- }
return getPTXParamAlign(DirectCallee, Ty, Idx, DL);
}
diff --git a/llvm/lib/Target/NVPTX/NVPTXUtilities.h b/llvm/lib/Target/NVPTX/NVPTXUtilities.h
index 6785883e5af1e..66ea690631a76 100644
--- a/llvm/lib/Target/NVPTX/NVPTXUtilities.h
+++ b/llvm/lib/Target/NVPTX/NVPTXUtilities.h
@@ -32,30 +32,24 @@ class TargetMachine;
Function *getMaybeBitcastedCallee(const CallBase *CB);
-/// Since function arguments are passed via .param space, we may want to
-/// increase their alignment in a way that ensures that we can effectively
-/// vectorize their loads & stores. We can increase alignment only if the
-/// function has internal or private linkage as for other linkage types callers
-/// may already rely on default alignment. To allow using 128-bit vectorized
-/// loads/stores, this function ensures that alignment is 16 or greater.
-Align getPTXPromotedParamTypeAlign(const Function *F, Type *ArgTy,
- const DataLayout &DL);
-
-Align getDeviceByValParamAlign(const Function *F, Type *ArgTy,
- Align InitialAlign, const DataLayout &DL);
-
-/// Get the alignment for a function parameter or return value.
-/// \p AttrIdx is the AttributeList index (e.g. FirstArgIndex + argNo, or
-/// ReturnIndex for return values). Checks for an explicit alignment attribute,
-/// then falls back to getPromotedParamTypeAlign, incorporating byval param
-/// alignment when applicable.
+/// ABI alignment of \p ArgTy in .param space, capped at the PTX maximum of 128.
+Align getPTXParamTypeAlign(Type *ArgTy, const DataLayout &DL);
+
+/// The .param-space alignment for a byval parameter or call argument: the
+/// (possibly promoted) parameter alignment, raised to the ptxas byval minimum.
+Align getDeviceByValParamAlign(const Function *F, Type *ArgTy, unsigned AttrIdx,
+ const DataLayout &DL);
+Align getDeviceByValParamAlign(const CallBase *CB, Type *ArgTy,
+ unsigned AttrIdx, const DataLayout &DL);
+
+/// Alignment for a function parameter or return value at AttributeList index
+/// \p AttrIdx (FirstArgIndex + argNo, or ReturnIndex). Prefers an explicit
+/// stackalign, else the ABI type alignment, folding in the byval `align`.
Align getPTXParamAlign(const Function *F, Type *Ty, unsigned AttrIdx,
const DataLayout &DL);
-/// Get the alignment for a call-site argument or return value. Resolves the
-/// callee and delegates to the Function overload of getParamAlign. For
-/// indirect calls with no resolvable callee, falls back to
-/// getPromotedParamTypeAlign.
+/// Alignment for a call-site argument or return value. Prefers an explicit
+/// stackalign on the call, else resolves the direct callee.
Align getPTXParamAlign(const CallBase *CB, Type *Ty, unsigned AttrIdx,
const DataLayout &DL);
diff --git a/llvm/test/CodeGen/NVPTX/promote-param-align.ll b/llvm/test/CodeGen/NVPTX/promote-param-align.ll
new file mode 100644
index 0000000000000..a98ffa07f287f
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/promote-param-align.ll
@@ -0,0 +1,217 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=nvptx-promote-param-align -S | FileCheck %s
+
+target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
+target triple = "nvptx64-nvidia-cuda"
+
+%struct.S = type { i32, i32 }
+
+ at gfp = internal global ptr null
+
+; Kernels aren't promoted, but byval params still get an explicit ABI align (4).
+define ptx_kernel void @kernel_byval(ptr byval(%struct.S) align 1 %s) {
+; CHECK-LABEL: define ptx_kernel void @kernel_byval(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; External functions aren't promoted, but the byval param still gets ABI align.
+define void @extern_byval(ptr byval(%struct.S) %s) {
+; CHECK-LABEL: define void @extern_byval(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Internal byval: explicit ABI align (4) plus alignstack(16) promotion.
+define internal void @internal_byval(ptr byval(%struct.S) %s) {
+; CHECK-LABEL: define internal void @internal_byval(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 alignstack(16) [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Already aligned above 16: unchanged.
+define internal void @internal_byval_high_align(ptr byval(%struct.S) align 32 %s) {
+; CHECK-LABEL: define internal void @internal_byval_high_align(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 32 [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Internal aggregate (non-byval) parameter: promoted to alignstack(16).
+define internal void @internal_aggr(%struct.S %s) {
+; CHECK-LABEL: define internal void @internal_aggr(
+; CHECK-SAME: [[STRUCT_S:%.*]] alignstack(16) [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Scalars gain nothing from vectorization: untouched.
+define internal void @internal_scalar(i32 %x, ptr %p) {
+; CHECK-LABEL: define internal void @internal_scalar(
+; CHECK-SAME: i32 [[X:%.*]], ptr [[P:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Internal aggregate return value: promoted to alignstack(16).
+define internal %struct.S @internal_ret() {
+; CHECK-LABEL: define internal alignstack(16) %struct.S @internal_ret() {
+; CHECK-NEXT: ret [[STRUCT_S:%.*]] zeroinitializer
+;
+ ret %struct.S zeroinitializer
+}
+
+; An explicit stackalign takes precedence.
+define internal void @internal_explicit_stackalign(%struct.S alignstack(8) %s) {
+; CHECK-LABEL: define internal void @internal_explicit_stackalign(
+; CHECK-SAME: [[STRUCT_S:%.*]] alignstack(8) [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; The promoted alignment is propagated to constant-offset byval loads.
+define internal void @internal_load_propagation(ptr byval(%struct.S) %s) {
+; CHECK-LABEL: define internal void @internal_load_propagation(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 alignstack(16) [[S:%.*]]) {
+; CHECK-NEXT: [[V1:%.*]] = load i32, ptr [[S]], align 16
+; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds [[STRUCT_S]], ptr [[S]], i64 0, i32 1
+; CHECK-NEXT: [[V2:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT: ret void
+;
+ %v1 = load i32, ptr %s, align 1
+ %gep = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1
+ %v2 = load i32, ptr %gep, align 1
+ ret void
+}
+
+; External byval loads get the ABI align (4) propagated, avoiding byte-wise loads.
+define void @extern_load_propagation(ptr byval(%struct.S) %s) {
+; CHECK-LABEL: define void @extern_load_propagation(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
+; CHECK-NEXT: [[V1:%.*]] = load i32, ptr [[S]], align 4
+; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds [[STRUCT_S]], ptr [[S]], i64 0, i32 1
+; CHECK-NEXT: [[V2:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT: ret void
+;
+ %v1 = load i32, ptr %s, align 1
+ %gep = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1
+ %v2 = load i32, ptr %gep, align 1
+ ret void
+}
+
+; Kernel byval loads get the ABI align (4) propagated.
+define ptx_kernel void @kernel_load_propagation(ptr byval(%struct.S) align 1 %s) {
+; CHECK-LABEL: define ptx_kernel void @kernel_load_propagation(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
+; CHECK-NEXT: [[V1:%.*]] = load i32, ptr [[S]], align 4
+; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds [[STRUCT_S]], ptr [[S]], i64 0, i32 1
+; CHECK-NEXT: [[V2:%.*]] = load i32, ptr [[GEP]], align 4
+; CHECK-NEXT: ret void
+;
+ %v1 = load i32, ptr %s, align 1
+ %gep = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1
+ %v2 = load i32, ptr %gep, align 1
+ ret void
+}
+
+; The promoted alignment is mirrored onto direct call sites.
+define internal %struct.S @callee(%struct.S %a) {
+; CHECK-LABEL: define internal alignstack(16) %struct.S @callee(
+; CHECK-SAME: [[STRUCT_S:%.*]] alignstack(16) [[A:%.*]]) {
+; CHECK-NEXT: ret [[STRUCT_S]] [[A]]
+;
+ ret %struct.S %a
+}
+
+define void @caller() {
+; CHECK-LABEL: define void @caller() {
+; CHECK-NEXT: [[R:%.*]] = call alignstack(16) [[STRUCT_S:%.*]] @[[CALLEE:[a-zA-Z0-9_$\"\\.-]*[a-zA-Z_$\"\\.-][a-zA-Z0-9_$\"\\.-]*]]([[STRUCT_S]] alignstack(16) zeroinitializer)
+; CHECK-NEXT: ret void
+;
+ %r = call %struct.S @callee(%struct.S zeroinitializer)
+ ret void
+}
+
+; A byval call matching the callee's ABI is promoted and mirrored onto the call.
+define internal void @byval_match(ptr byval(%struct.S) %s) {
+; CHECK-LABEL: define internal void @byval_match(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 alignstack(16) [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+define void @byval_match_caller(ptr %p) {
+; CHECK-LABEL: define void @byval_match_caller(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT: call void @byval_match(ptr byval([[STRUCT_S:%.*]]) alignstack(16) [[P]])
+; CHECK-NEXT: ret void
+;
+ call void @byval_match(ptr byval(%struct.S) %p)
+ ret void
+}
+
+; A byval call whose align differs from the callee is not promoted.
+define internal void @byval_abi_mismatch(ptr byval(%struct.S) align 8 %s) {
+; CHECK-LABEL: define internal void @byval_abi_mismatch(
+; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 8 [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+define void @byval_abi_mismatch_caller(ptr %p) {
+; CHECK-LABEL: define void @byval_abi_mismatch_caller(
+; CHECK-SAME: ptr [[P:%.*]]) {
+; CHECK-NEXT: call void @byval_abi_mismatch(ptr byval([[STRUCT_S:%.*]]) align 4 [[P]])
+; CHECK-NEXT: ret void
+;
+ call void @byval_abi_mismatch(ptr byval(%struct.S) align 4 %p)
+ ret void
+}
+
+; Address-taken functions aren't promoted (an indirect call wouldn't see them).
+define internal void @addr_taken(%struct.S %s) {
+; CHECK-LABEL: define internal void @addr_taken(
+; CHECK-SAME: [[STRUCT_S:%.*]] [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+define void @take_address() {
+; CHECK-LABEL: define void @take_address() {
+; CHECK-NEXT: store ptr @addr_taken, ptr @gfp, align 8
+; CHECK-NEXT: ret void
+;
+ store ptr @addr_taken, ptr @gfp
+ ret void
+}
+
+; A type-mismatched call is lowered as indirect, so the callee isn't promoted.
+define internal void @type_mismatch(%struct.S %s) {
+; CHECK-LABEL: define internal void @type_mismatch(
+; CHECK-SAME: [[STRUCT_S:%.*]] [[S:%.*]]) {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+define void @mismatch_caller() {
+; CHECK-LABEL: define void @mismatch_caller() {
+; CHECK-NEXT: call void @type_mismatch(i64 0)
+; CHECK-NEXT: ret void
+;
+ call void @type_mismatch(i64 0)
+ ret void
+}
diff --git a/llvm/test/CodeGen/NVPTX/set-byval-param-align.ll b/llvm/test/CodeGen/NVPTX/set-byval-param-align.ll
deleted file mode 100644
index 9d3843af8a555..0000000000000
--- a/llvm/test/CodeGen/NVPTX/set-byval-param-align.ll
+++ /dev/null
@@ -1,120 +0,0 @@
-; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=nvptx-set-byval-param-align -S | FileCheck %s
-
-target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"
-target triple = "nvptx64-nvidia-cuda"
-
-%struct.S = type { i32, i32 }
-
-; Kernel byval with no alignment: should get ABI align of i32 (4).
-define ptx_kernel void @kernel_no_align(ptr byval(i32) %s) {
-; CHECK-LABEL: define ptx_kernel void @kernel_no_align(
-; CHECK-SAME: ptr byval(i32) align 4 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Kernel byval with low alignment: should be bumped to ABI align (4).
-define ptx_kernel void @kernel_low_align(ptr byval(%struct.S) align 1 %s) {
-; CHECK-LABEL: define ptx_kernel void @kernel_low_align(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Kernel byval with adequate alignment: should stay unchanged.
-define ptx_kernel void @kernel_high_align(ptr byval(%struct.S) align 8 %s) {
-; CHECK-LABEL: define ptx_kernel void @kernel_high_align(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 8 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; External device function: gets ABI alignment (4).
-define void @device_extern_no_align(ptr byval(%struct.S) %s) {
-; CHECK-LABEL: define void @device_extern_no_align(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Internal device function: gets max(16, ABI) = 16.
-define internal void @device_internal_no_align(ptr byval(%struct.S) %s) {
-; CHECK-LABEL: define internal void @device_internal_no_align(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 16 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Internal device function with adequate alignment: stays unchanged.
-define internal void @device_internal_high_align(ptr byval(%struct.S) align 32 %s) {
-; CHECK-LABEL: define internal void @device_internal_high_align(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 32 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Large vector type: ABI alignment exceeds 128 bytes, capped to PTX maximum.
-define ptx_kernel void @kernel_align_cap(ptr byval(<256 x i32>) %s) {
-; CHECK-LABEL: define ptx_kernel void @kernel_align_cap(
-; CHECK-SAME: ptr byval(<256 x i32>) align 128 [[S:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Non-byval args should be untouched.
-define ptx_kernel void @no_byval(ptr %p, i32 %x) {
-; CHECK-LABEL: define ptx_kernel void @no_byval(
-; CHECK-SAME: ptr [[P:%.*]], i32 [[X:%.*]]) {
-; CHECK-NEXT: ret void
-;
- ret void
-}
-
-; Load propagation: loads at known offsets get alignment upgraded.
-define ptx_kernel void @load_propagation(ptr byval(%struct.S) %s) {
-; CHECK-LABEL: define ptx_kernel void @load_propagation(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
-; CHECK-NEXT: [[V1:%.*]] = load i32, ptr [[S]], align 4
-; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds [[STRUCT_S]], ptr [[S]], i64 0, i32 1
-; CHECK-NEXT: [[V2:%.*]] = load i32, ptr [[GEP]], align 4
-; CHECK-NEXT: ret void
-;
- %v1 = load i32, ptr %s, align 1
- %gep = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1
- %v2 = load i32, ptr %gep, align 1
- ret void
-}
-
-; Load propagation through bitcast.
-define ptx_kernel void @load_through_bitcast(ptr byval(%struct.S) %s) {
-; CHECK-LABEL: define ptx_kernel void @load_through_bitcast(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 4 [[S:%.*]]) {
-; CHECK-NEXT: [[V1:%.*]] = load i64, ptr [[S]], align 4
-; CHECK-NEXT: ret void
-;
- %v1 = load i64, ptr %s, align 1
- ret void
-}
-
-; Internal function gets higher alignment, which propagates to loads.
-define internal void @internal_load_propagation(ptr byval(%struct.S) %s) {
-; CHECK-LABEL: define internal void @internal_load_propagation(
-; CHECK-SAME: ptr byval([[STRUCT_S:%.*]]) align 16 [[S:%.*]]) {
-; CHECK-NEXT: [[V1:%.*]] = load i32, ptr [[S]], align 16
-; CHECK-NEXT: [[GEP:%.*]] = getelementptr inbounds [[STRUCT_S]], ptr [[S]], i64 0, i32 1
-; CHECK-NEXT: [[V2:%.*]] = load i32, ptr [[GEP]], align 4
-; CHECK-NEXT: ret void
-;
- %v1 = load i32, ptr %s, align 1
- %gep = getelementptr inbounds %struct.S, ptr %s, i64 0, i32 1
- %v2 = load i32, ptr %gep, align 1
- ret void
-}
diff --git a/llvm/utils/gn/secondary/llvm/lib/Target/NVPTX/BUILD.gn b/llvm/utils/gn/secondary/llvm/lib/Target/NVPTX/BUILD.gn
index b183efb52836d..11b340a80d5be 100644
--- a/llvm/utils/gn/secondary/llvm/lib/Target/NVPTX/BUILD.gn
+++ b/llvm/utils/gn/secondary/llvm/lib/Target/NVPTX/BUILD.gn
@@ -59,11 +59,11 @@ static_library("LLVMNVPTXCodeGen") {
"NVPTXMarkKernelPtrsGlobal.cpp",
"NVPTXPeephole.cpp",
"NVPTXPrologEpilogPass.cpp",
+ "NVPTXPromoteParamAlign.cpp",
"NVPTXProxyRegErasure.cpp",
"NVPTXRegisterInfo.cpp",
"NVPTXReplaceImageHandles.cpp",
"NVPTXSelectionDAGInfo.cpp",
- "NVPTXSetByValParamAlign.cpp",
"NVPTXSubtarget.cpp",
"NVPTXTagInvariantLoads.cpp",
"NVPTXTargetMachine.cpp",
>From d5cd3ee20bd6ed4427bc5d44044e851f5ced15a9 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Tue, 7 Jul 2026 09:09:30 -0700
Subject: [PATCH 2/2] [NVPTX] Promote interal function alignments in IR pass
---
llvm/lib/Target/NVPTX/NVPTXUtilities.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
index 2b636f0ac8cad..1cfb552a0dd42 100644
--- a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
@@ -60,8 +60,8 @@ Align llvm::getDeviceByValParamAlign(const Function *F, Type *ArgTy,
getByValParamAlignFloor(F));
}
-Align getDeviceByValParamAlign(const CallBase *CB, Type *ArgTy,
- unsigned AttrIdx, const DataLayout &DL) {
+Align llvm::getDeviceByValParamAlign(const CallBase *CB, Type *ArgTy,
+ unsigned AttrIdx, const DataLayout &DL) {
Align ParamAlign = getPTXParamAlign(CB, ArgTy, AttrIdx, DL);
// For an indirect call getPTXParamAlign can't see the call's own byval
More information about the llvm-commits
mailing list