Optimize kernel builder JIT cache (#1206)

NVIDIA · Feb 13, 2024 · bcafaa5 · bcafaa5
1 parent 56c1b88
commit bcafaa5
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 20 deletions.
diff --git a/python/runtime/cudaq/algorithms/py_observe.cpp b/python/runtime/cudaq/algorithms/py_observe.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -39,7 +39,6 @@ observe_result pyObserve(kernel_builder<> &kernel, spin_op &spin_operator,
 
   // TODO: would like to handle errors in the case that
   // `kernel.num_qubits() >= spin_operator.num_qubits()`
-  kernel.jitCode();
   auto name = kernel.name();
 
   // Launch the observation task
@@ -68,9 +67,6 @@ async_observe_result pyObserveAsync(kernel_builder<> &kernel,
   // TODO: would like to handle errors in the case that
   // `kernel.num_qubits() >= spin_operator.num_qubits()`
 
-  // JIT the code
-  kernel.jitCode();
-
   // Get the kernel name
   auto name = kernel.name();
 
@@ -169,7 +165,6 @@ pyObserveN(kernel_builder<> &kernel, spin_op &op, py::args args = {},
   auto &platform = cudaq::get_platform();
   if (noise)
     platform.set_noise(&*noise);
-  kernel.jitCode();
   auto name = kernel.name();
   std::vector<observe_result> results;
   for (std::size_t currentIter = 0; auto &a : argSet) {

diff --git a/python/runtime/cudaq/algorithms/py_sample.cpp b/python/runtime/cudaq/algorithms/py_sample.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -31,7 +31,6 @@ sample_result pySample(kernel_builder<> &builder, py::args args = {},
   auto validatedArgs = validateInputArguments(builder, args);
 
   cudaq::info("Sampling the provided pythonic kernel.");
-  builder.jitCode();
   auto kernelName = builder.name();
 
   // Map py::args to OpaqueArguments handle
@@ -57,7 +56,6 @@ pySampleN(kernel_builder<> &kernel, py::args args = {},
           std::optional<noise_model> noise = std::nullopt) {
   auto argSet = createArgumentSet(args);
   auto N = argSet.size();
-  kernel.jitCode();
   auto name = kernel.name();
   auto &platform = cudaq::get_platform();
   if (noise)
@@ -90,8 +88,7 @@ async_sample_result pySampleAsync(kernel_builder<> &builder,
   auto &platform = cudaq::get_platform();
   cudaq::info("Asynchronously sampling the provided pythonic kernel.");
 
-  // JIT the code and get the kernel name
-  builder.jitCode();
+  // Get the kernel name
   auto kernelName = builder.name();
 
   // Create the argument holder and pack the runtime arguments

diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -35,7 +35,6 @@ void extractStateData(py::buffer_info &info, complex *data) {
 state pyGetState(kernel_builder<> &kernel, py::args args) {
   // Ensure the user input is correct.
   auto validatedArgs = validateInputArguments(kernel, args);
-  kernel.jitCode();
   OpaqueArguments argData;
   packArgs(argData, validatedArgs);
   return details::extractState(

diff --git a/runtime/cudaq/builder/kernel_builder.cpp b/runtime/cudaq/builder/kernel_builder.cpp
@@ -1,5 +1,5 @@
 /*******************************************************************************
- * Copyright (c) 2022 - 2023 NVIDIA Corporation & Affiliates.                  *
+ * Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -760,12 +760,11 @@ jitCode(ImplicitLocOpBuilder &builder, ExecutionEngine *jit,
   auto currentModule = function->getParentOfType<ModuleOp>();
 
   // Create a unique hash from that ModuleOp
-  std::string modulePrintOut;
-  {
-    llvm::raw_string_ostream os(modulePrintOut);
-    currentModule.print(os);
-  }
-  auto moduleHash = std::hash<std::string>{}(modulePrintOut);
+  auto hash = llvm::hash_code{0};
+  currentModule.walk([&hash](Operation *op) {
+    hash = llvm::hash_combine(hash, OperationEquivalence::computeHash(op));
+  });
+  auto moduleHash = static_cast<size_t>(hash);
 
   if (jit) {
     // Have we added more instructions