TensorFlow has segmentation fault in tfg-translate

@@ -219,6 +219,7 @@ tf_cc_binary(
         ":graphdef_export",
         ":graphdef_import",
         ":load_proto",
+        "//tensorflow/compiler/mlir:init_mlir",
         "//tensorflow/core:ops",  # Ops need to be registered for import.
         "//tensorflow/core/ir:Dialect",
         "@llvm-project//mlir:IR",

@@ -19,6 +19,7 @@ limitations under the License.
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Tools/mlir-translate/MlirTranslateMain.h"  // from @llvm-project
 #include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
 #include "tensorflow/core/ir/dialect.h"
 #include "tensorflow/core/ir/importexport/graphdef_export.h"
 #include "tensorflow/core/ir/importexport/graphdef_import.h"
@@ -63,6 +64,7 @@ TranslateFromMLIRRegistration mlir_to_graphdef(
 
 int main(int argc, char **argv) {
   mlir::registerAsmPrinterCLOptions();
+  tensorflow::InitMlir y(&argc, &argv);
   return failed(
       mlir::mlirTranslateMain(argc, argv, "Graph(Def)<->TFG Translation Tool"));
 }

@@ -86,11 +86,11 @@ class ValueMapManager {
     return ::tensorflow::OkStatus();
   }
 
-  Value GetValueOrCreatePlaceholder(StringRef full_name) {
+  tensorflow::StatusOr<Value> GetValueOrCreatePlaceholder(StringRef full_name) {
     StringRef node_name;
     StringRef output_name = "";
     bool is_control_dep = full_name[0] == '^';
-    int output_num = 0;
+    size_t output_num = 0;
     if (is_control_dep) full_name = full_name.drop_front();
     {
       size_t colon_sep = full_name.find_first_of(':');
@@ -105,8 +105,16 @@ class ValueMapManager {
         // NOLINTNEXTLINE: type matching the API taking a reference.
         unsigned long long value;
         if (!llvm::getAsUnsignedInteger(output_name.drop_front(colon_sep + 1),
-                                        10, value))
-          output_num = value;
+                                        10, value)) {
+          if (LLVM_LIKELY(
+                  value <=
+                  std::numeric_limits<llvm::SmallVectorSizeType<Value>>::max() -
+                      1))
+            output_num = value;
+          else
+            return InvalidArgument("Output index ", value,
+                                   " is invalid (too large)");
+        }
         output_name = output_name.take_front(colon_sep);
       }
     }
@@ -171,8 +179,9 @@ Status ImportNodes(ValueMapManager value_manager,
     for (const std::string& input : node.input()) {
       if (input.empty())
         return InvalidArgument("Node '", node.name(), "' has an empty input");
-      state.operands.push_back(
-          value_manager.GetValueOrCreatePlaceholder(input));
+      TF_ASSIGN_OR_RETURN(Value value,
+                          value_manager.GetValueOrCreatePlaceholder(input));
+      state.operands.push_back(value);
     }
     // Retrieve the entry in the nodes_map for this node and infer the result
     // count from what was inferred during the first traversal above.
@@ -470,8 +479,9 @@ Status ImportGenericFunction(
       return InvalidArgument("Function '", func.signature().name(),
                              "' has empty result name");
     }
-    ret_vals[position->second] =
-        value_manager.GetValueOrCreatePlaceholder(ret_val.second);
+    TF_ASSIGN_OR_RETURN(
+        ret_vals[position->second],
+        value_manager.GetValueOrCreatePlaceholder(ret_val.second));
   }
   for (const auto& ret_val : func.control_ret()) {
     auto position = control_output_to_position.find(ret_val.first);
@@ -485,8 +495,8 @@ Status ImportGenericFunction(
       return InvalidArgument("Function '", func.signature().name(),
                              "' has empty control result name");
     }
-    Value result = value_manager.GetValueOrCreatePlaceholder(
-        (Twine("^") + ret_val.second).str());
+    TF_ASSIGN_OR_RETURN(Value result, value_manager.GetValueOrCreatePlaceholder(
+                                          (Twine("^") + ret_val.second).str()));
     if (!result.getType().isa<ControlType>())
       return InvalidArgument("failed to map returned value ", ret_val.second,
                              ", isn't a control output");

@@ -0,0 +1,52 @@
+# RUN: not tfg-translate -graphdef-to-mlir %s 2>&1 | FileCheck %s
+
+# CHECK: Output index 18446744073709551615 is invalid (too large)
+
+library {
+  function {
+    signature {
+      name: "foo"
+      attr {
+        name: "T"
+        type: "type"
+      }
+    }
+    node_def {
+      name: "two"
+      op: "Const"
+      attr {
+        key: "dtype"
+        value {
+          type: DT_INT64
+        }
+      }
+      attr {
+        key: "value"
+        value {
+          tensor {
+            dtype: DT_INT64
+            tensor_shape {}
+            int64_val: 2
+          }
+        }
+      }
+    }
+    node_def {
+      name: "a"
+      op: "Cast"
+      input: "two:output:18446744073709551615"
+      attr {
+        key: "DstT"
+        value {
+          placeholder: "T"
+        }
+      }
+      attr {
+        key: "SrcT"
+        value {
+          type: DT_INT64
+        }
+      }
+    }
+  }
+}