[Refactor] Remove redundant code in diagnostic driver

Author: Weiming Zhao

Diff for: armory/diagnostic/driver.cc

@@ -31,86 +31,17 @@
 #include "halo/lib/transforms/onnxextension_legalizer.h"
 #include "halo/lib/transforms/tfextension_legalizer.h"
 #include "halo/lib/transforms/type_legalizer.h"
+#include "halo/utils/cl_options.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 
 using namespace halo;
 
-static llvm::cl::list<std::string> ModelFiles(
-    llvm::cl::Positional, llvm::cl::desc("model file name."),
-    llvm::cl::OneOrMore);
-
-static llvm::cl::opt<Parser::Format> ModelFormat(
-    "x",
-    llvm::cl::desc(
-        "format of the following input model files. Permissible formats "
-        "include: TENSORFLOW CAFFE ONNX MXNET. If unspecified, the format is "
-        "guessed base on file's extension."),
-    llvm::cl::init(Parser::Format::INVALID));
-
-static llvm::cl::opt<std::string> EntryFunctionName(
-    "entry-func-name", llvm::cl::desc("name of entry function"),
-    llvm::cl::init(""));
-
 static llvm::cl::opt<bool> PrintDiagnosticReport(
     "print-diagnostic-report", llvm::cl::desc("Print diagnostic report"),
     llvm::cl::init(false));
 
-/// Guess the model format based on input file extension.gg
-static Parser::Format InferFormat(
-    const llvm::cl::list<std::string>& model_files, size_t file_idx) {
-  llvm::StringRef ext = llvm::sys::path::extension(model_files[file_idx]);
-  auto format = llvm::StringSwitch<Parser::Format>(ext)
-                    .Case(".pb", Parser::Format::TENSORFLOW)
-                    .Case(".pbtxt", Parser::Format::TENSORFLOW)
-                    .Case(".prototxt", Parser::Format::TENSORFLOW)
-                    .Case(".onnx", Parser::Format::ONNX)
-                    .Case(".json", Parser::Format::MXNET)
-                    .Default(Parser::Format::INVALID);
-  // Check the next input file to see if it is caffe.
-  if (format == Parser::Format::TENSORFLOW &&
-      (file_idx + 1 < model_files.size()) &&
-      llvm::sys::path::extension(model_files[file_idx + 1]) == ".caffemodel") {
-    format = Parser::Format::CAFFE;
-  }
-  return format;
-}
-
-static Status ParseModels(const llvm::cl::list<std::string>& model_files,
-                          const llvm::cl::opt<Parser::Format>& model_format,
-                          const llvm::cl::opt<std::string>& entry_func_name,
-                          const armory::Opts& opts, Module* module) {
-  std::set<std::string> func_names;
-  for (size_t i = 0, e = model_files.size(); i < e; ++i) {
-    Parser::Format format = model_format;
-    if (format == Parser::Format::INVALID) {
-      format = InferFormat(model_files, i);
-    }
-
-    FunctionBuilder func_builder(module);
-    // Use stem of the input model as function name.
-    std::string func_name = entry_func_name.empty()
-                                ? llvm::sys::path::stem(model_files[i]).str()
-                                : entry_func_name.getValue();
-    while (func_names.count(func_name) != 0) {
-      func_name.append("_").append(std::to_string(i));
-    }
-    func_names.insert(func_name);
-    Function* func = func_builder.CreateFunction(func_name);
-    std::vector<std::string> files{model_files[i]};
-    if (format == Parser::Format::CAFFE || format == Parser::Format::MXNET) {
-      HLCHECK(i + 1 < e);
-      files.push_back(model_files[++i]);
-    }
-    if (Status status = Parser::Parse(func, format, files, opts);
-        status != Status::SUCCESS) {
-      return status;
-    }
-  }
-  return Status::SUCCESS;
-}
-
 int main(int argc, char** argv) {
   llvm::cl::ParseCommandLineOptions(argc, argv);
   GlobalContext ctx;
@@ -119,8 +50,9 @@ int main(int argc, char** argv) {
   Module m(ctx, "diagnostic_module");
 
   armory::Opts opts(PrintDiagnosticReport);
-  if (ParseModels(ModelFiles, ModelFormat, EntryFunctionName, opts, &m) !=
-      Status::SUCCESS) {
+  Parser::Format format = Parser::Format::INVALID;
+  if (ParseModels(ModelFiles, ModelFormat, EntryFunctionName, opts, &m,
+                  &format) != Status::SUCCESS) {
     return 1;
   }