message: Let tf-mlir-translate support -split-input-file
  Having a -split-input-file mode is useful in tf-mlir-translate. It would allow us to put logically related tests in the same test file for better organization.
diff: @@ -24,6 +24,7 @@ limitations under the License.
  #include "mlir/IR/MLIRContext.h" // TF:local_config_mlir
  #include "mlir/Support/FileUtilities.h" // TF:local_config_mlir
  #include "mlir/Support/LogicalResult.h" // TF:local_config_mlir
  +#include "mlir/Support/ToolUtilities.h" // TF:local_config_mlir
  #include "mlir/Support/TranslateClParser.h" // TF:local_config_mlir
  #include "tensorflow/compiler/mlir/init_mlir.h"
  #include "tensorflow/compiler/mlir/tensorflow/translate/tf_mlir_translate.h"
  @@ -40,6 +41,13 @@ static llvm::cl::opt<std::string> output_filename(
  "o", llvm::cl::desc("Output filename"), llvm::cl::value_desc("filename"),
  llvm::cl::init("-"));
  +// NOLINTNEXTLINE
  +static llvm::cl::opt<bool> splitInputFile(
  + "split-input-file",
  + llvm::cl::desc("Split the input file into pieces and process each chunk "
  + "independently"),
  + llvm::cl::init(false));
  +
  // NOLINTNEXTLINE
  static llvm::cl::opt<bool> import_saved_model(
  "savedmodel-to-mlir",
  @@ -85,13 +93,12 @@ int main(int argc, char** argv) {
  return 1;
  }
  - mlir::MLIRContext context;
  -
  if (import_saved_model) {
  std::unordered_set<std::string> tags =
  absl::StrSplit(saved_model_tags, ',');
  std::vector<std::string> exported_names =
  absl::StrSplit(saved_model_exported_names, ',', absl::SkipEmpty());
  + mlir::MLIRContext context;
  auto module = tensorflow::SavedModelToMlirImport(
  input_filename, tags, absl::Span<std::string>(exported_names),
  @@ -107,12 +114,23 @@ int main(int argc, char** argv) {
  return 1;
  }
  - llvm::SourceMgr source_mgr;
  - source_mgr.AddNewSourceBuffer(std::move(input), llvm::SMLoc());
  - mlir::SourceMgrDiagnosticHandler diagnostic_handler(source_mgr, &context);
  + // Processes the memory buffer with a new MLIRContext.
  + auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
  + llvm::raw_ostream& os) {
  + llvm::SourceMgr sourceMgr;
  + sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), llvm::SMLoc());
  + mlir::MLIRContext context;
  + mlir::SourceMgrDiagnosticHandler diagnostic_handler(sourceMgr, &context);
  + return (*requested_translation)(sourceMgr, os, &context);
  + };
  - if (failed((*requested_translation)(source_mgr, output->os(), &context)))
  + if (splitInputFile) {
  + if (failed(mlir::splitAndProcessBuffer(std::move(input), processBuffer,
  + output->os())))
  return 1;
  + } else {
  + if (failed(processBuffer(std::move(input), output->os()))) return 1;
  + }
  }
  output->keep();
  