Add hybridLaunchKernel dispatch (NVIDIA#2111)

* Add hybridLaunchKernel Dispatch kernel launch based on whether JIT+synthesis would be required. * Fix tests * Modify error message
oqc-tech · Aug 21, 2024 · 18c305b · 18c305b
1 parent e7afc56
commit 18c305b
Show file tree

Hide file tree

Showing 10 changed files with 30 additions and 18 deletions.
diff --git a/include/cudaq/Optimizer/Transforms/Passes.td b/include/cudaq/Optimizer/Transforms/Passes.td
@@ -332,7 +332,7 @@ def GenerateKernelExecution : Pass<"kernel-execution", "mlir::ModuleOp"> {
       /*default=*/"\"-\"", "Name of output file.">,
     Option<"startingArgIdx", "starting-arg-idx", "std::size_t", /*default=*/"0",
       "The starting argument index for the argsCreator.">,
-    Option<"codegenKind", "codegen", "std::size_t", /*default=*/"1",
+    Option<"codegenKind", "codegen", "std::size_t", /*default=*/"0",
       "Set the kind of code to generate for the launches.">
   ];
 }

diff --git a/runtime/cudaq/platform/default/DefaultQuantumPlatform.cpp b/runtime/cudaq/platform/default/DefaultQuantumPlatform.cpp
@@ -39,11 +39,6 @@ class DefaultQPU : public cudaq::QPU {
     kernelFunc(args);
   }
 
-  void launchKernel(const std::string &name,
-                    const std::vector<void *> &) override {
-    throw std::runtime_error("Wrong kernel launch point.");
-  }
-
   /// Overrides setExecutionContext to forward it to the ExecutionManager
   void setExecutionContext(cudaq::ExecutionContext *context) override {
     ScopedTraceWithContext("DefaultPlatform::setExecutionContext",

diff --git a/runtime/cudaq/platform/mqpu/custatevec/GPUEmulatedQPU.cpp b/runtime/cudaq/platform/mqpu/custatevec/GPUEmulatedQPU.cpp
@@ -44,11 +44,6 @@ class GPUEmulatedQPU : public cudaq::QPU {
     kernelFunc(args);
   }
 
-  void launchKernel(const std::string &name,
-                    const std::vector<void *> &rawArgs) override {
-    throw std::runtime_error("not implemented");
-  }
-
   /// Overrides setExecutionContext to forward it to the ExecutionManager
   void setExecutionContext(cudaq::ExecutionContext *context) override {
     cudaSetDevice(qpu_id);

diff --git a/runtime/cudaq/platform/qpu.h b/runtime/cudaq/platform/qpu.h
@@ -174,8 +174,17 @@ class QPU : public registry::RegisteredType<QPU> {
   /// as a struct-packed void pointer and its corresponding size.
   virtual void launchKernel(const std::string &name, void (*kernelFunc)(void *),
                             void *args, std::uint64_t, std::uint64_t) = 0;
+
+  /// Launch the kernel with given name and argument arrays.
+  // This is intended for remote QPUs whereby we need to JIT-compile the kernel
+  // with argument synthesis. Remote QPU implementation to override this.
   virtual void launchKernel(const std::string &name,
-                            const std::vector<void *> &rawArgs) = 0;
+                            const std::vector<void *> &rawArgs) {
+    if (!isRemote())
+      throw std::runtime_error("Wrong kernel launch point: Attempt to launch "
+                               "kernel in streamlined for JIT mode on local "
+                               "simulated QPU. This is not supported.");
+  }
 
   /// Launch serialized code for remote execution. Subtypes that support this
   /// should override this function.

diff --git a/runtime/cudaq/platform/quantum_platform.cpp b/runtime/cudaq/platform/quantum_platform.cpp
@@ -223,3 +223,16 @@ void cudaq::streamlinedLaunchKernel(const char *kernelName,
   std::string kernName = kernelName;
   platform.launchKernel(kernName, rawArgs);
 }
+
+void cudaq::hybridLaunchKernel(const char *kernelName, void (*kernel)(void *),
+                               void *args, std::uint64_t argsSize,
+                               std::uint64_t resultOffset,
+                               const std::vector<void *> &rawArgs) {
+  ScopedTraceWithContext("hybridLaunchKernel", kernelName);
+  auto &platform = *cudaq::getQuantumPlatformInternal();
+  const std::string kernName = kernelName;
+  if (platform.is_remote(platform.get_current_qpu()))
+    platform.launchKernel(kernName, rawArgs);
+  else
+    platform.launchKernel(kernName, kernel, args, argsSize, resultOffset);
+}
diff --git a/test/Quake-QIR/argument.qke b/test/Quake-QIR/argument.qke
@@ -6,7 +6,7 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --kernel-execution --canonicalize %s | \
+// RUN: cudaq-opt --kernel-execution=codegen=1 --canonicalize %s | \
 // RUN: cudaq-translate --convert-to=qir | FileCheck %s
 
 // NB: the mangled name map is required for the kernel-execution pass.

diff --git a/test/Quake-QIR/return_values.qke b/test/Quake-QIR/return_values.qke
@@ -6,7 +6,7 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --add-dealloc --kernel-execution --canonicalize %s | \
+// RUN: cudaq-opt --add-dealloc --kernel-execution=codegen=1 --canonicalize %s | \
 // RUN: cudaq-translate --convert-to=qir | FileCheck %s
 
 // NB: the mangled name map is required for the kernel-execution pass.

diff --git a/test/Quake/kernel_exec-1.qke b/test/Quake/kernel_exec-1.qke
@@ -6,9 +6,9 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --kernel-execution %s | FileCheck %s
+// RUN: cudaq-opt --kernel-execution=codegen=1 %s | FileCheck %s
 // RUN: cudaq-opt --kernel-execution=codegen=2 %s | FileCheck --check-prefix=STREAM %s
-// RUN: cudaq-opt --kernel-execution=codegen=0 %s | FileCheck --check-prefix=HYBRID %s
+// RUN: cudaq-opt --kernel-execution %s | FileCheck --check-prefix=HYBRID %s
 
 module attributes {quake.mangled_name_map = {
   __nvqpp__mlirgen__ghz = "_ZN3ghzclEi"}} {

diff --git a/test/Quake/kernel_exec-2.qke b/test/Quake/kernel_exec-2.qke
@@ -6,7 +6,7 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --kernel-execution %s | FileCheck %s
+// RUN: cudaq-opt --kernel-execution=codegen=1 %s | FileCheck %s
 
 module attributes {quake.mangled_name_map = {
 __nvqpp__mlirgen__function_hawaiian = "shirt",

diff --git a/test/Quake/return_vector.qke b/test/Quake/return_vector.qke
@@ -6,7 +6,7 @@
 // the terms of the Apache License 2.0 which accompanies this distribution.   //
 // ========================================================================== //
 
-// RUN: cudaq-opt --add-dealloc --kernel-execution --canonicalize %s | \
+// RUN: cudaq-opt --add-dealloc --kernel-execution=codegen=1 --canonicalize %s | \
 // RUN: FileCheck %s
 
 // NB: the mangled name map is required for the kernel-execution pass.