[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

86–87

nit: the comment needs an update.

91

Should we use the attribute name here and below?

92–93

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

1188–1197

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

91

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

92–93

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

1188–1197

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

91

Never mind I thought the getFunctionEntryCountAttrName is a static function.

mlir/lib/Target/LLVMIR/ModuleImport.cpp

1188–1197

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

 /// conversion to an MLIR profiling attribute exists and succeeds. Returns
 /// 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.
   auto *name = dyn_cast<llvm::MDString>(node->getOperand(0));
-  if (!name || !name->getString().equals("branch_weights"))
+  if (!name)
     return failure();
 
-  // Copy the branch weights to an array.
+  // Handle function entry count metadata.
+  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();
+  }
+
+  if (!name->getString().equals("branch_weights"))
+    return failure();
+
+  // Handle branch weights metadata.
   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));
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 

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