[mlir][llvm] Add MD_prof import error handling

This commit adds additional checks and warning messages to the MD_prof
import. As LLVM does not verify most metadata, the import has the be
resilient towards ill-formatted inputs.

Reviewed By: gysit

Differential Revision: https://reviews.llvm.org/D143492

Author: Dinistro

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

@@ -84,36 +84,48 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
                                       LLVM::ModuleImport &moduleImport) {
   // Return success for empty metadata nodes since there is nothing to import.
   if (!node->getNumOperands())
-    return success();
+    return op->emitWarning() << "expected non-empty profiling metadata node";
 
   auto *name = dyn_cast<llvm::MDString>(node->getOperand(0));
   if (!name)
-    return failure();
+    return op->emitWarning()
+           << "expected profiling metadata node to have a string identifier";
 
   // Handle function entry count metadata.
   if (name->getString().equals("function_entry_count")) {
+    auto emitNodeWarning = [&]() {
+      return op->emitWarning()
+             << "expected function_entry_count to hold a single i64 value";
+    };
+
     // TODO support function entry count metadata with GUID fields.
     if (node->getNumOperands() != 2)
-      return failure();
+      return emitNodeWarning();
 
     llvm::ConstantInt *entryCount =
-        llvm::mdconst::extract<llvm::ConstantInt>(node->getOperand(1));
+        llvm::mdconst::dyn_extract<llvm::ConstantInt>(node->getOperand(1));
+    if (!entryCount)
+      return emitNodeWarning();
     if (auto funcOp = dyn_cast<LLVMFuncOp>(op)) {
       funcOp.setFunctionEntryCount(entryCount->getZExtValue());
       return success();
     }
-    return failure();
+    return op->emitWarning()
+           << "expected function_entry_count to be attached to a function";
   }
 
   if (!name->getString().equals("branch_weights"))
-    return failure();
+    return op->emitWarning()
+           << "unknown profiling metadata node " << name->getString();
 
   // 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));
+        llvm::mdconst::dyn_extract<llvm::ConstantInt>(node->getOperand(i));
+    if (!branchWeight)
+      return op->emitWarning() << "expected branch weights to be integers";
     branchWeights.push_back(branchWeight->getZExtValue());
   }
 
@@ -124,7 +136,10 @@ static LogicalResult setProfilingAttr(OpBuilder &builder, llvm::MDNode *node,
             builder.getI32VectorAttr(branchWeights));
         return success();
       })
-      .Default([](auto) { return failure(); });
+      .Default([op](auto) {
+        return op->emitWarning()
+               << op->getName() << " does not support branch weights";
+      });
 }
 
 /// Searches the symbol reference pointing to the metadata operation that

mlir/test/Target/LLVMIR/Import/import-failure.ll

@@ -56,8 +56,7 @@ define void @unhandled_intrinsic(ptr %arg1, ptr %arg2) {
 
 ; // -----
 
-; CHECK:      import-failure.ll
-; CHECK-SAME: warning: unhandled metadata: !0 = !{!"unknown metadata"} on br i1 %arg1, label %bb1, label %bb2, !prof !0
+; CHECK: warning: unhandled metadata: !0 = !{!"unknown metadata"} on br i1 %arg1, label %bb1, label %bb2, !prof !0
 define i64 @unhandled_metadata(i1 %arg1, i64 %arg2) {
 entry:
   br i1 %arg1, label %bb1, label %bb2, !prof !0
@@ -71,8 +70,7 @@ bb2:
 
 ; // -----
 
-; CHECK:      import-failure.ll
-; CHECK-SAME: warning: unhandled function metadata: !0 = !{!"unknown metadata"} on define void @unhandled_func_metadata(i1 %arg1, i64 %arg2) !prof !0
+; CHECK: warning: unhandled function metadata: !0 = !{!"unknown metadata"} on define void @unhandled_func_metadata(i1 %arg1, i64 %arg2) !prof !0
 define void @unhandled_func_metadata(i1 %arg1, i64 %arg2) !prof !0 {
   ret void
 }
@@ -439,3 +437,126 @@ end:
 !0 = distinct !{!0, !1, !2}
 !1 = !{!"llvm.loop.disable_nonforced"}
 !2 = !{!"llvm.loop.typo"}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected non-empty profiling metadata node
+; CHECK:      warning: unhandled metadata: !0 = !{}
+define void @cond_br(i1 %arg) {
+entry:
+  br i1 %arg, label %bb1, label %bb2, !prof !0
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected profiling metadata node to have a string identifier
+; CHECK:      import-failure.ll:{{.*}} warning: unhandled metadata: !0 = !{i32 64}
+define void @cond_br(i1 %arg) {
+entry:
+  br i1 %arg, label %bb1, label %bb2, !prof !0
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{i32 64}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected function_entry_count to hold a single i64 value
+; CHECK:      warning: unhandled function metadata: !0 = !{!"function_entry_count"}
+define void @cond_br(i1 %arg) !prof !0 {
+entry:
+  br i1 %arg, label %bb1, label %bb2
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{!"function_entry_count"}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected function_entry_count to hold a single i64 value
+; CHECK:      warning: unhandled function metadata: !0 = !{!"function_entry_count", !"string"}
+define void @cond_br(i1 %arg) !prof !0 {
+entry:
+  br i1 %arg, label %bb1, label %bb2
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{!"function_entry_count", !"string"}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected function_entry_count to be attached to a function
+; CHECK:      warning: unhandled metadata: !0 = !{!"function_entry_count", i64 42}
+define void @cond_br(i1 %arg) {
+entry:
+  br i1 %arg, label %bb1, label %bb2, !prof !0
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{!"function_entry_count", i64 42}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: unknown profiling metadata node unknown_prof_type
+; CHECK:      warning: unhandled metadata: !0 = !{!"unknown_prof_type"}
+define void @cond_br(i1 %arg) {
+entry:
+  br i1 %arg, label %bb1, label %bb2, !prof !0
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{!"unknown_prof_type"}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: expected branch weights to be integers
+; CHECK:      warning: unhandled metadata: !0 = !{!"branch_weights", !"foo"}
+define void @cond_br(i1 %arg) {
+entry:
+  br i1 %arg, label %bb1, label %bb2, !prof !0
+bb1:
+  ret void
+bb2:
+  ret void
+}
+
+!0 = !{!"branch_weights", !"foo"}
+
+; // -----
+
+; CHECK:      import-failure.ll
+; CHECK-SAME: warning: llvm.func does not support branch weights
+; CHECK:      import-failure.ll:{{.*}} warning: unhandled function metadata: !0 = !{!"branch_weights", i32 64}
+define void @cond_br(i1 %arg) !prof !0 {
+  ret void
+}
+
+!0 = !{!"branch_weights", i32 64}