Merge pull request #521 from lamhoangtung/raise_error_with_better_pytorch_traceback

feat: Show PyTorch code of unsupported operators

Author: narendasan

core/conversion/conversion.cpp

@@ -436,8 +436,8 @@ std::string ConvertBlockToEngine(const torch::jit::Block* b, ConversionInfo buil
   return engine;
 }
 
-std::set<std::string> GetUnsupportedOpsInBlock(const torch::jit::Block* b) {
-  std::set<std::string> unsupported_ops;
+std::unordered_map<c10::OperatorName, std::string> GetUnsupportedOpsInBlock(const torch::jit::Block* b) {
+  std::unordered_map<c10::OperatorName, std::string> unsupported_ops;
   for (const auto n : b->nodes()) {
     if (n->kind() != torch::jit::prim::Loop && n->kind() != torch::jit::prim::If && !OpSupported(n)) {
       auto schema = n->maybeSchema();
@@ -446,7 +446,7 @@ std::set<std::string> GetUnsupportedOpsInBlock(const torch::jit::Block* b) {
           "Unable to get schema for Node " << util::node_info(n) << " (conversion.VerifyCoverterSupportForBlock)");
       std::stringstream ss;
       ss << *schema;
-      unsupported_ops.insert(ss.str());
+      unsupported_ops[schema->operator_name()] = ss.str();
     }
     for (const auto sub_b : n->blocks()) {
       auto sub_b_unsupported_ops = GetUnsupportedOpsInBlock(sub_b);
@@ -488,12 +488,27 @@ bool VerifyConverterSupportForBlock(const torch::jit::Block* b) {
     unsupported_msg << "Method requested cannot be compiled by TRTorch.\nUnsupported operators listed below:"
                     << std::endl;
     for (auto s : unsupported_ops) {
-      unsupported_msg << "  -  " << s << std::endl;
+      unsupported_msg << "  - " << s.second << std::endl;
     }
     unsupported_msg << "You can either implement converters for these ops in your application or request implementation"
                     << std::endl;
     unsupported_msg << "https://www.github.com/nvidia/TRTorch/issues" << std::endl;
+    unsupported_msg << std::endl << "In Module:" << std::endl;
+
     LOG_ERROR(unsupported_msg.str());
+
+    for (const auto n : b->nodes()) {
+      auto schema = n->maybeSchema();
+      if (schema) {
+        for (const auto& x : unsupported_ops) {
+          if (x.first == schema->operator_name()) {
+            LOG_ERROR(
+                "Unsupported operator: " << *schema << std::endl
+                                         << trtorch::core::util::GetPyTorchSourceCode(n) << std::endl);
+          }
+        }
+      }
+    }
     return false;
   }
 

core/util/jit_util.h

@@ -47,6 +47,11 @@ inline c10::FunctionSchema GenerateGraphSchema(std::string method_name, std::sha
   return c10::FunctionSchema(method_name, method_name, args, returns);
 }
 
+inline std::string GetPyTorchSourceCode(const torch::jit::Node* n) {
+  std::string source_code = n->sourceRange().str();
+  return source_code;
+}
+
 } // namespace util
 } // namespace core
 } // namespace trtorch