PennyLaneAI / catalyst

A JIT compiler for hybrid quantum programs in PennyLane
https://docs.pennylane.ai/projects/catalyst
Apache License 2.0
142 stars 36 forks source link

Avoid printing the IR unless the user specifically requests keep_intermediate #1327

Open erick-xanadu opened 4 days ago

erick-xanadu commented 4 days ago
diff --git a/mlir/lib/Driver/CompilerDriver.cpp b/mlir/lib/Driver/CompilerDriver.cpp
index d75406bf6..2114a7a78 100644
--- a/mlir/lib/Driver/CompilerDriver.cpp
+++ b/mlir/lib/Driver/CompilerDriver.cpp
@@ -691,7 +691,9 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &
             return failure();
         }
         output.outIR.clear();
-        outIRStream << *mlirModule;
+        if (options.keepIntermediate) {
+            outIRStream << *mlirModule;
+        }
         optTiming.stop();
     }

@@ -716,7 +718,9 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &
             dumpToFile(options, outFile, tmp);
         }
         output.outIR.clear();
-        outIRStream << *llvmModule;
+        if (options.keepIntermediate) {
+            outIRStream << *llvmModule;
+        }
         translateTiming.stop();
     }

@@ -772,7 +776,9 @@ LogicalResult QuantumDriverMain(const CompilerOptions &options, CompilerOutput &

         TimingScope outputTiming = llcTiming.nest("compileObject");
         output.outIR.clear();
-        outIRStream << *llvmModule;
+        if (options.keepIntermediate) {
+            outIRStream << *llvmModule;
+        }

         if (failed(timer::timer(compileObjectFile, "compileObjFile", /* add_endl */ true, options,
                                 std::move(llvmModule), targetMachine, options.getObjectFile()))) {

Essentially this, but we need to make sure that all tests pass. A lot of tests rely on the attribute .mlir and .qir to be available even when not using keep_intermediate=True.

Why should we make this change? The textual representation may be too big. In some cases up to GB of disk are used just to store the LLVM-IR. Dumping all of this also takes time.

dime10 commented 4 days ago

I think originally this came for free because we were passing the textual IR to llc as a file. Are we now invoking it directly on the IR in memory, or do we need to go via a bitcode file?

erick-xanadu commented 4 days ago

Are we now invoking it directly on the IR in memory, or do we need to go via a bitcode file?

Most tests pass with this change. I believe we don't really need to go to disk.