[MLIR][LLVM] Add CG Profile module flags support

mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td

@@ -1322,7 +1322,7 @@ def LLVM_DereferenceableAttr : LLVM_Attr<"Dereferenceable", "dereferenceable"> {
 }
 
 //===----------------------------------------------------------------------===//
-// ModuleFlagAttr
+// ModuleFlagAttr & related
 //===----------------------------------------------------------------------===//
 
 def ModuleFlagAttr
@@ -1332,14 +1332,22 @@ def ModuleFlagAttr
     Represents a single entry of llvm.module.flags metadata
     (llvm::Module::ModuleFlagEntry in LLVM). The first element is a behavior
     flag described by `ModFlagBehaviorAttr`, the second is a string ID
-    and third is the value of the flag. Current supported types of values:
-      - Integer constants
-      - Strings
+    and third is the value of the flag. Supported keys and values include:
+      - Arbitrary `key`s holding integer constants or strings.
+      - Domain specific keys (e.g "CG Profile"), holding lists of supported
+        module flag values (e.g. `llvm.cgprofile_entry`).
 
     Example:
     ```mlir
-      #llvm.mlir.module_flag<error, "wchar_size", 4>
-      #llvm.mlir.module_flag<error, "probe-stack", "inline-asm">
+      llvm.module_flags [
+          #llvm.mlir.module_flag<error, "wchar_size", 4>,
+          #llvm.mlir.module_flag<error, "probe-stack", "inline-asm">,
+          #llvm.mlir.module_flag<append, "CG Profile", [
+            #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+            #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+            #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+          ]
+      >]
     ```
   }];
   let parameters = (ins "ModFlagBehavior":$behavior,
@@ -1349,6 +1357,25 @@ def ModuleFlagAttr
   let genVerifyDecl = 1;
 }
 
+def ModuleFlagCGProfileEntryAttr
+    : LLVM_Attr<"ModuleFlagCGProfileEntry", "cgprofile_entry"> {
+  let summary = "CG profile module flag entry";
+  let description = [{
+    Describes a single entry for a CG profile module flag. Example:
+    ```mlir
+      llvm.module_flags [
+        #llvm.mlir.module_flag<append, "CG Profile",
+          [#llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+           ...
+          ]>]
+    ```
+  }];
+  let parameters = (ins "FlatSymbolRefAttr":$from,
+                        "FlatSymbolRefAttr":$to,
+                        "uint64_t":$count);
+  let assemblyFormat = "`<` struct(params) `>`";
+}
+
 //===----------------------------------------------------------------------===//
 // LLVM_DependentLibrariesAttr
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.td

@@ -88,6 +88,11 @@ def LLVM_Dialect : Dialect {
       return "llvm.dependent_libraries";
     }
 
+    /// Names of known llvm module flag keys.
+    static StringRef getModuleFlagKeyCGProfileName() {
+      return "CG Profile";
+    }
+
     /// Returns `true` if the given type is compatible with the LLVM dialect.
     static bool isCompatibleType(Type);
 

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

@@ -380,8 +380,20 @@ LogicalResult
 ModuleFlagAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                        LLVM::ModFlagBehavior flagBehavior, StringAttr key,
                        Attribute value) {
-  if (!isa<IntegerAttr, StringAttr>(value))
-    return emitError()
-           << "only integer and string values are currently supported";
-  return success();
+  if (key == LLVMDialect::getModuleFlagKeyCGProfileName()) {
+    auto arrayAttr = dyn_cast<ArrayAttr>(value);
+    if ((!arrayAttr) || (!llvm::all_of(arrayAttr, [](Attribute attr) {
+          return isa<ModuleFlagCGProfileEntryAttr>(attr);
+        })))
+      return emitError()
+             << "'CG Profile' key expects an array of '#llvm.cgprofile_entry'";
+    return success();
+  }
+
+  if (isa<IntegerAttr, StringAttr>(value))
+    return success();
+
+  return emitError() << "only integer and string values are currently "
+                        "supported for unknown key '"
+                     << key << "'";
 }

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

@@ -271,6 +271,31 @@ static void convertLinkerOptionsOp(ArrayAttr options,
   linkerMDNode->addOperand(listMDNode);
 }
 
