[mlir][llvm] Add MD_prof import error handling

[Dinistro]

Summary

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.

[gysit]

Thanks LGTM!

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

96

nit: emitValueWarning

[Dinistro]

address NIT comment

[Dinistro]

rebased

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

 /// Converts the given profiling metadata `node` to an MLIR profiling attribute
 /// and attaches it to the imported operation if the translation succeeds.
 /// Returns failure otherwise.
 static LogicalResult setProfilingAttr(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 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());
   }
 
   // Attach the branch weights to the operations that support it.
   return llvm::TypeSwitch<Operation *, LogicalResult>(op)
       .Case<CondBrOp, SwitchOp, CallOp, InvokeOp>([&](auto branchWeightOp) {
         branchWeightOp.setBranchWeightsAttr(
             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
 /// maps to the given TBAA metadata `node` and attaches it to the imported
 /// operation if the lookup succeeds. Returns failure otherwise.
 static LogicalResult setTBAAAttr(const llvm::MDNode *node, Operation *op,
                                  LLVM::ModuleImport &moduleImport) {
   SymbolRefAttr tbaaTagSym = moduleImport.lookupTBAAAttr(node);

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

 ; CHECK-SAME: error: unhandled intrinsic: call void @llvm.gcroot(ptr %arg1, ptr %arg2)
 define void @unhandled_intrinsic(ptr %arg1, ptr %arg2) {
   call void @llvm.gcroot(ptr %arg1, ptr %arg2)
   ret void
 }
 
 ; // -----
 
-; 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
 bb1:
   ret i64 %arg2
 bb2:
   ret i64 %arg2
 }
 
 !0 = !{!"unknown metadata"}
 
 ; // -----
 
-; 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
 }
 
 !0 = !{!"unknown metadata"}
 
 ; // -----
 
entry:
   br label %end, !llvm.loop !0
 end:
   ret void
 }
 
 !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}