[Mlir-commits] [mlir] [mlir][llvm] Preserve function entry count metadata (PR #204707)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Sun Jun 21 13:17:12 PDT 2026


https://github.com/Kuhai9801 updated https://github.com/llvm/llvm-project/pull/204707

>From 9a6508a95be9dfc0dcaf7f59e3599d622ecc9d20 Mon Sep 17 00:00:00 2001
From: "Cyne Jarvis J. Zarceno" <cynejarviszarceno at gmail.com>
Date: Fri, 19 Jun 2026 14:40:54 +0800
Subject: [PATCH 1/2] [mlir][llvm] Preserve function entry count metadata

Preserve function entry count profile metadata when importing and exporting MLIR LLVM dialect functions.

Represent the synthetic bit and import GUID operands on llvm.func, validate GUID operands before mutating the operation, and export the metadata through llvm::Function::setEntryCount.

Add import, export, verifier, synthetic+imports, malformed operand, and round-trip coverage.
---
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td   | 11 +++
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    | 28 +++++++
 .../LLVMIR/LLVMIRToLLVMTranslation.cpp        | 66 ++++++++++++++--
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 23 +++++-
 mlir/test/Dialect/LLVMIR/invalid.mlir         | 55 +++++++++++++
 .../LLVMIR/Import/function-attributes.ll      | 78 +++++++++++++++++++
 .../Import/function-entry-count-roundtrip.ll  | 21 +++++
 mlir/test/Target/LLVMIR/llvmir.mlir           | 42 +++++++++-
 8 files changed, 311 insertions(+), 13 deletions(-)
 create mode 100644 mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 9d112e5ea227e..557e9881bd921 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -2018,6 +2018,15 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
       llvm.return
     }
     ```
+
+    The `function_entry_count` attribute models function-level `!prof`
+    entry-count metadata. The `function_entry_count_synthetic` unit attribute
+    selects the `"synthetic_function_entry_count"` metadata label. The
+    `function_entry_count_imports` attribute represents the non-empty set of
+    trailing import GUID operands on real `"function_entry_count"` metadata.
+    Those GUID operands are unsigned 64-bit LLVM GUID bit patterns stored in a
+    signed i64 array attribute in unsigned sorted-unique order. This attribute
+    must not be combined with `function_entry_count_synthetic`.
   }];
 
   let arguments = (ins
@@ -2035,6 +2044,8 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<DictArrayAttr>:$arg_attrs,
     OptionalAttr<DictArrayAttr>:$res_attrs,
     OptionalAttr<I64Attr>:$function_entry_count,
+    UnitAttr:$function_entry_count_synthetic,
+    OptionalAttr<DenseI64ArrayAttr>:$function_entry_count_imports,
     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..1898ad752f80c 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"
@@ -3266,6 +3267,33 @@ LogicalResult LLVMFuncOp::verify() {
   if (failed(verifyComdat(*this, getComdat())))
     return failure();
 
+  if (!getFunctionEntryCountAttr()) {
+    if (getFunctionEntryCountSynthetic())
+      return emitOpError() << "requires function_entry_count when "
+                              "function_entry_count_synthetic is set";
+    if (getFunctionEntryCountImportsAttr())
+      return emitOpError() << "requires function_entry_count when "
+                              "function_entry_count_imports is set";
+  }
+  if (DenseI64ArrayAttr imports = getFunctionEntryCountImportsAttr()) {
+    if (getFunctionEntryCountSynthetic())
+      return emitOpError() << "does not support function_entry_count_imports "
+                              "with function_entry_count_synthetic";
+
+    ArrayRef<int64_t> values = imports.asArrayRef();
+    if (values.empty())
+      return emitOpError() << "requires function_entry_count_imports to be "
+                              "non-empty when set";
+
+    for (auto [previous, current] :
+         llvm::zip_equal(values.drop_back(), values.drop_front())) {
+      if (static_cast<uint64_t>(previous) >= static_cast<uint64_t>(current))
+        return emitOpError()
+               << "requires function_entry_count_imports to be sorted and "
+                  "unique by unsigned GUID value";
+    }
+  }
+
   if (isExternal()) {
     if (getLinkage() != LLVM::Linkage::External &&
         getLinkage() != LLVM::Linkage::ExternWeak)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
index e9cd335835263..d15b940ca342c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
@@ -17,12 +17,17 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/ModuleImport.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InlineAsm.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/MemoryModelRelaxationAnnotations.h"
+#include <optional>
+
+#include <algorithm>
+#include <optional>
 
 using namespace mlir;
 using namespace mlir::LLVM;
@@ -102,6 +107,14 @@ getSupportedMetadataImpl(llvm::LLVMContext &llvmContext) {
 /// 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.
+static std::optional<uint64_t> getUInt64Metadata(llvm::Metadata *metadata) {
+  llvm::ConstantInt *constant =
+      llvm::mdconst::dyn_extract<llvm::ConstantInt>(metadata);
+  if (!constant)
+    return std::nullopt;
+  return constant->getValue().tryZExtValue();
+}
+
 static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
                                       Operation *op,
                                       LLVM::ModuleImport &moduleImport) {
@@ -112,27 +125,64 @@ 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) {
+  if (profName == llvm::MDProfLabels::FunctionEntryCount ||
+      profName == llvm::MDProfLabels::SyntheticFunctionEntryCount) {
+    if (node->getNumOperands() < 2)
+      return failure();
+
+    bool isSynthetic =
+        profName == llvm::MDProfLabels::SyntheticFunctionEntryCount;
 
-    // TODO support function entry count metadata with GUID fields.
-    if (node->getNumOperands() != 2)
+    // LLVM's semantic import-GUID API only reads trailing GUID operands from
+    // "function_entry_count" metadata. Do not model trailing operands on
+    // "synthetic_function_entry_count" as import GUIDs in MLIR.
+    if (isSynthetic && node->getNumOperands() > 2)
       return failure();
 
-    llvm::ConstantInt *entryCount =
-        llvm::mdconst::dyn_extract<llvm::ConstantInt>(node->getOperand(1));
-    if (!entryCount)
+    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);
+    }
+
+    // Import GUIDs are semantically a set in LLVM. Canonicalize them as
+    // unsigned sorted-unique values before storing the bit patterns in MLIR.
+    llvm::sort(importGUIDValues);
+    importGUIDValues.erase(
+        std::unique(importGUIDValues.begin(), importGUIDValues.end()),
+        importGUIDValues.end());
+
     if (auto funcOp = dyn_cast<LLVMFuncOp>(op)) {
-      funcOp.setFunctionEntryCount(entryCount->getZExtValue());
+      funcOp.setFunctionEntryCount(*entryCountValue);
+      if (isSynthetic)
+        funcOp.setFunctionEntryCountSynthetic(true);
+      if (!importGUIDValues.empty()) {
+        SmallVector<int64_t> importGUIDs;
+        importGUIDs.reserve(importGUIDValues.size());
+        for (uint64_t guid : importGUIDValues)
+          importGUIDs.push_back(static_cast<int64_t>(guid));
+        funcOp.setFunctionEntryCountImportsAttr(
+            DenseI64ArrayAttr::get(builder.getContext(), importGUIDs));
+      }
       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 442e6e16e955e..1dbed59dae278 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"
@@ -2003,9 +2004,25 @@ LogicalResult ModuleTranslation::convertFunctionSignatures() {
     // Convert function kernel attributes to metadata.
     convertFunctionKernelAttributes(function, llvmFunc, *this);
 
-    // Convert function_entry_count attribute to metadata.
-    if (std::optional<uint64_t> entryCount = function.getFunctionEntryCount())
-      llvmFunc->setEntryCount(entryCount.value());
+    // Convert function_entry_count attributes to metadata.
+    if (std::optional<uint64_t> entryCount = function.getFunctionEntryCount()) {
+      llvm::Function::ProfileCount profileCount(
+          entryCount.value(), function.getFunctionEntryCountSynthetic()
+                                  ? llvm::Function::PCT_Synthetic
+                                  : llvm::Function::PCT_Real);
+      std::optional<llvm::DenseSet<llvm::GlobalValue::GUID>> importGUIDs;
+      if (DenseI64ArrayAttr imports =
+              function.getFunctionEntryCountImportsAttr()) {
+        importGUIDs.emplace();
+        for (int64_t guid : imports.asArrayRef()) {
+          // The MLIR attribute preserves the unsigned GUID bit pattern in a
+          // signed i64 element.
+          importGUIDs->insert(static_cast<uint64_t>(guid));
+        }
+      }
+      llvmFunc->setEntryCount(profileCount,
+                              importGUIDs ? &*importGUIDs : nullptr);
+    }
 
     // Convert result attributes.
     if (ArrayAttr allResultAttrs = function.getAllResultAttrs()) {
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index d5ea5c8de862e..6acc7a410a5ce 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -44,6 +44,61 @@ llvm.mlir.global_dtors dtors = [@dtor], priorities = [0 : i32], data = [0 : i32]
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// expected-error at +1{{requires function_entry_count when function_entry_count_synthetic is set}}
+llvm.func @function_entry_count_synthetic_requires_count() attributes {
+  function_entry_count_synthetic
+}
+
+// -----
+
+// expected-error at +1{{requires function_entry_count when function_entry_count_imports is set}}
+llvm.func @function_entry_count_imports_requires_count() attributes {
+  function_entry_count_imports = array<i64: 1234>
+}
+
+// -----
+
+// expected-error at +1{{does not support function_entry_count_imports with function_entry_count_synthetic}}
+llvm.func @function_entry_count_imports_requires_real_count() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64: 1234>,
+  function_entry_count_synthetic
+}
+
+// -----
+
+// expected-error at +1{{requires function_entry_count_imports to be non-empty when set}}
+llvm.func @function_entry_count_imports_non_empty() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64>
+}
+
+// -----
+
+// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
+llvm.func @function_entry_count_imports_sorted() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64: 9, 4>
+}
+
+// -----
+
+// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
+llvm.func @function_entry_count_imports_unique() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64: 4, 4>
+}
+
+// -----
+
+// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
+llvm.func @function_entry_count_imports_unsigned_order() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64: -1, 4>
+}
+
+// -----
+
 // Check that parser errors are properly produced and do not crash the compiler.
 
 // -----
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
index 5d664519e100c..e5584ee6d3824 100644
--- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll
+++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
@@ -180,6 +180,84 @@ define void @entry_count() !prof !1 {
 
 ; // -----
 
+; CHECK-LABEL: @synthetic_entry_count
+; CHECK-SAME:  attributes {function_entry_count = 7 : i64
+; CHECK-SAME:  function_entry_count_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 = 7 : i64
+; CHECK-SAME:  function_entry_count_imports = array<i64: 4, 1234, -1>
+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-NOT: function_entry_count
+; expected-warning @unknown {{unhandled function metadata}}
+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. LLVM's semantic getEntryCount()
+; treats real uint64_t(-1) as unknown, but translation preserves the metadata.
+; CHECK-LABEL: @entry_count_negative_count
+; CHECK-SAME:  attributes {function_entry_count = -1 : i64}
+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/Import/function-entry-count-roundtrip.ll b/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll
new file mode 100644
index 0000000000000..2d7a0cec4bd4e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll
@@ -0,0 +1,21 @@
+; RUN: mlir-translate -import-llvm %s | mlir-translate -mlir-to-llvmir | FileCheck %s
+
+define void @synthetic() !prof !0 {
+  ret void
+}
+
+define void @with_import_guid() !prof !1 {
+  ret void
+}
+
+!0 = !{!"synthetic_function_entry_count", i64 7}
+!1 = !{!"function_entry_count", i64 7, i64 1234}
+
+; CHECK: define void @synthetic()
+; CHECK-SAME: !prof ![[SYNTH:[0-9]+]]
+
+; CHECK: define void @with_import_guid()
+; CHECK-SAME: !prof ![[IMPORTS:[0-9]+]]
+
+; CHECK-DAG: ![[SYNTH]] = !{!"synthetic_function_entry_count", i64 7}
+; CHECK-DAG: ![[IMPORTS]] = !{!"function_entry_count", i64 7, i64 1234}
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 5eca8f19154a1..7a9aedb262333 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1893,12 +1893,50 @@ llvm.func @my_allocator(i64) attributes {passthrough = [["allocsize", "429496729
 // -----
 
 // CHECK-LABEL: @functionEntryCount
-// CHECK-SAME: !prof ![[PROF_ID:[0-9]*]]
+// CHECK-SAME: !prof ![[PROF_ID:[0-9]+]]
 llvm.func @functionEntryCount() attributes {function_entry_count = 4242 : i64} {
   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 = 7 : i64,
+  function_entry_count_synthetic
+} {
+  llvm.return
+}
+
+// CHECK-DAG: ![[SYNTH_PROF_ID]] = !{!"synthetic_function_entry_count", i64 7}
+
+// -----
+
+// CHECK-LABEL: @functionEntryCountWithImports
+// CHECK-SAME: !prof ![[IMPORTS_PROF_ID:[0-9]+]]
+llvm.func @functionEntryCountWithImports() attributes {
+  function_entry_count = 7 : i64,
+  function_entry_count_imports = array<i64: 4, 1234, -1>
+} {
+  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 = -1 : i64
+} {
+  llvm.return
+}
+
+// CHECK-DAG: ![[NEG_PROF_ID]] = !{!"function_entry_count", i64 -1}
 
 // -----
 

>From 6f14005d0391f56ee21bdc08fca2b234e916eda0 Mon Sep 17 00:00:00 2001
From: "Cyne Jarvis J. Zarceno" <cynejarviszarceno at gmail.com>
Date: Mon, 22 Jun 2026 04:16:45 +0800
Subject: [PATCH 2/2] [mlir][llvm] Refine function entry count metadata

---
 .../mlir/Dialect/LLVMIR/LLVMAttrDefs.td       | 20 +++++++
 mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td | 18 ++++++
 mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td   | 14 ++---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    | 31 +----------
 .../LLVMIR/LLVMIRToLLVMTranslation.cpp        | 32 ++---------
 mlir/lib/Target/LLVMIR/ModuleTranslation.cpp  | 20 +++----
 mlir/test/Dialect/LLVMIR/invalid.mlir         | 55 -------------------
 .../LLVMIR/Import/function-attributes.ll      | 13 ++---
 .../Import/function-entry-count-roundtrip.ll  | 13 ++++-
 mlir/test/Target/LLVMIR/llvmir.mlir           | 24 ++++++--
 10 files changed, 91 insertions(+), 149 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
index 9474afb897119..b571f563609c3 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
@@ -96,6 +96,26 @@ 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
+    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,
+                        "ProfileCountType":$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..cb69c53e60d8e 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMEnums.td
@@ -83,6 +83,24 @@ def AsmATTOrIntel : LLVM_EnumAttr<
   let cppNamespace = "::mlir::LLVM";
 }
 
+//===----------------------------------------------------------------------===//
+// ProfileCountType
+//===----------------------------------------------------------------------===//
+
+def ProfileCountReal : LLVM_EnumAttrCase<
+  /*string cppSym=*/"Real", /*string irSym=*/"real",
+  /*string llvmSym=*/"PCT_Real", /*int val=*/0>;
+def ProfileCountSynthetic : LLVM_EnumAttrCase<
+  /*string cppSym=*/"Synthetic", /*string irSym=*/"synthetic",
+  /*string llvmSym=*/"PCT_Synthetic", /*int val=*/1>;
+def ProfileCountType : LLVM_EnumAttr<
+  /*string name=*/"ProfileCountType",
+  /*string llvmName=*/"::llvm::Function::ProfileCountType",
+  /*string description=*/"real or synthetic function entry count",
+  /*list<LLVM_EnumAttrCase> 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 557e9881bd921..4bdcf5f6c1cc5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -2020,13 +2020,9 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     ```
 
     The `function_entry_count` attribute models function-level `!prof`
-    entry-count metadata. The `function_entry_count_synthetic` unit attribute
-    selects the `"synthetic_function_entry_count"` metadata label. The
-    `function_entry_count_imports` attribute represents the non-empty set of
-    trailing import GUID operands on real `"function_entry_count"` metadata.
-    Those GUID operands are unsigned 64-bit LLVM GUID bit patterns stored in a
-    signed i64 array attribute in unsigned sorted-unique order. This attribute
-    must not be combined with `function_entry_count_synthetic`.
+    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
@@ -2043,9 +2039,7 @@ def LLVM_LLVMFuncOp : LLVM_Op<"func", [
     OptionalAttr<ArrayAttr>:$passthrough,
     OptionalAttr<DictArrayAttr>:$arg_attrs,
     OptionalAttr<DictArrayAttr>:$res_attrs,
-    OptionalAttr<I64Attr>:$function_entry_count,
-    UnitAttr:$function_entry_count_synthetic,
-    OptionalAttr<DenseI64ArrayAttr>:$function_entry_count_imports,
+    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 1898ad752f80c..f3dc00e177b0f 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -3045,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()) {
@@ -3267,33 +3269,6 @@ LogicalResult LLVMFuncOp::verify() {
   if (failed(verifyComdat(*this, getComdat())))
     return failure();
 
-  if (!getFunctionEntryCountAttr()) {
-    if (getFunctionEntryCountSynthetic())
-      return emitOpError() << "requires function_entry_count when "
-                              "function_entry_count_synthetic is set";
-    if (getFunctionEntryCountImportsAttr())
-      return emitOpError() << "requires function_entry_count when "
-                              "function_entry_count_imports is set";
-  }
-  if (DenseI64ArrayAttr imports = getFunctionEntryCountImportsAttr()) {
-    if (getFunctionEntryCountSynthetic())
-      return emitOpError() << "does not support function_entry_count_imports "
-                              "with function_entry_count_synthetic";
-
-    ArrayRef<int64_t> values = imports.asArrayRef();
-    if (values.empty())
-      return emitOpError() << "requires function_entry_count_imports to be "
-                              "non-empty when set";
-
-    for (auto [previous, current] :
-         llvm::zip_equal(values.drop_back(), values.drop_front())) {
-      if (static_cast<uint64_t>(previous) >= static_cast<uint64_t>(current))
-        return emitOpError()
-               << "requires function_entry_count_imports to be sorted and "
-                  "unique by unsigned GUID value";
-    }
-  }
-
   if (isExternal()) {
     if (getLinkage() != LLVM::Linkage::External &&
         getLinkage() != LLVM::Linkage::ExternWeak)
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
index d15b940ca342c..549df5fecc9e4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp
@@ -17,7 +17,6 @@
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/ModuleImport.h"
 
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/InlineAsm.h"
@@ -26,9 +25,6 @@
 #include "llvm/IR/MemoryModelRelaxationAnnotations.h"
 #include <optional>
 
-#include <algorithm>
-#include <optional>
-
 using namespace mlir;
 using namespace mlir::LLVM;
 using namespace mlir::LLVM::detail;
@@ -136,12 +132,6 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
     bool isSynthetic =
         profName == llvm::MDProfLabels::SyntheticFunctionEntryCount;
 
-    // LLVM's semantic import-GUID API only reads trailing GUID operands from
-    // "function_entry_count" metadata. Do not model trailing operands on
-    // "synthetic_function_entry_count" as import GUIDs in MLIR.
-    if (isSynthetic && node->getNumOperands() > 2)
-      return failure();
-
     std::optional<uint64_t> entryCountValue =
         getUInt64Metadata(node->getOperand(1));
     if (!entryCountValue)
@@ -157,25 +147,11 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
       importGUIDValues.push_back(*guidValue);
     }
 
-    // Import GUIDs are semantically a set in LLVM. Canonicalize them as
-    // unsigned sorted-unique values before storing the bit patterns in MLIR.
-    llvm::sort(importGUIDValues);
-    importGUIDValues.erase(
-        std::unique(importGUIDValues.begin(), importGUIDValues.end()),
-        importGUIDValues.end());
-
     if (auto funcOp = dyn_cast<LLVMFuncOp>(op)) {
-      funcOp.setFunctionEntryCount(*entryCountValue);
-      if (isSynthetic)
-        funcOp.setFunctionEntryCountSynthetic(true);
-      if (!importGUIDValues.empty()) {
-        SmallVector<int64_t> importGUIDs;
-        importGUIDs.reserve(importGUIDValues.size());
-        for (uint64_t guid : importGUIDValues)
-          importGUIDs.push_back(static_cast<int64_t>(guid));
-        funcOp.setFunctionEntryCountImportsAttr(
-            DenseI64ArrayAttr::get(builder.getContext(), importGUIDs));
-      }
+      funcOp.setFunctionEntryCountAttr(FunctionEntryCountAttr::get(
+          builder.getContext(), *entryCountValue,
+          isSynthetic ? ProfileCountType::Synthetic : ProfileCountType::Real,
+          importGUIDValues));
       return success();
     }
     return op->emitWarning()
diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
index 1dbed59dae278..2fec39b313a1d 100644
--- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
@@ -2004,21 +2004,17 @@ LogicalResult ModuleTranslation::convertFunctionSignatures() {
     // Convert function kernel attributes to metadata.
     convertFunctionKernelAttributes(function, llvmFunc, *this);
 
-    // Convert function_entry_count attributes to metadata.
-    if (std::optional<uint64_t> entryCount = function.getFunctionEntryCount()) {
+    // Convert function_entry_count attribute to metadata.
+    if (FunctionEntryCountAttr entryCount =
+            function.getFunctionEntryCountAttr()) {
       llvm::Function::ProfileCount profileCount(
-          entryCount.value(), function.getFunctionEntryCountSynthetic()
-                                  ? llvm::Function::PCT_Synthetic
-                                  : llvm::Function::PCT_Real);
+          entryCount.getEntryCount(),
+          convertProfileCountTypeToLLVM(entryCount.getCountType()));
       std::optional<llvm::DenseSet<llvm::GlobalValue::GUID>> importGUIDs;
-      if (DenseI64ArrayAttr imports =
-              function.getFunctionEntryCountImportsAttr()) {
+      ArrayRef<uint64_t> imports = entryCount.getImports();
+      if (!imports.empty()) {
         importGUIDs.emplace();
-        for (int64_t guid : imports.asArrayRef()) {
-          // The MLIR attribute preserves the unsigned GUID bit pattern in a
-          // signed i64 element.
-          importGUIDs->insert(static_cast<uint64_t>(guid));
-        }
+        importGUIDs->insert(imports.begin(), imports.end());
       }
       llvmFunc->setEntryCount(profileCount,
                               importGUIDs ? &*importGUIDs : nullptr);
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index 6acc7a410a5ce..d5ea5c8de862e 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -44,61 +44,6 @@ llvm.mlir.global_dtors dtors = [@dtor], priorities = [0 : i32], data = [0 : i32]
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// expected-error at +1{{requires function_entry_count when function_entry_count_synthetic is set}}
-llvm.func @function_entry_count_synthetic_requires_count() attributes {
-  function_entry_count_synthetic
-}
-
-// -----
-
-// expected-error at +1{{requires function_entry_count when function_entry_count_imports is set}}
-llvm.func @function_entry_count_imports_requires_count() attributes {
-  function_entry_count_imports = array<i64: 1234>
-}
-
-// -----
-
-// expected-error at +1{{does not support function_entry_count_imports with function_entry_count_synthetic}}
-llvm.func @function_entry_count_imports_requires_real_count() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_imports = array<i64: 1234>,
-  function_entry_count_synthetic
-}
-
-// -----
-
-// expected-error at +1{{requires function_entry_count_imports to be non-empty when set}}
-llvm.func @function_entry_count_imports_non_empty() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_imports = array<i64>
-}
-
-// -----
-
-// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
-llvm.func @function_entry_count_imports_sorted() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_imports = array<i64: 9, 4>
-}
-
-// -----
-
-// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
-llvm.func @function_entry_count_imports_unique() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_imports = array<i64: 4, 4>
-}
-
-// -----
-
-// expected-error at +1{{requires function_entry_count_imports to be sorted and unique by unsigned GUID value}}
-llvm.func @function_entry_count_imports_unsigned_order() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_imports = array<i64: -1, 4>
-}
-
-// -----
-
 // Check that parser errors are properly produced and do not crash the compiler.
 
 // -----
diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll
index e5584ee6d3824..ae8ee3fd07709 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, count_type = real>}
 define void @entry_count() !prof !1 {
   ret void
 }
@@ -181,8 +181,7 @@ define void @entry_count() !prof !1 {
 ; // -----
 
 ; CHECK-LABEL: @synthetic_entry_count
-; CHECK-SAME:  attributes {function_entry_count = 7 : i64
-; CHECK-SAME:  function_entry_count_synthetic
+; 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
 }
@@ -192,8 +191,7 @@ define void @synthetic_entry_count() !prof !2 {
 ; // -----
 
 ; CHECK-LABEL: @entry_count_imports
-; CHECK-SAME:  attributes {function_entry_count = 7 : i64
-; CHECK-SAME:  function_entry_count_imports = array<i64: 4, 1234, -1>
+; CHECK-SAME:  attributes {function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = real, imports = [1234, 18446744073709551615, 4, 1234]>}
 define void @entry_count_imports() !prof !3 {
   ret void
 }
@@ -203,8 +201,7 @@ define void @entry_count_imports() !prof !3 {
 ; // -----
 
 ; CHECK-LABEL: @synthetic_entry_count_imports
-; CHECK-NOT: function_entry_count
-; expected-warning @unknown {{unhandled function metadata}}
+; 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
 }
@@ -249,7 +246,7 @@ define void @entry_count_too_wide_import() !prof !7 {
 ; Preserve the raw i64 metadata bit pattern. LLVM's semantic getEntryCount()
 ; treats real uint64_t(-1) as unknown, but translation preserves the metadata.
 ; CHECK-LABEL: @entry_count_negative_count
-; CHECK-SAME:  attributes {function_entry_count = -1 : i64}
+; CHECK-SAME:  attributes {function_entry_count = #llvm.function_entry_count<entry_count = 18446744073709551615, count_type = real>}
 define void @entry_count_negative_count() !prof !8 {
   ret void
 }
diff --git a/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll b/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll
index 2d7a0cec4bd4e..125684902f398 100644
--- a/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll
+++ b/mlir/test/Target/LLVMIR/Import/function-entry-count-roundtrip.ll
@@ -8,8 +8,13 @@ define void @with_import_guid() !prof !1 {
   ret void
 }
 
+define void @synthetic_with_import_guid() !prof !2 {
+  ret void
+}
+
 !0 = !{!"synthetic_function_entry_count", i64 7}
-!1 = !{!"function_entry_count", i64 7, i64 1234}
+!1 = !{!"function_entry_count", i64 7, i64 1234, i64 4, i64 1234}
+!2 = !{!"synthetic_function_entry_count", i64 7, i64 1234}
 
 ; CHECK: define void @synthetic()
 ; CHECK-SAME: !prof ![[SYNTH:[0-9]+]]
@@ -17,5 +22,9 @@ define void @with_import_guid() !prof !1 {
 ; CHECK: define void @with_import_guid()
 ; CHECK-SAME: !prof ![[IMPORTS:[0-9]+]]
 
+; CHECK: define void @synthetic_with_import_guid()
+; CHECK-SAME: !prof ![[SYNTH_IMPORTS:[0-9]+]]
+
 ; CHECK-DAG: ![[SYNTH]] = !{!"synthetic_function_entry_count", i64 7}
-; CHECK-DAG: ![[IMPORTS]] = !{!"function_entry_count", i64 7, i64 1234}
+; CHECK-DAG: ![[IMPORTS]] = !{!"function_entry_count", i64 7, i64 4, i64 1234}
+; CHECK-DAG: ![[SYNTH_IMPORTS]] = !{!"synthetic_function_entry_count", i64 7, i64 1234}
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 7a9aedb262333..1733143b16c27 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1894,7 +1894,9 @@ 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} {
+llvm.func @functionEntryCount() attributes {
+  function_entry_count = #llvm.function_entry_count<entry_count = 4242, count_type = real>
+} {
   llvm.return
 }
 
@@ -1905,8 +1907,7 @@ llvm.func @functionEntryCount() attributes {function_entry_count = 4242 : i64} {
 // CHECK-LABEL: @syntheticFunctionEntryCount
 // CHECK-SAME: !prof ![[SYNTH_PROF_ID:[0-9]+]]
 llvm.func @syntheticFunctionEntryCount() attributes {
-  function_entry_count = 7 : i64,
-  function_entry_count_synthetic
+  function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = synthetic>
 } {
   llvm.return
 }
@@ -1915,11 +1916,22 @@ llvm.func @syntheticFunctionEntryCount() attributes {
 
 // -----
 
+// 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 = 7 : i64,
-  function_entry_count_imports = array<i64: 4, 1234, -1>
+  function_entry_count = #llvm.function_entry_count<entry_count = 7, count_type = real, imports = [1234, 4, 18446744073709551615, 1234]>
 } {
   llvm.return
 }
@@ -1931,7 +1943,7 @@ llvm.func @functionEntryCountWithImports() attributes {
 // CHECK-LABEL: @functionEntryCountNegativeCount
 // CHECK-SAME: !prof ![[NEG_PROF_ID:[0-9]+]]
 llvm.func @functionEntryCountNegativeCount() attributes {
-  function_entry_count = -1 : i64
+  function_entry_count = #llvm.function_entry_count<entry_count = 18446744073709551615, count_type = real>
 } {
   llvm.return
 }



More information about the Mlir-commits mailing list