+static llvm::Metadata *
+convertModuleFlagValue(StringRef key, ArrayAttr arrayAttr,
+                       llvm::IRBuilderBase &builder,
+                       LLVM::ModuleTranslation &moduleTranslation) {
+  llvm::LLVMContext &context = builder.getContext();
+  llvm::MDBuilder mdb(context);
+  SmallVector<llvm::Metadata *> nodes;
+
+  if (key == LLVMDialect::getModuleFlagKeyCGProfileName()) {
+    for (auto entry : arrayAttr.getAsRange<ModuleFlagCGProfileEntryAttr>()) {
+      llvm::Function *fromFn =
+          moduleTranslation.lookupFunction(entry.getFrom().getValue());
+      llvm::Function *toFn =
+          moduleTranslation.lookupFunction(entry.getTo().getValue());
+      llvm::Metadata *vals[] = {
+          llvm::ValueAsMetadata::get(fromFn), llvm::ValueAsMetadata::get(toFn),
+          mdb.createConstant(llvm::ConstantInt::get(
+              llvm::Type::getInt64Ty(context), entry.getCount()))};
+      nodes.push_back(llvm::MDNode::get(context, vals));
+    }
+    return llvm::MDTuple::getDistinct(context, nodes);
+  }
+  return nullptr;
+}
+
 static void convertModuleFlagsOp(ArrayAttr flags, llvm::IRBuilderBase &builder,
                                  LLVM::ModuleTranslation &moduleTranslation) {
   llvm::Module *llvmModule = moduleTranslation.getLLVMModule();
@@ -286,6 +311,11 @@ static void convertModuleFlagsOp(ArrayAttr flags, llvm::IRBuilderBase &builder,
                   llvm::Type::getInt32Ty(builder.getContext()),
                   intAttr.getInt()));
             })
+            .Case<ArrayAttr>([&](auto arrayAttr) {
+              return convertModuleFlagValue(flagAttr.getKey().getValue(),
+                                            arrayAttr, builder,
+                                            moduleTranslation);
+            })
             .Default([](auto) { return nullptr; });
 
     assert(valueMetadata && "expected valid metadata");

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -519,6 +519,39 @@ void ModuleImport::addDebugIntrinsic(llvm::CallInst *intrinsic) {
   debugIntrinsics.insert(intrinsic);
 }
 
