[nvq++][rt] Improve binary sizes by not globing MLIR dialects/passes.

By not globing all MLIR dialects/passes, many of which CUDAQ does not use, we can significantly improve binary sizes: ``` Filename Size (old) Size (new) Improvement -------- ----------- ----------- ----------- cudaq-opt 161.9 MB 39.6 MB -122.2 MB cudaq-quake 187.6 MB 108.6 MB -79.0 MB cudaq-translate 138.2 MB 76.1 MB -62.0 MB libcudaq-mlir-runtime.so 179.9 MB 111.8 MB -68.1 MB ``` This change also improves compilation time when building `llvm` along with CUDAQ, instead of using a pre-built version for llvm. It also improves the `--help` output for the `cudaq-*` tools, as they won't be "polluted" with many passes/options that are not relevant in the context of `nvq++`. Improving the python bindings situation, which can possibly profit from a similar change, will be left to a later commit. Signed-off-by: boschmitt <[email protected]>
NVIDIA · Nov 27, 2024 · b1db6fe · b1db6fe
1 parent a595985
commit b1db6fe
Show file tree

Hide file tree

Showing 15 changed files with 99 additions and 72 deletions.
diff --git a/include/cudaq/Optimizer/InitAllDialects.h b/include/cudaq/Optimizer/InitAllDialects.h
@@ -10,13 +10,34 @@
 
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 
 namespace cudaq {
 
-// Add all the MLIR dialects to the provided registry.
+// Add all required dialects to the provided registry.
 inline void registerAllDialects(mlir::DialectRegistry &registry) {
   // clang-format off
   registry.insert<
+    // MLIR dialects
+    mlir::AffineDialect,
+    mlir::arith::ArithDialect,
+    mlir::cf::ControlFlowDialect,
+    mlir::complex::ComplexDialect,
+    mlir::func::FuncDialect,
+    mlir::LLVM::LLVMDialect,
+    mlir::math::MathDialect,
+    mlir::memref::MemRefDialect,
+    mlir::scf::SCFDialect,
+
+    // NVQ++ dialects
     cudaq::cc::CCDialect,
     quake::QuakeDialect
   >();

diff --git a/include/cudaq/Optimizer/InitAllPasses.h b/include/cudaq/Optimizer/InitAllPasses.h
@@ -0,0 +1,34 @@
+/****************************************************************-*- C++ -*-****
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#pragma once
+
+#include "cudaq/Optimizer/CodeGen/Passes.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "mlir/Transforms/Passes.h"
+
+namespace cudaq {
+
+inline void registerAllPasses() {
+  // General MLIR passes
+  mlir::registerTransformsPasses();
+
+  // NVQPP passes
+  opt::registerOptCodeGenPasses();
+  opt::registerOptTransformsPasses();
+  opt::registerAggressiveEarlyInlining();
+
+  // Pipelines
+  opt::registerUnrollingPipeline();
+  opt::registerToExecutionManagerCCPipeline();
+  opt::registerTargetPipelines();
+  opt::registerWireSetToProfileQIRPipeline();
+  opt::registerMappingPipeline();
+}
+
+} // namespace cudaq
diff --git a/lib/Frontend/nvqpp/CMakeLists.txt b/lib/Frontend/nvqpp/CMakeLists.txt
@@ -8,9 +8,6 @@
 
 set(LIBRARY_NAME cudaq-mlirgen)
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-switch")
 add_mlir_library(${LIBRARY_NAME}
     # bridge from AST to Quake
@@ -25,8 +22,6 @@ add_mlir_library(${LIBRARY_NAME}
     QuakeDialect
     OptimBuilder
 
-    ${dialect_libs}
-    ${conversion_libs}
     MLIRIR
     MLIRParser
     MLIRPass
@@ -38,6 +33,7 @@ add_mlir_library(${LIBRARY_NAME}
     MLIRTargetLLVMIRExport
     MLIRLLVMCommonConversion
     MLIRLLVMToLLVMIRTranslation
+    MLIRMathDialect
     clangBasic
     clangAST
     clangFrontend

diff --git a/lib/Optimizer/Builder/CMakeLists.txt b/lib/Optimizer/Builder/CMakeLists.txt
@@ -16,6 +16,7 @@ add_mlir_dialect_library(OptimBuilder
 
   LINK_LIBS PUBLIC
     MLIRIR
+    MLIRParser
     CCDialect
     QuakeDialect
 )
diff --git a/lib/Optimizer/CodeGen/CMakeLists.txt b/lib/Optimizer/CodeGen/CMakeLists.txt
@@ -43,7 +43,27 @@ add_cudaq_library(OptCodeGen
 
   LINK_LIBS PUBLIC
     CCDialect
-    MLIRIR
     OptimBuilder
     QuakeDialect
+
+    MLIRArithTransforms
+    MLIRIR
+    MLIRLLVMDialect
+    MLIROpenACCDialect
+    MLIROpenMPDialect
+    MLIRTransforms
+
+    # Conversions
+    MLIRAffineToStandard
+    MLIRArithToLLVM
+    MLIRComplexToLibm
+    MLIRComplexToLLVM
+    MLIRControlFlowToLLVM
+    MLIRFuncToLLVM
+    MLIRMathToFuncs
+    MLIRMathToLLVM
+    MLIRSCFToControlFlow
+
+    # Translation
+    MLIRTargetLLVMIRExport
 )
diff --git a/runtime/common/CMakeLists.txt b/runtime/common/CMakeLists.txt
@@ -83,9 +83,6 @@ endif()
 ## here called cudaq-mlir-runtime that cudaq-builder and cudaq-rest-qpu 
 ## can link to and get MLIR/LLVM things uniformly
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-
 add_library(cudaq-mlir-runtime
   SHARED
     ArgumentConversion.cpp
@@ -113,8 +110,6 @@ target_include_directories(cudaq-mlir-runtime
 
 target_link_libraries(cudaq-mlir-runtime 
   PUBLIC
-    ${dialect_libs}
-    ${conversion_libs}
     nvqir
     CCDialect
     QuakeDialect

diff --git a/runtime/common/RuntimeMLIR.cpp b/runtime/common/RuntimeMLIR.cpp
@@ -14,6 +14,8 @@
 #include "cudaq/Optimizer/CodeGen/QIRFunctionNames.h"
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
+#include "cudaq/Optimizer/InitAllPasses.h"
 #include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/MC/SubtargetFeature.h"
@@ -25,8 +27,6 @@
 #include "llvm/Support/TargetSelect.h"
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
-#include "mlir/InitAllDialects.h"
-#include "mlir/InitAllPasses.h"
 #include "mlir/InitAllTranslations.h"
 #include "mlir/Parser/Parser.h"
 #include "mlir/Pass/PassManager.h"
@@ -75,24 +75,16 @@ std::unique_ptr<MLIRContext> initializeMLIR() {
   if (!mlirLLVMInitialized) {
     llvm::InitializeNativeTarget();
     llvm::InitializeNativeTargetAsmPrinter();
-    registerAllPasses();
-    cudaq::opt::registerOptCodeGenPasses();
-    cudaq::opt::registerOptTransformsPasses();
+    cudaq::registerAllPasses();
     registerToQIRTranslation();
     registerToOpenQASMTranslation();
     registerToIQMJsonTranslation();
-    cudaq::opt::registerAggressiveEarlyInlining();
-    cudaq::opt::registerUnrollingPipeline();
-    cudaq::opt::registerTargetPipelines();
-    cudaq::opt::registerWireSetToProfileQIRPipeline();
-    cudaq::opt::registerMappingPipeline();
     mlirLLVMInitialized = true;
   }
 
   DialectRegistry registry;
-  registry.insert<quake::QuakeDialect, cc::CCDialect>();
   cudaq::opt::registerCodeGenDialect(registry);
-  registerAllDialects(registry);
+  cudaq::registerAllDialects(registry);
   auto context = std::make_unique<MLIRContext>(registry);
   context->loadAllAvailableDialects();
   registerLLVMDialectTranslation(*context);

diff --git a/runtime/test/test_argument_conversion.cpp b/runtime/test/test_argument_conversion.cpp
@@ -15,8 +15,8 @@
 #include "common/ArgumentConversion.h"
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
 #include "cudaq/qis/pauli_word.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Parser/Parser.h"
 #include <numeric>
 
@@ -519,8 +519,7 @@ void test_combinations(mlir::MLIRContext *ctx) {
 
 int main() {
   mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  registry.insert<cudaq::cc::CCDialect, quake::QuakeDialect>();
+  cudaq::registerAllDialects(registry);
   mlir::MLIRContext context(registry);
   context.loadAllAvailableDialects();
   test_scalars(&context);

diff --git a/tools/cudaq-opt/CMakeLists.txt b/tools/cudaq-opt/CMakeLists.txt
@@ -6,18 +6,13 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-type-limits")
 
 add_llvm_executable(cudaq-opt cudaq-opt.cpp)
 
 llvm_update_compile_flags(cudaq-opt)
 target_link_libraries(cudaq-opt
   PUBLIC
-  ${dialect_libs}
-  ${conversion_libs}
-
   MLIRLLVMDialect
   MLIROptLib
   MLIRIR

diff --git a/tools/cudaq-opt/cudaq-opt.cpp b/tools/cudaq-opt/cudaq-opt.cpp
@@ -10,6 +10,8 @@
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Common/InlinerInterface.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
+#include "cudaq/Optimizer/InitAllPasses.h"
 #include "cudaq/Optimizer/Transforms/Passes.h"
 #include "cudaq/Support/Plugin.h"
 #include "cudaq/Support/Version.h"
@@ -20,8 +22,6 @@
 #include "llvm/Support/ToolOutputFile.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/InitAllDialects.h"
-#include "mlir/InitAllPasses.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
@@ -55,15 +55,7 @@ int main(int argc, char **argv) {
   // nvidia/cuda-quantum
   llvm::setBugReportMsg(cudaq::bugReportMsg);
 
-  mlir::registerAllPasses();
-  cudaq::opt::registerOptCodeGenPasses();
-  cudaq::opt::registerOptTransformsPasses();
-  cudaq::opt::registerAggressiveEarlyInlining();
-  cudaq::opt::registerUnrollingPipeline();
-  cudaq::opt::registerToExecutionManagerCCPipeline();
-  cudaq::opt::registerTargetPipelines();
-  cudaq::opt::registerWireSetToProfileQIRPipeline();
-  cudaq::opt::registerMappingPipeline();
+  cudaq::registerAllPasses();
 
   // See if we have been asked to load a pass plugin,
   // if so load it.
@@ -82,8 +74,7 @@ int main(int argc, char **argv) {
   }
 
   mlir::DialectRegistry registry;
-  registry.insert<quake::QuakeDialect, cudaq::cc::CCDialect>();
-  registerAllDialects(registry);
+  cudaq::registerAllDialects(registry);
   registerInlinerExtension(registry);
   return mlir::asMainReturnCode(
       mlir::MlirOptMain(argc, argv, "nvq++ optimizer\n", registry));

diff --git a/tools/cudaq-quake/CMakeLists.txt b/tools/cudaq-quake/CMakeLists.txt
@@ -6,18 +6,13 @@
 # the terms of the Apache License 2.0 which accompanies this distribution.     #
 # ============================================================================ #
 
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-
 set(LLVM_LINK_COMPONENTS Support ${LLVM_TARGETS_TO_BUILD})
 
 add_llvm_executable(cudaq-quake cudaq-quake.cpp)
 
 llvm_update_compile_flags(cudaq-quake)
 target_link_libraries(cudaq-quake
   PRIVATE
-  ${dialect_libs}
-  ${conversion_libs}
   MLIRIR
   MLIRParser
   MLIRPass
@@ -29,6 +24,11 @@ target_link_libraries(cudaq-quake
   MLIRTargetLLVMIRExport
   MLIRLLVMCommonConversion
   MLIRLLVMToLLVMIRTranslation
+
+  MLIRAffineDialect
+  MLIRMemRefDialect
+  MLIRSCFDialect
+
   clangCodeGen
   clangFrontendTool
   clangFrontend

diff --git a/tools/cudaq-quake/cudaq-quake.cpp b/tools/cudaq-quake/cudaq-quake.cpp
@@ -16,6 +16,7 @@
 #include "cudaq/Optimizer/Builder/Runtime.h"
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeDialect.h"
+#include "cudaq/Optimizer/InitAllDialects.h"
 #include "cudaq/Optimizer/Support/Verifier.h"
 #include "cudaq/Support/Version.h"
 #include "nvqpp_config.h"
@@ -32,9 +33,9 @@
 #include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/InitAllPasses.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
 #include <filesystem>
 
 using namespace llvm;
@@ -313,13 +314,10 @@ int main(int argc, char **argv) {
     return ec.value();
   }
 
-  mlir::registerAllPasses();
-
   // Read the code into a memory buffer and setup MLIR.
   auto cplusplusCode = fileOrError.get()->getBuffer();
   mlir::DialectRegistry registry;
-  mlir::registerAllDialects(registry);
-  registry.insert<cudaq::cc::CCDialect, quake::QuakeDialect>();
+  cudaq::registerAllDialects(registry);
   mlir::MLIRContext context(registry);
   // TODO: Consider only loading the dialects we know we'll use.
   context.loadAllAvailableDialects();

diff --git a/tools/cudaq-translate/CMakeLists.txt b/tools/cudaq-translate/CMakeLists.txt
@@ -7,20 +7,13 @@
 # ============================================================================ #
 
 set(LLVM_LINK_COMPONENTS Support ${LLVM_TARGETS_TO_BUILD})
-
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
-get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
-get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-type-limits")
 
 add_llvm_executable(cudaq-translate cudaq-translate.cpp)
 
 llvm_update_compile_flags(cudaq-translate)
 target_link_libraries(cudaq-translate
   PRIVATE
-  ${dialect_libs}
-  ${translation_libs}
-  ${conversion_libs}
   MLIRIR
   MLIRParser
   MLIRPass