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 <[email protected]>
onnx · Nov 7, 2024 · b856fc5 · b856fc5
1 parent 33b466e
commit b856fc5
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 8 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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} : () -> ()
+}
+