From b856fc57872a358edabb1a130e8cc7a0252b5f01 Mon Sep 17 00:00:00 2001
From: Haruki Imai <imaihal@jp.ibm.com>
Date: Thu, 7 Nov 2024 15:12:56 +0900
Subject: [PATCH] Option to not emit the full MLIR (only emit .tmp file)
 (#2997)

When emitting MLIR, there are two versions of IR, <name>.onnx.mlir and <name>.tmp . Since the constant values
are embedded in <name>.onnx.mlir, we got memory and disk pressure especially in large models. This option specify
not emit full MLIR(<name>.onnx.mlir). This option works with emitting MLIR options such as --EmitONNXIR and
 --EmitMLIR.

Signed-off-by: Haruki Imai <imaihal@jp.ibm.com>
---
 src/Compiler/CompilerOptions.cpp            | 13 ++++++++++++-
 src/Compiler/CompilerOptions.hpp            |  1 +
 src/Compiler/CompilerUtils.cpp              | 15 ++++++++-------
 test/mlir/driver/do_not_emit_full_mlir.mlir |  9 +++++++++
 4 files changed, 30 insertions(+), 8 deletions(-)
 create mode 100644 test/mlir/driver/do_not_emit_full_mlir.mlir

diff --git a/src/Compiler/CompilerOptions.cpp b/src/Compiler/CompilerOptions.cpp
index d71d5c333a..730f2ac50f 100644
--- a/src/Compiler/CompilerOptions.cpp
+++ b/src/Compiler/CompilerOptions.cpp
@@ -4,7 +4,7 @@
 
 //===------------------------ CompilerOptions.cpp -------------------------===//
 //
-// Copyright 2022, 2023 The IBM Research Authors.
+// Copyright 2022, 2024 The IBM Research Authors.
 //
 // =============================================================================
 //
@@ -49,6 +49,7 @@ EmissionTargetType emissionTarget;                     // onnx-mlir only
 bool invokeOnnxVersionConverter;                       // onnx-mlir only
 bool preserveLocations;                                // onnx-mlir only
 bool printIR;                                          // onnx-mlir only
+bool doNotEmitFullMLIRCode;                            // onnx-mlir only
 bool preserveBitcode;                                  // onnx-mlir only
 bool preserveLLVMIR;                                   // onnx-mlir only
 bool preserveMLIR;                                     // onnx-mlir only
@@ -281,6 +282,16 @@ static llvm::cl::opt<bool, true> printIROpt("printIR",
     llvm::cl::desc("Print the IR to stdout:."), llvm::cl::location(printIR),
     llvm::cl::init(false), llvm::cl::cat(OnnxMlirOptions));
 
+static llvm::cl::opt<bool, true> doNotEmitFullMLIRCodeOpt(
+    "do-not-emit-full-mlir-code",
+    llvm::cl::desc(
+        "Do not emit the MLIR the constant values are embeded "
+        "(<name>onnx.mlir). Emit only the MLIR without the constants "
+        "(<name>.tmp). Need to be used with emitting MLIR options such as "
+        "--EmitONNXIR and --EmitMLIR."),
+    llvm::cl::location(doNotEmitFullMLIRCode), llvm::cl::init(false),
+    llvm::cl::cat(OnnxMlirOptions));
+
 static llvm::cl::opt<bool, true> preserveBitcodeOpt("preserveBitcode",
     llvm::cl::desc("Preserve the bitcode files (optimized and unoptimized)."),
     llvm::cl::location(preserveBitcode), llvm::cl::init(false),
diff --git a/src/Compiler/CompilerOptions.hpp b/src/Compiler/CompilerOptions.hpp
index f216943e57..f7ebc2c6e5 100644
--- a/src/Compiler/CompilerOptions.hpp
+++ b/src/Compiler/CompilerOptions.hpp
@@ -97,6 +97,7 @@ extern bool printIR;                                          // onnx-mlir only
 extern bool preserveBitcode;                                  // onnx-mlir only
 extern bool preserveLLVMIR;                                   // onnx-mlir only
 extern bool preserveMLIR;                                     // onnx-mlir only
+extern bool doNotEmitFullMLIRCode;                            // onnx-mlir only
 extern bool useOnnxModelTypes;                                // onnx-mlir only
 extern int repeatOnnxTransform;                               // onnx-mlir only
 extern std::string shapeInformation;                          // onnx-mlir only
diff --git a/src/Compiler/CompilerUtils.cpp b/src/Compiler/CompilerUtils.cpp
index 105763107e..ac88eb213e 100644
--- a/src/Compiler/CompilerUtils.cpp
+++ b/src/Compiler/CompilerUtils.cpp
@@ -835,13 +835,14 @@ static int emitOutputFiles(std::string outputNameNoExt,
     // Emit the version with all constants included.
     std::string outputNameWithExt =
         getTargetFilename(outputNameNoExt, emissionTarget);
-    int rc = outputCode(module, outputNameWithExt);
-    if (VerboseOutput)
-      llvm::outs() << "Full MLIR code written to:\n"
-                   << "\t" << outputNameWithExt << "\n\n";
-    if (rc != CompilerSuccess)
-      return rc;
-
+    if (!doNotEmitFullMLIRCode) {
+      int rc = outputCode(module, outputNameWithExt);
+      if (VerboseOutput)
+        llvm::outs() << "Full MLIR code written to:\n"
+                     << "\t" << outputNameWithExt << "\n\n";
+      if (rc != CompilerSuccess)
+        return rc;
+    }
     // Elide element attributes if larger than 100.
     if (emissionTarget == EmitONNXBasic || emissionTarget == EmitONNXIR ||
         emissionTarget == EmitMLIR) {
diff --git a/test/mlir/driver/do_not_emit_full_mlir.mlir b/test/mlir/driver/do_not_emit_full_mlir.mlir
new file mode 100644
index 0000000000..794a7a0934
--- /dev/null
+++ b/test/mlir/driver/do_not_emit_full_mlir.mlir
@@ -0,0 +1,9 @@
+// RUN: onnx-mlir --EmitMLIR --do-not-emit-full-mlir-code %s -o %t && test ! -f %t.onnx.mlir && rm %t.tmp
+
+module {
+  func.func @main_graph(%arg0: tensor<?xf32>) -> tensor<?xf32> {
+    onnx.Return %arg0 : tensor<?xf32>
+  }
+  "onnx.EntryPoint"() {func = @main_graph} : () -> ()
+}
+