[mlir][llvm] Preserve function entry count metadata (#204707)
Fixes #202374. Preserve LLVM function `!prof` metadata for: - `synthetic_function_entry_count` - `function_entry_count` with import GUID operands The importer represents the entry count, count kind, and trailing import GUID operands in a single `#llvm.function_entry_count` attribute, validates malformed operands before mutating the operation, and exports the metadata through `llvm::Function::setEntryCount`. Tests cover import, export, malformed metadata operands, synthetic counts, synthetic counts with trailing operands, import GUID round-tripping, and raw i64 bit-pattern preservation. Checks: - `git diff --check` - `git clang-format --diff origin/main -- ...` - `llvm-as -disable-output` smoke check Codex was used in making this PR.
This commit is contained in:
1sgtpepper
GitHub
parent
fba1a4b129
commit
8518c74b5e
8 changed files with 225 additions and 17 deletions
|
|
@ -96,6 +96,28 @@ def FramePointerKindAttr : LLVM_Attr<"FramePointerKind", "framePointerKind"> {
|
|||
let assemblyFormat = "`<` $framePointerKind `>`";
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// FunctionEntryCountAttr
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
def LLVM_FunctionEntryCountAttr
|
||||
: LLVM_Attr<"FunctionEntryCount", "function_entry_count"> {
|
||||
let summary = "LLVM function entry count profile metadata";
|
||||
let description = [{
|
||||
Models function-level `!prof` entry-count metadata. The `entry_count` field
|
||||
stores the unsigned 64-bit counter bit pattern. The `count_type` field
|
||||
defaults to real and selects whether the metadata is emitted as
|
||||
`"function_entry_count"` or `"synthetic_function_entry_count"`. The optional
|
||||
`imports` field stores the trailing import GUID operands used by ThinLTO
|
||||
sample PGO.
|
||||
}];
|
||||
let parameters = (ins "uint64_t":$entry_count,
|
||||
DefaultValuedParameter<"ProfileCountType",
|
||||
"ProfileCountType::Real">:$count_type,
|
||||
OptionalArrayRefParameter<"uint64_t">:$imports);
|
||||
let assemblyFormat = "`<` struct(params) `>`";
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// Loop Attributes
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
|
|
|||
|
|
@ -83,6 +83,21 @@ def AsmATTOrIntel : LLVM_EnumAttr<
|
|||
let cppNamespace = "::mlir::LLVM";
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// ProfileCountType
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
||||
def ProfileCountReal : I64EnumAttrCase<
|
||||
/*string sym=*/"Real", /*int val=*/0, /*string str=*/"real">;
|
||||
def ProfileCountSynthetic : I64EnumAttrCase<
|
||||
/*string sym=*/"Synthetic", /*int val=*/1, /*string str=*/"synthetic">;
|
||||
def ProfileCountTypeAttr : I64EnumAttr<
|
||||
/*string name=*/"ProfileCountType",
|
||||
/*string description=*/"real or synthetic function entry count",
|
||||
/*list<I64EnumAttrCase> cases=*/[ProfileCountReal, ProfileCountSynthetic]> {
|
||||
let cppNamespace = "::mlir::LLVM";
|
||||
}
|
||||
|
||||
//===----------------------------------------------------------------------===//
|
||||
// Atomic Operations
|
||||
//===----------------------------------------------------------------------===//
|
||||
|
|
|
|||
|
|
@ -2018,6 +2018,11 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
|
|||
llvm.return
|
||||
}
|
||||
```
|
||||
|
||||
The `function_entry_count` attribute models function-level `!prof`
|
||||
entry-count metadata. It stores the entry count, whether the count is real
|
||||
or synthetic, and any trailing import GUID operands in a single
|
||||
`#llvm.function_entry_count` attribute.
|
||||
}];
|
||||
|
||||
let arguments = (ins
|
||||
|
|
@ -2034,7 +2039,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
|
|||
OptionalAttr<ArrayAttr>:$passthrough,
|
||||
OptionalAttr<DictArrayAttr>:$arg_attrs,
|
||||
OptionalAttr<DictArrayAttr>:$res_attrs,
|
||||
OptionalAttr<I64Attr>:$function_entry_count,
|
||||
OptionalAttr<LLVM_FunctionEntryCountAttr>:$function_entry_count,
|
||||
OptionalAttr<LLVM_MemoryEffectsAttr>:$memory_effects,
|
||||
DefaultValuedAttr<Visibility, "mlir::LLVM::Visibility::Default">:$visibility_,
|
||||
UnitAttr:$arm_streaming, UnitAttr:$arm_locally_streaming,
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@
|
|||
|
||||
#include "llvm/ADT/APFloat.h"
|
||||
#include "llvm/ADT/DenseSet.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/ADT/TypeSwitch.h"
|
||||
#include "llvm/IR/DataLayout.h"
|
||||
#include "llvm/Support/Error.h"
|
||||
|
|
@ -3044,7 +3045,9 @@ void LLVMFuncOp::build(OpBuilder &builder, OperationState &result,
|
|||
result.addAttribute(getComdatAttrName(result.name), comdat);
|
||||
if (functionEntryCount)
|
||||
result.addAttribute(getFunctionEntryCountAttrName(result.name),
|
||||
builder.getI64IntegerAttr(functionEntryCount.value()));
|
||||
FunctionEntryCountAttr::get(
|
||||
builder.getContext(), *functionEntryCount,
|
||||
ProfileCountType::Real, ArrayRef<uint64_t>{}));
|
||||
#ifndef NDEBUG
|
||||
std::optional<NamedAttribute> duplicate = result.attributes.findDuplicate();
|
||||
if (duplicate.has_value()) {
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@
|
|||
#include "llvm/IR/Instructions.h"
|
||||
#include "llvm/IR/IntrinsicInst.h"
|
||||
#include "llvm/IR/MemoryModelRelaxationAnnotations.h"
|
||||
#include <optional>
|
||||
|
||||
using namespace mlir;
|
||||
using namespace mlir::LLVM;
|
||||
|
|
@ -99,6 +100,14 @@ getSupportedMetadataImpl(llvm::LLVMContext &llvmContext) {
|
|||
return convertibleMetadata;
|
||||
}
|
||||
|
||||
/// Extracts an LLVM metadata constant as an unsigned 64-bit integer.
|
||||
static std::optional<uint64_t> getUInt64Metadata(llvm::Metadata *metadata) {
|
||||
auto *constant = llvm::mdconst::dyn_extract<llvm::ConstantInt>(metadata);
|
||||
if (!constant)
|
||||
return std::nullopt;
|
||||
return constant->getValue().tryZExtValue();
|
||||
}
|
||||
|
||||
/// Converts the given profiling metadata `node` to an MLIR profiling attribute
|
||||
/// and attaches it to the imported operation if the translation succeeds.
|
||||
/// Returns failure otherwise.
|
||||
|
|
@ -112,27 +121,45 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
|
|||
auto *name = dyn_cast<llvm::MDString>(node->getOperand(0));
|
||||
if (!name)
|
||||
return failure();
|
||||
StringRef profName = name->getString();
|
||||
|
||||
// Handle function entry count metadata.
|
||||
if (name->getString() == llvm::MDProfLabels::FunctionEntryCount) {
|
||||
|
||||
// TODO support function entry count metadata with GUID fields.
|
||||
if (node->getNumOperands() != 2)
|
||||
if (profName == llvm::MDProfLabels::FunctionEntryCount ||
|
||||
profName == llvm::MDProfLabels::SyntheticFunctionEntryCount) {
|
||||
if (node->getNumOperands() < 2)
|
||||
return failure();
|
||||
|
||||
llvm::ConstantInt *entryCount =
|
||||
llvm::mdconst::dyn_extract<llvm::ConstantInt>(node->getOperand(1));
|
||||
if (!entryCount)
|
||||
bool isSynthetic =
|
||||
profName == llvm::MDProfLabels::SyntheticFunctionEntryCount;
|
||||
ProfileCountType profileCountType =
|
||||
isSynthetic ? ProfileCountType::Synthetic : ProfileCountType::Real;
|
||||
|
||||
std::optional<uint64_t> entryCountValue =
|
||||
getUInt64Metadata(node->getOperand(1));
|
||||
if (!entryCountValue)
|
||||
return failure();
|
||||
|
||||
SmallVector<uint64_t> importGUIDValues;
|
||||
importGUIDValues.reserve(node->getNumOperands() - 2);
|
||||
for (unsigned idx = 2, e = node->getNumOperands(); idx < e; ++idx) {
|
||||
std::optional<uint64_t> guidValue =
|
||||
getUInt64Metadata(node->getOperand(idx));
|
||||
if (!guidValue)
|
||||
return failure();
|
||||
importGUIDValues.push_back(*guidValue);
|
||||
}
|
||||
|
||||
if (auto funcOp = dyn_cast<LLVMFuncOp>(op)) {
|
||||
funcOp.setFunctionEntryCount(entryCount->getZExtValue());
|
||||
funcOp.setFunctionEntryCountAttr(
|
||||
FunctionEntryCountAttr::get(builder.getContext(), *entryCountValue,
|
||||
profileCountType, importGUIDValues));
|
||||
return success();
|
||||
}
|
||||
return op->emitWarning()
|
||||
<< "expected function_entry_count to be attached to a function";
|
||||
}
|
||||
|
||||
if (name->getString() != llvm::MDProfLabels::BranchWeights)
|
||||
if (profName != llvm::MDProfLabels::BranchWeights)
|
||||
return failure();
|
||||
// The branch_weights metadata must have at least 2 operands.
|
||||
if (node->getNumOperands() < 2)
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@
|
|||
#include "mlir/Target/LLVMIR/LLVMTranslationInterface.h"
|
||||
#include "mlir/Target/LLVMIR/TypeToLLVM.h"
|
||||
|
||||
#include "llvm/ADT/DenseSet.h"
|
||||
#include "llvm/ADT/STLExtras.h"
|
||||
#include "llvm/ADT/StringExtras.h"
|
||||
#include "llvm/ADT/TypeSwitch.h"
|
||||
|
|
@ -2018,8 +2019,19 @@ LogicalResult ModuleTranslation::convertFunctionSignatures() {
|
|||
convertFunctionKernelAttributes(function, llvmFunc, *this);
|
||||
|
||||
// Convert function_entry_count attribute to metadata.
|
||||
if (std::optional<uint64_t> entryCount = function.getFunctionEntryCount())
|
||||
llvmFunc->setEntryCount(entryCount.value());
|
||||
if (auto entryCount = function.getFunctionEntryCountAttr()) {
|
||||
ArrayRef<uint64_t> imports = entryCount.getImports();
|
||||
llvm::DenseSet<llvm::GlobalValue::GUID> importGUIDs;
|
||||
if (!imports.empty())
|
||||
importGUIDs.insert(imports.begin(), imports.end());
|
||||
llvm::MDBuilder metadataBuilder(llvmFunc->getContext());
|
||||
llvmFunc->setMetadata(
|
||||
llvm::LLVMContext::MD_prof,
|
||||
metadataBuilder.createFunctionEntryCount(
|
||||
entryCount.getEntryCount(),
|
||||
entryCount.getCountType() == ProfileCountType::Synthetic,
|
||||
imports.empty() ? nullptr : &importGUIDs));
|
||||
}
|
||||
|
||||
// Convert result attributes.
|
||||
if (ArrayAttr allResultAttrs = function.getAllResultAttrs()) {
|
||||
|
|
|
|||
|
|
@ -171,7 +171,7 @@ declare range(i64 0, 4097) i64 @func_res_attr_range()
|
|||
; // -----
|
||||
|
||||
; CHECK-LABEL: @entry_count
|
||||
; CHECK-SAME: attributes {function_entry_count = 4242 : i64}
|
||||
; CHECK-SAME: attributes {function_entry_count = #llvm.function_entry_count<entry_count = 4242>}
|
||||
define void @entry_count() !prof !1 {
|
||||
ret void
|
||||
}
|
||||
|
|
@ -180,6 +180,80 @@ define void @entry_count() !prof !1 {
|
|||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @synthetic_entry_count
|
||||
; CHECK-SAME: attributes {function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = synthetic>}
|
||||
define void @synthetic_entry_count() !prof !2 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!2 = !{!"synthetic_function_entry_count", i64 7}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @entry_count_imports
|
||||
; CHECK-SAME: attributes {function_entry_count = #llvm.function_entry_count<entry_count = 7, imports = 1234, 18446744073709551615, 4, 1234>}
|
||||
define void @entry_count_imports() !prof !3 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!3 = !{!"function_entry_count", i64 7, i64 1234, i64 -1, i64 4, i64 1234}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @synthetic_entry_count_imports
|
||||
; CHECK-SAME: attributes {function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = synthetic, imports = 1234>}
|
||||
define void @synthetic_entry_count_imports() !prof !4 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!4 = !{!"synthetic_function_entry_count", i64 7, i64 1234}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @entry_count_malformed_import
|
||||
; CHECK-NOT: function_entry_count
|
||||
; expected-warning @unknown {{unhandled function metadata}}
|
||||
define void @entry_count_malformed_import() !prof !5 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!5 = !{!"function_entry_count", i64 7, !"bad"}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @entry_count_too_wide_count
|
||||
; CHECK-NOT: function_entry_count
|
||||
; expected-warning @unknown {{unhandled function metadata}}
|
||||
define void @entry_count_too_wide_count() !prof !6 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!6 = !{!"function_entry_count", i128 18446744073709551616}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @entry_count_too_wide_import
|
||||
; CHECK-NOT: function_entry_count
|
||||
; expected-warning @unknown {{unhandled function metadata}}
|
||||
define void @entry_count_too_wide_import() !prof !7 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!7 = !{!"function_entry_count", i64 7, i128 18446744073709551616}
|
||||
|
||||
; // -----
|
||||
|
||||
; Preserve the raw i64 metadata bit pattern.
|
||||
; CHECK-LABEL: @entry_count_negative_count
|
||||
; CHECK-SAME: attributes {function_entry_count = #llvm.function_entry_count<entry_count = 18446744073709551615>}
|
||||
define void @entry_count_negative_count() !prof !8 {
|
||||
ret void
|
||||
}
|
||||
|
||||
!8 = !{!"function_entry_count", i64 -1}
|
||||
|
||||
; // -----
|
||||
|
||||
; CHECK-LABEL: @func_memory
|
||||
; CHECK-SAME: attributes {memory_effects = #llvm.memory_effects<other = readwrite, argMem = none, inaccessibleMem = readwrite, errnoMem = readwrite, targetMem0 = readwrite, targetMem1 = readwrite>}
|
||||
; CHECK: llvm.return
|
||||
|
|
|
|||
|
|
@ -1893,12 +1893,62 @@ llvm.func @my_allocator(i64) attributes {passthrough = [["allocsize", "429496729
|
|||
// -----
|
||||
|
||||
// CHECK-LABEL: @functionEntryCount
|
||||
// CHECK-SAME: !prof ![[PROF_ID:[0-9]*]]
|
||||
llvm.func @functionEntryCount() attributes {function_entry_count = 4242 : i64} {
|
||||
// CHECK-SAME: !prof ![[PROF_ID:[0-9]+]]
|
||||
llvm.func @functionEntryCount() attributes {
|
||||
function_entry_count = #llvm.function_entry_count<entry_count = 4242>
|
||||
} {
|
||||
llvm.return
|
||||
}
|
||||
|
||||
// CHECK: ![[PROF_ID]] = !{!"function_entry_count", i64 4242}
|
||||
// CHECK-DAG: ![[PROF_ID]] = !{!"function_entry_count", i64 4242}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @syntheticFunctionEntryCount
|
||||
// CHECK-SAME: !prof ![[SYNTH_PROF_ID:[0-9]+]]
|
||||
llvm.func @syntheticFunctionEntryCount() attributes {
|
||||
function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = synthetic>
|
||||
} {
|
||||
llvm.return
|
||||
}
|
||||
|
||||
// CHECK-DAG: ![[SYNTH_PROF_ID]] = !{!"synthetic_function_entry_count", i64 7}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @syntheticFunctionEntryCountWithImports
|
||||
// CHECK-SAME: !prof ![[SYNTH_IMPORTS_PROF_ID:[0-9]+]]
|
||||
llvm.func @syntheticFunctionEntryCountWithImports() attributes {
|
||||
function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = synthetic, imports = 1234, 4, 1234>
|
||||
} {
|
||||
llvm.return
|
||||
}
|
||||
|
||||
// CHECK-DAG: ![[SYNTH_IMPORTS_PROF_ID]] = !{!"synthetic_function_entry_count", i64 7, i64 4, i64 1234}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @functionEntryCountWithImports
|
||||
// CHECK-SAME: !prof ![[IMPORTS_PROF_ID:[0-9]+]]
|
||||
llvm.func @functionEntryCountWithImports() attributes {
|
||||
function_entry_count = #llvm.function_entry_count<entry_count = 7, imports = 1234, 4, 18446744073709551615, 1234>
|
||||
} {
|
||||
llvm.return
|
||||
}
|
||||
|
||||
// CHECK-DAG: ![[IMPORTS_PROF_ID]] = !{!"function_entry_count", i64 7, i64 4, i64 1234, i64 -1}
|
||||
|
||||
// -----
|
||||
|
||||
// CHECK-LABEL: @functionEntryCountNegativeCount
|
||||
// CHECK-SAME: !prof ![[NEG_PROF_ID:[0-9]+]]
|
||||
llvm.func @functionEntryCountNegativeCount() attributes {
|
||||
function_entry_count = #llvm.function_entry_count<entry_count = 18446744073709551615>
|
||||
} {
|
||||
llvm.return
|
||||
}
|
||||
|
||||
// CHECK-DAG: ![[NEG_PROF_ID]] = !{!"function_entry_count", i64 -1}
|
||||
|
||||
// -----
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue