[mlir] Add function_entry_count to LLVMFuncOp

[Dinistro]

Summary

This commit introduces the function_entry_count metadata field to the
LLVMFuncOp and adds both the corresponding import and export
funtionalities.
The import of the function metadata uses the same infrastructure as the
instruction metadata, i.e., it dispatches through a dialect interface.

[gysit]

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1996

nit: functionEntryCounts -> functionEntryCount

2010

I think there is a getFunctionEntryCountAttrName name? Also should't this be singular without an s?

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

87

nit: the comment needs an update.

92

Should we use the attribute name here and below?

93–94

nit: "TODO support function entry count metadata with a global identifier field." That way it is clear there is no specific reason why this is not imported.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1191–1200

Any chance this could be factored out in a helper method assuming instructions and functions are both values?

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

893

nit: auto -> Int64Attr?

[Dinistro]

address review comments

[Dinistro]

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

92

I do not quite follow. This function does not belong to a specific operation, so one has no access to the attribute name accessors

93–94

I'm not sure if these GUIDs would be preserved when exporting back to LLVM. But I can add it as a TODO and we can later on check if this is indeed the case or not.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1191–1200

I tired, but unfortunately Value::getAllMetadata is protected and cannot be accessed. Furthermore, llvm::Value has no location, so such a function would take up to 4 arguments.

[Dinistro]

drop unnecessary comment

[gysit]

LGTM

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

92

Never mind I thought the getFunctionEntryCountAttrName is a static function.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1191–1200

ok let's keep it then.

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

let arguments = (ins
     TypeAttrOf<LLVM_FunctionType>:$function_type,
     DefaultValuedAttr<Linkage, "Linkage::External">:$linkage,
     UnitAttr:$dso_local,
     DefaultValuedAttr<CConv, "CConv::C">:$CConv,
     OptionalAttr<FlatSymbolRefAttr>:$personality,
     OptionalAttr<StrAttr>:$garbageCollector,
     OptionalAttr<ArrayAttr>:$passthrough,
     OptionalAttr<DictArrayAttr>:$arg_attrs,
-    OptionalAttr<DictArrayAttr>:$res_attrs
+    OptionalAttr<DictArrayAttr>:$res_attrs,
+    OptionalAttr<I64Attr>:$function_entry_count
   );
 
   let regions = (region AnyRegion:$body);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "StringRef":$name, "Type":$type,
       CArg<"Linkage", "Linkage::External">:$linkage,
       CArg<"bool", "false">:$dsoLocal,
       CArg<"CConv", "CConv::C">:$cconv,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs,