+static Attribute convertCGProfileModuleFlagValue(ModuleOp mlirModule,
+                                                 llvm::MDTuple *mdTuple) {
+  auto getFunctionSymbol = [&](const llvm::MDOperand &funcMDO) {
+    auto *f = cast<llvm::ValueAsMetadata>(funcMDO);
+    auto *llvmFn = cast<llvm::Function>(f->getValue()->stripPointerCasts());
+    return FlatSymbolRefAttr::get(mlirModule->getContext(), llvmFn->getName());
+  };
+
+  // Each tuple element becomes one ModuleFlagCGProfileEntryAttr.
+  SmallVector<Attribute> cgProfile;
+  for (unsigned i = 0; i < mdTuple->getNumOperands(); i++) {
+    const llvm::MDOperand &mdo = mdTuple->getOperand(i);
+    auto *cgEntry = cast<llvm::MDNode>(mdo);
+    llvm::Constant *llvmConstant =
+        cast<llvm::ConstantAsMetadata>(cgEntry->getOperand(2))->getValue();
+    uint64_t count = cast<llvm::ConstantInt>(llvmConstant)->getZExtValue();
+    cgProfile.push_back(ModuleFlagCGProfileEntryAttr::get(
+        mlirModule->getContext(), getFunctionSymbol(cgEntry->getOperand(0)),
+        getFunctionSymbol(cgEntry->getOperand(1)), count));
+  }
+  return ArrayAttr::get(mlirModule->getContext(), cgProfile);
+}
+
+/// Invoke specific handlers for each known module flag value, returns nullptr
+/// if the key is unknown or unimplemented.
+static Attribute convertModuleFlagValueFromMDTuple(ModuleOp mlirModule,
+                                                   StringRef key,
+                                                   llvm::MDTuple *mdTuple) {
+  if (key == LLVMDialect::getModuleFlagKeyCGProfileName())
+    return convertCGProfileModuleFlagValue(mlirModule, mdTuple);
+  return nullptr;
+}
+
 LogicalResult ModuleImport::convertModuleFlagsMetadata() {
   SmallVector<llvm::Module::ModuleFlagEntry> llvmModuleFlags;
   llvmModule->getModuleFlagsMetadata(llvmModuleFlags);
@@ -530,7 +563,12 @@ LogicalResult ModuleImport::convertModuleFlagsMetadata() {
       valAttr = builder.getI32IntegerAttr(constInt->getZExtValue());
     } else if (auto *mdString = dyn_cast<llvm::MDString>(val)) {
       valAttr = builder.getStringAttr(mdString->getString());
-    } else {
+    } else if (auto *mdTuple = dyn_cast<llvm::MDTuple>(val)) {
+      valAttr = convertModuleFlagValueFromMDTuple(mlirModule, key->getString(),
+                                                  mdTuple);
+    }
+
+    if (!valAttr) {
       emitWarning(mlirModule.getLoc())
           << "unsupported module flag value: " << diagMD(val, llvmModule.get());
       continue;

mlir/test/Dialect/LLVMIR/invalid.mlir

@@ -1784,6 +1784,22 @@ module {
 
 // -----
 
+module {
+  // expected-error@below {{'CG Profile' key expects an array of '#llvm.cgprofile_entry'}}
+  llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", [
+    "yo"
+  ]>]
+}
+
+// -----
+
+module {
+  // expected-error@below {{'CG Profile' key expects an array of '#llvm.cgprofile_entry'}}
+  llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", 3 : i64>]
+}
+
+// -----
+
 llvm.func @t0() -> !llvm.ptr {
   %0 = llvm.blockaddress <function = @t0, tag = <id = 1>> : !llvm.ptr
   llvm.blocktag <id = 1>

mlir/test/Dialect/LLVMIR/module-roundtrip.mlir

@@ -6,7 +6,12 @@ module {
                      #llvm.mlir.module_flag<max, "PIE Level", 2 : i32>,
                      #llvm.mlir.module_flag<max, "uwtable", 2 : i32>,
                      #llvm.mlir.module_flag<max, "frame-pointer", 1 : i32>,
-                     #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">]
+                     #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">,
+                     #llvm.mlir.module_flag<append, "CG Profile", [
+                       #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+                       #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+                       #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+                    ]>]
 }
 
 // CHECK: llvm.module_flags [
@@ -15,4 +20,9 @@ module {
 // CHECK-SAME: #llvm.mlir.module_flag<max, "PIE Level", 2 : i32>,
 // CHECK-SAME: #llvm.mlir.module_flag<max, "uwtable", 2 : i32>,
 // CHECK-SAME: #llvm.mlir.module_flag<max, "frame-pointer", 1 : i32>,
-// CHECK-SAME: #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">]
+// CHECK-SAME: #llvm.mlir.module_flag<override, "probe-stack", "inline-asm">,
+// CHECK-SAME: #llvm.mlir.module_flag<append, "CG Profile", [
+// CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+// CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+// CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+// CHECK-SAME: ]>]

mlir/test/Target/LLVMIR/Import/module-flags.ll

@@ -25,3 +25,22 @@
 !12 = !{ i32 2, !"qux", i32 42 }
 !13 = !{ i32 3, !"qux", !{ !"foo", i32 1 }}
 !llvm.module.flags = !{ !10, !11, !12, !13 }
+
+; // -----
+
+declare void @from(i32)
+declare void @to()
+
+!llvm.module.flags = !{!20}
+
+!20 = !{i32 5, !"CG Profile", !21}
+!21 = distinct !{!22, !23, !24}
+!22 = !{ptr @from, ptr @to, i64 222}
+!23 = !{ptr @from, ptr @from, i64 222}
+!24 = !{ptr @to, ptr @from, i64 222}
+
+; CHECK: llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", [
+; CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+; CHECK-SAME: #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+; CHECK-SAME: #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+; CHECK-SAME: ]>]

mlir/test/Target/LLVMIR/llvmir.mlir

@@ -2838,6 +2838,25 @@ module {
 
 // -----
 
+llvm.module_flags [#llvm.mlir.module_flag<append, "CG Profile", [
+  #llvm.cgprofile_entry<from = @from, to = @to, count = 222>,
+  #llvm.cgprofile_entry<from = @from, to = @from, count = 222>,
+  #llvm.cgprofile_entry<from = @to, to = @from, count = 222>
+]>]
+llvm.func @from(i32)
+llvm.func @to()
+
+// CHECK: !llvm.module.flags = !{![[#CGPROF:]], ![[#DBG:]]}
+
+// CHECK: ![[#CGPROF]] = !{i32 5, !"CG Profile", ![[#LIST:]]}
+// CHECK: ![[#LIST]] = distinct !{![[#ENTRY_A:]], ![[#ENTRY_B:]], ![[#ENTRY_C:]]}
+// CHECK: ![[#ENTRY_A]] = !{ptr @from, ptr @to, i64 222}
+// CHECK: ![[#ENTRY_B]] = !{ptr @from, ptr @from, i64 222}
+// CHECK: ![[#ENTRY_C]] = !{ptr @to, ptr @from, i64 222}
+// CHECK: ![[#DBG]] = !{i32 2, !"Debug Info Version", i32 3}
+
+// -----
+
 module attributes {llvm.dependent_libraries = ["foo", "bar"]} {}
 
 // CHECK: !llvm.dependent-libraries =  !{![[#LIBFOO:]], ![[#LIBBAR:]]}