[llvm] [NVPTX] Add NVVMUpgradeAnnotations pass to cleanup legacy annotations (PR #119261)
Alex MacLean via llvm-commits
llvm-commits at lists.llvm.org
Mon Dec 9 12:18:00 PST 2024
https://github.com/AlexMaclean created https://github.com/llvm/llvm-project/pull/119261
None
>From f9f30a77f5e7232f968a3063c34338c9dfc7bac5 Mon Sep 17 00:00:00 2001
From: Alex Maclean <amaclean at nvidia.com>
Date: Fri, 8 Nov 2024 22:39:34 +0000
Subject: [PATCH] [NVPTX] Add NVVMUpgradeAnnotations pass to cleanup legacy
annotations
---
llvm/lib/Target/NVPTX/CMakeLists.txt | 1 +
llvm/lib/Target/NVPTX/NVPTX.h | 5 +
llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp | 4 +
llvm/lib/Target/NVPTX/NVPTXUtilities.cpp | 9 +-
.../Target/NVPTX/NVVMUpgradeAnnotations.cpp | 130 ++++++++++++++++++
.../CodeGen/NVPTX/upgrade-nvvm-annotations.ll | 30 ++++
6 files changed, 177 insertions(+), 2 deletions(-)
create mode 100644 llvm/lib/Target/NVPTX/NVVMUpgradeAnnotations.cpp
create mode 100644 llvm/test/CodeGen/NVPTX/upgrade-nvvm-annotations.ll
diff --git a/llvm/lib/Target/NVPTX/CMakeLists.txt b/llvm/lib/Target/NVPTX/CMakeLists.txt
index 693365161330f5..bb2e4ad48b51d8 100644
--- a/llvm/lib/Target/NVPTX/CMakeLists.txt
+++ b/llvm/lib/Target/NVPTX/CMakeLists.txt
@@ -39,6 +39,7 @@ set(NVPTXCodeGen_sources
NVVMReflect.cpp
NVPTXProxyRegErasure.cpp
NVPTXCtorDtorLowering.cpp
+ NVVMUpgradeAnnotations.cpp
)
add_llvm_target(NVPTXCodeGen
diff --git a/llvm/lib/Target/NVPTX/NVPTX.h b/llvm/lib/Target/NVPTX/NVPTX.h
index ca915cd3f3732f..53418148be3615 100644
--- a/llvm/lib/Target/NVPTX/NVPTX.h
+++ b/llvm/lib/Target/NVPTX/NVPTX.h
@@ -52,6 +52,7 @@ FunctionPass *createNVPTXLowerUnreachablePass(bool TrapUnreachable,
bool NoTrapAfterNoreturn);
MachineFunctionPass *createNVPTXPeephole();
MachineFunctionPass *createNVPTXProxyRegErasurePass();
+ModulePass *createNVVMUpgradeAnnotationsPass();
struct NVVMIntrRangePass : PassInfoMixin<NVVMIntrRangePass> {
PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
@@ -74,6 +75,10 @@ struct NVPTXCopyByValArgsPass : PassInfoMixin<NVPTXCopyByValArgsPass> {
PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
};
+struct NVVMUpgradeAnnotationsPass : PassInfoMixin<NVVMUpgradeAnnotationsPass> {
+ PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
namespace NVPTX {
enum DrvInterface {
NVCL,
diff --git a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
index a5c5e9420ee737..b4fd36625adc9c 100644
--- a/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXTargetMachine.cpp
@@ -254,6 +254,8 @@ void NVPTXTargetMachine::registerPassBuilderCallbacks(PassBuilder &PB) {
PB.registerPipelineStartEPCallback(
[this](ModulePassManager &PM, OptimizationLevel Level) {
+ PM.addPass(NVVMUpgradeAnnotationsPass());
+
FunctionPassManager FPM;
FPM.addPass(NVVMReflectPass(Subtarget.getSmVersion()));
// Note: NVVMIntrRangePass was causing numerical discrepancies at one
@@ -349,6 +351,8 @@ void NVPTXPassConfig::addIRPasses() {
AAR.addAAResult(WrapperPass->getResult());
}));
+ addPass(createNVVMUpgradeAnnotationsPass());
+
// NVVMReflectPass is added in addEarlyAsPossiblePasses, so hopefully running
// it here does nothing. But since we need it for correctness when lowering
// to NVPTX, run it here too, in case whoever built our pass pipeline didn't
diff --git a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
index 98bffd92a087b6..04e83576cbf958 100644
--- a/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXUtilities.cpp
@@ -311,11 +311,16 @@ std::optional<unsigned> getMaxNReg(const Function &F) {
}
bool isKernelFunction(const Function &F) {
+ if (F.getCallingConv() == CallingConv::PTX_Kernel)
+ return true;
+
+ if (F.hasFnAttribute("nvvm.kernel"))
+ return true;
+
if (const auto X = findOneNVVMAnnotation(&F, "kernel"))
return (*X == 1);
- // There is no NVVM metadata, check the calling convention
- return F.getCallingConv() == CallingConv::PTX_Kernel;
+ return false;
}
MaybeAlign getAlign(const Function &F, unsigned Index) {
diff --git a/llvm/lib/Target/NVPTX/NVVMUpgradeAnnotations.cpp b/llvm/lib/Target/NVPTX/NVVMUpgradeAnnotations.cpp
new file mode 100644
index 00000000000000..ca550434835a2c
--- /dev/null
+++ b/llvm/lib/Target/NVPTX/NVVMUpgradeAnnotations.cpp
@@ -0,0 +1,130 @@
+//===- NVVMUpgradeAnnotations.cpp - Upgrade NVVM Annotations --------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This pass replaces deprecated metadata in nvvm.annotation with a more modern
+// IR representation.
+//
+//===----------------------------------------------------------------------===//
+
+#include "NVPTX.h"
+#include "llvm/ADT/SmallSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/CodeGen/Passes.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/Pass.h"
+#include <cstdint>
+
+#define DEBUG_TYPE "nvvm-upgrade-annotations"
+
+using namespace llvm;
+
+namespace llvm {
+void initializeNVVMUpgradeAnnotationsLegacyPassPass(PassRegistry &);
+} // namespace llvm
+
+namespace {
+
+class NVVMUpgradeAnnotationsLegacyPass : public ModulePass {
+public:
+ static char ID;
+ NVVMUpgradeAnnotationsLegacyPass() : ModulePass(ID) {
+ initializeNVVMUpgradeAnnotationsLegacyPassPass(
+ *PassRegistry::getPassRegistry());
+ }
+ bool runOnModule(Module &M) override;
+};
+} // namespace
+
+char NVVMUpgradeAnnotationsLegacyPass::ID = 0;
+
+bool static autoUpgradeAnnotation(Function *F, StringRef K, const Metadata *V) {
+ if (K == "kernel") {
+ assert(mdconst::extract<ConstantInt>(V)->getZExtValue() == 1);
+ F->addFnAttr("nvvm.kernel");
+ return true;
+ }
+ if (K == "align") {
+ const uint64_t AlignBits = mdconst::extract<ConstantInt>(V)->getZExtValue();
+ const unsigned Idx = (AlignBits >> 16);
+ const Align StackAlign = Align(AlignBits & 0xFFFF);
+ // TODO: Skip adding the stackalign attribute for returns, for now.
+ if (!Idx)
+ return false;
+ F->addAttributeAtIndex(
+ Idx, Attribute::getWithStackAlignment(F->getContext(), StackAlign));
+ return true;
+ }
+
+ return false;
+}
+
+// Iterate over nvvm.annotations rewriting them as appropiate.
+void static upgradeNVAnnotations(Module &M) {
+ NamedMDNode *NamedMD = M.getNamedMetadata("nvvm.annotations");
+ if (!NamedMD)
+ return;
+
+ SmallVector<MDNode *, 8> NewNodes;
+ SmallSet<const MDNode *, 8> SeenNodes;
+ for (MDNode *MD : NamedMD->operands()) {
+ if (SeenNodes.contains(MD))
+ continue;
+ SeenNodes.insert(MD);
+
+ Function *F = mdconst::dyn_extract_or_null<Function>(MD->getOperand(0));
+ if (!F)
+ continue;
+
+ assert(MD && "Invalid MDNode for annotation");
+ assert((MD->getNumOperands() % 2) == 1 && "Invalid number of operands");
+
+ SmallVector<Metadata *, 8> NewOperands;
+ // start index = 1, to skip the global variable key
+ // increment = 2, to skip the value for each property-value pairs
+ for (unsigned j = 1, je = MD->getNumOperands(); j < je; j += 2) {
+ MDString *K = cast<MDString>(MD->getOperand(j));
+ const MDOperand &V = MD->getOperand(j + 1);
+ bool Upgraded = autoUpgradeAnnotation(F, K->getString(), V);
+ if (!Upgraded)
+ NewOperands.append({K, V});
+ }
+
+ if (!NewOperands.empty()) {
+ NewOperands.insert(NewOperands.begin(), MD->getOperand(0));
+ NewNodes.push_back(MDNode::get(M.getContext(), NewOperands));
+ }
+ }
+
+ NamedMD->clearOperands();
+ for (MDNode *N : NewNodes)
+ NamedMD->addOperand(N);
+}
+
+PreservedAnalyses NVVMUpgradeAnnotationsPass::run(Module &M,
+ ModuleAnalysisManager &AM) {
+ upgradeNVAnnotations(M);
+ return PreservedAnalyses::all();
+}
+
+bool NVVMUpgradeAnnotationsLegacyPass::runOnModule(Module &M) {
+ upgradeNVAnnotations(M);
+ return false;
+}
+
+INITIALIZE_PASS(NVVMUpgradeAnnotationsLegacyPass, DEBUG_TYPE,
+ "NVVMUpgradeAnnotations", false, false)
+
+ModulePass *llvm::createNVVMUpgradeAnnotationsPass() {
+ return new NVVMUpgradeAnnotationsLegacyPass();
+}
diff --git a/llvm/test/CodeGen/NVPTX/upgrade-nvvm-annotations.ll b/llvm/test/CodeGen/NVPTX/upgrade-nvvm-annotations.ll
new file mode 100644
index 00000000000000..68dc2353858cb3
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/upgrade-nvvm-annotations.ll
@@ -0,0 +1,30 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals all --version 5
+; RUN: opt < %s -mtriple=nvptx64-unknown-unknown -O0 -S | FileCheck %s
+
+define i32 @foo(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @foo(
+; CHECK-SAME: i32 alignstack(8) [[A:%.*]], i32 alignstack(16) [[B:%.*]]) {
+; CHECK-NEXT: ret i32 0
+;
+ ret i32 0
+}
+
+define i32 @bar(i32 %a, i32 %b) {
+; CHECK-LABEL: define i32 @bar(
+; CHECK-SAME: i32 [[A:%.*]], i32 [[B:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: ret i32 0
+;
+ ret i32 0
+}
+
+!nvvm.annotations = !{!0, !1, !2}
+
+!0 = !{ptr @foo, !"align", i32 u0x00000008, !"align", i32 u0x00010008, !"align", i32 u0x00020010}
+!1 = !{null, !"align", i32 u0x00000008, !"align", i32 u0x00010008, !"align", i32 u0x00020008}
+!2 = !{ptr @bar, !"kernel", i32 1}
+
+;.
+; CHECK: attributes #[[ATTR0]] = { "nvvm.kernel" }
+;.
+; CHECK: [[META0:![0-9]+]] = !{ptr @foo, !"align", i32 8}
+;.
More information about the llvm-commits
mailing list