diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td @@ -1330,7 +1330,8 @@ OptionalAttr:$garbageCollector, OptionalAttr:$passthrough, OptionalAttr:$arg_attrs, - OptionalAttr:$res_attrs + OptionalAttr:$res_attrs, + OptionalAttr:$function_entry_count ); let regions = (region AnyRegion:$body); @@ -1343,7 +1344,8 @@ CArg<"bool", "false">:$dsoLocal, CArg<"CConv", "CConv::C">:$cconv, CArg<"ArrayRef", "{}">:$attrs, - CArg<"ArrayRef", "{}">:$argAttrs)> + CArg<"ArrayRef", "{}">:$argAttrs, + CArg<"Optional", "{}">:$functionEntryCount)> ]; let extraClassDeclaration = [{ diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp @@ -1992,7 +1992,8 @@ StringRef name, Type type, LLVM::Linkage linkage, bool dsoLocal, CConv cconv, ArrayRef attrs, - ArrayRef argAttrs) { + ArrayRef argAttrs, + Optional functionEntryCounts) { result.addRegion(); result.addAttribute(SymbolTable::getSymbolAttrName(), builder.getStringAttr(name)); @@ -2005,6 +2006,9 @@ result.attributes.append(attrs.begin(), attrs.end()); if (dsoLocal) result.addAttribute("dso_local", builder.getUnitAttr()); + if (functionEntryCounts) + result.addAttribute("function_entry_counts", + builder.getI64IntegerAttr(functionEntryCounts.value())); if (argAttrs.empty()) return; diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMIRToLLVMTranslation.cpp @@ -86,7 +86,26 @@ // Return failure for non-"branch_weights" metadata. auto *name = dyn_cast(node->getOperand(0)); - if (!name || !name->getString().equals("branch_weights")) + if (!name) + return failure(); + + if (name->getString().equals("function_entry_count")) { + // MLIR has no notion of GUIDs thus we have to ignore metadata with GUIDs + // attached. + if (node->getNumOperands() != 2) + return failure(); + + llvm::ConstantInt *entryCount = + llvm::mdconst::extract(node->getOperand(1)); + if (auto funcOp = dyn_cast(op)) { + funcOp.setFunctionEntryCount(entryCount->getZExtValue()); + return success(); + } + + return failure(); + } + + if (!name->getString().equals("branch_weights")) return failure(); // Copy the branch weights to an array. diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -1187,6 +1187,18 @@ // Handle Function attributes. processFunctionAttributes(func, funcOp); + // Convert non-debug metadata by using the dialect interface. + SmallVector> 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(); diff --git a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp --- a/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleTranslation.cpp @@ -889,6 +889,10 @@ if (function->getAttrOfType(LLVMDialect::getReadnoneAttrName())) llvmFunc->setDoesNotAccessMemory(); + // Convert function_entry_count attribute to metadata. + if (auto entryCount = function.getFunctionEntryCount()) + llvmFunc->setEntryCount(entryCount.value()); + // Convert result attributes. if (ArrayAttr allResultAttrs = function.getAllResultAttrs()) { llvm::AttrBuilder retAttrs(llvmFunc->getContext()); diff --git a/mlir/test/Target/LLVMIR/Import/function-attributes.ll b/mlir/test/Target/LLVMIR/Import/function-attributes.ll --- a/mlir/test/Target/LLVMIR/Import/function-attributes.ll +++ b/mlir/test/Target/LLVMIR/Import/function-attributes.ll @@ -38,3 +38,13 @@ 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} diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir --- a/mlir/test/Target/LLVMIR/llvmir.mlir +++ b/mlir/test/Target/LLVMIR/llvmir.mlir @@ -1545,6 +1545,16 @@ // ----- +// 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