-      CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs)>
+      CArg<"ArrayRef<DictionaryAttr>", "{}">:$argAttrs,
+      CArg<"Optional<uint64_t>", "{}">:$functionEntryCount)>
   ];
 
   let extraClassDeclaration = [{
     // Add an entry block to an empty function, and set up the block arguments
     // to match the signature of the function.
     Block *addEntryBlock();
 
     bool isVarArg() { return getFunctionType().isVarArg(); }

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

for (unsigned i = 0, e = type.getNumParams(); i < e; ++i)
     entry->addArgument(type.getParamType(i), getLoc());
   return entry;
 }
 
 void LLVMFuncOp::build(OpBuilder &builder, OperationState &result,
                        StringRef name, Type type, LLVM::Linkage linkage,
                        bool dsoLocal, CConv cconv,
                        ArrayRef<NamedAttribute> attrs,
-                       ArrayRef<DictionaryAttr> argAttrs) {
+                       ArrayRef<DictionaryAttr> argAttrs,
+                       Optional<uint64_t> functionEntryCount) {
   result.addRegion();
   result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder.getStringAttr(name));
   result.addAttribute(getFunctionTypeAttrName(result.name),
                       TypeAttr::get(type));
   result.addAttribute(getLinkageAttrName(result.name),
                       LinkageAttr::get(builder.getContext(), linkage));
   result.addAttribute(getCConvAttrName(result.name),
                       CConvAttr::get(builder.getContext(), cconv));
   result.attributes.append(attrs.begin(), attrs.end());
   if (dsoLocal)
-    result.addAttribute("dso_local", builder.getUnitAttr());
+    result.addAttribute(getDsoLocalAttrName(result.name),
+                        builder.getUnitAttr());
+  if (functionEntryCount)
+    result.addAttribute(getFunctionEntryCountAttrName(result.name),
+                        builder.getI64IntegerAttr(functionEntryCount.value()));
   if (argAttrs.empty())
     return;
 
   assert(type.cast<LLVMFunctionType>().getNumParams() == argAttrs.size() &&
          "expected as many argument attribute lists as arguments");
   function_interface_impl::addArgAndResultAttrs(
       builder, result, argAttrs, /*resultAttrs=*/std::nullopt,
       getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp

 /// failure otherwise.
 static LogicalResult setProfilingAttrs(OpBuilder &builder, llvm::MDNode *node,
                                        Operation *op,
                                        LLVM::ModuleImport &moduleImport) {
   // Return success for empty metadata nodes since there is nothing to import.
   if (!node->getNumOperands())
     return success();
 
-  // Return failure for non-"branch_weights" metadata.
+  // Return failure for empty metadata.
   auto *name = dyn_cast<llvm::MDString>(node->getOperand(0));
-  if (!name || !name->getString().equals("branch_weights"))
+  if (!name)
+    return failure();
+
+  if (name->getString().equals("function_entry_count")) {
+    // TODO support function entry count metadata with GUID fields.
+    if (node->getNumOperands() != 2)
+      return failure();
+
+    llvm::ConstantInt *entryCount =
+        llvm::mdconst::extract<llvm::ConstantInt>(node->getOperand(1));
+    if (auto funcOp = dyn_cast<LLVMFuncOp>(op)) {
+      funcOp.setFunctionEntryCount(entryCount->getZExtValue());
+      return success();
+    }
+
+    return failure();
+  }
+
+  // Return failure for unsupported metadata.
+  if (!name->getString().equals("branch_weights"))
     return failure();
 
   // Copy the branch weights to an array.
   SmallVector<int32_t> branchWeights;
   branchWeights.reserve(node->getNumOperands() - 1);
   for (unsigned i = 1, e = node->getNumOperands(); i != e; ++i) {
     llvm::ConstantInt *branchWeight =
         llvm::mdconst::extract<llvm::ConstantInt>(node->getOperand(i));

mlir/lib/Target/LLVMIR/ModuleImport.cpp

emitWarning(UnknownLoc::get(context),
                 "could not deduce personality, skipping it");
 
   if (func->hasGC())
     funcOp.setGarbageCollector(StringRef(func->getGC()));
 
   // Handle Function attributes.
   processFunctionAttributes(func, funcOp);
 
+  // Convert non-debug metadata by using the dialect interface.
+  SmallVector<std::pair<unsigned, llvm::MDNode *>> allMetadata;
+  func->getAllMetadata(allMetadata);
+  for (auto &[kind, node] : allMetadata) {
+    if (!iface.isConvertibleMetadata(kind))
+      continue;
+    if (failed(iface.setMetadataAttrs(builder, kind, node, funcOp, *this))) {
+      emitWarning(funcOp->getLoc())
+          << "unhandled function metadata (" << kind << ") " << diag(*func);
+    }
+  }
+
   if (func->isDeclaration())
     return success();
 
   // Eagerly create all blocks.
   for (llvm::BasicBlock &bb : *func) {
     Block *block =
         builder.createBlock(&funcOp.getBody(), funcOp.getBody().end());
     mapBlock(&bb, block);

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

for (auto function : getModuleBody(mlirModule).getOps<LLVMFuncOp>()) {
     llvmFunc->setLinkage(convertLinkageToLLVM(function.getLinkage()));
     mapFunction(function.getName(), llvmFunc);
     addRuntimePreemptionSpecifier(function.getDsoLocal(), llvmFunc);
 
     // Convert function attributes.
     if (function->getAttrOfType<UnitAttr>(LLVMDialect::getReadnoneAttrName()))
       llvmFunc->setDoesNotAccessMemory();
 
+    // Convert function_entry_count attribute to metadata.
+    if (std::optional<uint64_t> entryCount = function.getFunctionEntryCount())
+      llvmFunc->setEntryCount(entryCount.value());
+
     // Convert result attributes.
     if (ArrayAttr allResultAttrs = function.getAllResultAttrs()) {
       llvm::AttrBuilder retAttrs(llvmFunc->getContext());
       DictionaryAttr resultAttrs = allResultAttrs[0].cast<DictionaryAttr>();
       for (const NamedAttribute &attr : resultAttrs) {
         StringAttr name = attr.getName();
         if (name == LLVMDialect::getAlignAttrName()) {
           auto alignAmount = attr.getValue().cast<IntegerAttr>();

mlir/test/Target/LLVMIR/Import/function-attributes.ll

 ; CHECK-SAME:  !llvm.ptr {llvm.inalloca = i64}
 define void @func_arg_attrs(
     ptr byval(i64) %arg0,
     ptr byref(i64) %arg1,
     ptr sret(i64) %arg2,
     ptr inalloca(i64) %arg3) {
   ret void
 }
+
+; // -----
+
+; CHECK-LABEL: @entry_count
+; CHECK-SAME:  attributes {function_entry_count = 4242 : i64}
+define void @entry_count() !prof !1 {
+  ret void
+}
+
+!1 = !{!"function_entry_count", i64 4242}

mlir/test/Target/LLVMIR/llvmir.mlir

 // CHECK: attributes #[[ATTR_GROUP]] = {
 // CHECK-DAG: noinline
 // CHECK-DAG: alignstack=4
 // CHECK-DAG: null_pointer_is_valid
 // CHECK-DAG: "foo"="bar"
 
 // -----
 
+// CHECK-LABEL: @functionEntryCount
+// 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-LABEL: @constant_bf16
 llvm.func @constant_bf16() -> bf16 {
   %0 = llvm.mlir.constant(1.000000e+01 : bf16) : bf16
   llvm.return %0 : bf16
 }
 
 // CHECK: ret bfloat 0xR4120