[Mlir-commits] [mlir] 8518c74 - [mlir][llvm] Preserve function entry count metadata (#204707)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Jul 14 06:12:37 PDT 2026
Author: 1sgtpepper
Date: 2026-07-14T15:12:33+02:00
New Revision: 8518c74b5ef8950b9b604df588f29d9037fe1380
URL: https://github.com/llvm/llvm-project/commit/8518c74b5ef8950b9b604df588f29d9037fe1380
DIFF: https://github.com/llvm/llvm-project/commit/8518c74b5ef8950b9b604df588f29d9037fe1380.diff
LOG: [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.
Added:
Modified:
mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
mlir/test/Target/LLVMIR/Import/function-attributes.ll
mlir/test/Target/LLVMIR/llvmir.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 72497a5328f78..6e2d4628c6173 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
index 51ac465000341..5a8f195a91a10 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
@@ -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
//===----------------------------------------------------------------------===//
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 9d112e5ea227e..4bdcf5f6c1cc5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -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,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 58f569abff8ea..f3dc00e177b0f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -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()) {
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
index e9cd335835263..b6dee4e89ee4c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
@@ -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)
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index a9ca4a57e8e8a..b76a803f93ad9 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -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()) {
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
index 5d664519e100c..1c9c616ce3c5b 100644
--- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll
+++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
@@ -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
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 5eca8f19154a1..a665c72f4f54c 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -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}
// -----
More information about the Mlir-commits
mailing list