Allow the environment variable to override the default target (NVIDIA…

…#877) * * Allow the environment variable to override the default target -- Added an environment variable - CUDAQ_DEFAULT_SIMULATOR -- Check value against valid simulators -- Added basic tests * Update tools/nvqpp/nvq++.in Co-authored-by: Eric Schweitz <[email protected]> * * Check environment variable and GPU availability in the LinkedLibraryHolder constructor. * * Clean up code from Python cudaq module (moved to LinkedLibraryHolder) * Update python/utils/LinkedLibraryHolder.cpp Co-authored-by: Ben Howe <[email protected]> * * Addressing review comments - enhanced tests * * Explicitly set simulator along with target * * Addressing review comments - Python tests modified for clarity * Initializing current target to default in Python * Update python/utils/LinkedLibraryHolder.cpp Co-authored-by: Eric Schweitz <[email protected]> * Update python/utils/LinkedLibraryHolder.h Co-authored-by: Eric Schweitz <[email protected]> * Update python/utils/LinkedLibraryHolder.h Co-authored-by: Eric Schweitz <[email protected]> * * Removing 'NULL' check --------- Co-authored-by: Eric Schweitz <[email protected]> Co-authored-by: Ben Howe <[email protected]>
svetlanasieber · Nov 10, 2023 · df34f7f · df34f7f
1 parent 0ccf994
commit df34f7f
Show file tree

Hide file tree

Showing 8 changed files with 231 additions and 46 deletions.
diff --git a/python/cudaq/__init__.py b/python/cudaq/__init__.py
@@ -9,7 +9,6 @@
 import sys
 import os, os.path
 from ._packages import *
-from ._query_gpu import is_gpu_available
 
 if not "CUDAQ_DYNLIBS" in os.environ:
     try:
@@ -29,10 +28,7 @@
 from ._pycudaq import *
 from .domains import chemistry
 
-initKwargs = {'target': 'qpp-cpu'}
-if is_gpu_available():
-    initKwargs = {'target': 'nvidia'}
-
+initKwargs = {}
 if '-target' in sys.argv:
     initKwargs['target'] = sys.argv[sys.argv.index('-target') + 1]
 

diff --git a/python/cudaq/_query_gpu.py b/python/cudaq/_query_gpu.py
diff --git a/python/tests/utils/target_env_var_check.py b/python/tests/utils/target_env_var_check.py
@@ -0,0 +1,74 @@
+# ============================================================================ #
+# Copyright (c) 2022 - 2023 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.     #
+# ============================================================================ #
+
+# RUN: PYTHONPATH=../../ pytest -rP  %s
+
+import os
+os.environ["CUDAQ_DEFAULT_SIMULATOR"] = "density-matrix-cpu"
+
+import pytest
+
+import cudaq
+
+def test_default_target():
+    """Tests the default target set by environment variable"""
+
+    assert ("density-matrix-cpu" == cudaq.get_target().name)
+
+    kernel = cudaq.make_kernel()
+    qubits = kernel.qalloc(2)
+    kernel.h(qubits[0])
+    kernel.cx(qubits[0], qubits[1])
+    kernel.mz(qubits)
+
+    result = cudaq.sample(kernel)
+    result.dump()
+    assert '00' in result
+    assert '11' in result
+
+def test_env_var_with_emulate():
+    """Tests the target when emulating a hardware backend"""
+
+    assert ("density-matrix-cpu" == cudaq.get_target().name)
+    cudaq.set_target("quantinuum", emulate=True)
+    assert ("quantinuum" == cudaq.get_target().name)
+
+    kernel = cudaq.make_kernel()
+    qubits = kernel.qalloc(2)
+    kernel.h(qubits[0])
+    kernel.cx(qubits[0], qubits[1])
+    kernel.mz(qubits)
+
+    result = cudaq.sample(kernel)
+    result.dump()
+    assert '00' in result
+    assert '11' in result
+
+def test_target_override():
+    """Tests the target set by environment variable is overridden by user setting"""
+
+    cudaq.set_target("qpp-cpu")
+    assert("qpp-cpu" == cudaq.get_target().name)
+
+    kernel = cudaq.make_kernel()
+    qubits = kernel.qalloc(2)
+    kernel.h(qubits[0])
+    kernel.cx(qubits[0], qubits[1])
+    kernel.mz(qubits)
+
+    result = cudaq.sample(kernel)
+    result.dump()
+    assert '00' in result
+    assert '11' in result
+
+os.environ.pop("CUDAQ_DEFAULT_SIMULATOR")
+
+# leave for gdb debugging
+if __name__ == "__main__":
+    loc = os.path.abspath(__file__)
+    pytest.main([loc, "-rP"])
diff --git a/python/tests/utils/target_env_var_reset.py b/python/tests/utils/target_env_var_reset.py
@@ -0,0 +1,46 @@
+# ============================================================================ #
+# Copyright (c) 2022 - 2023 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.     #
+# ============================================================================ #
+
+# RUN: PYTHONPATH=../../ pytest -rP  %s
+
+import os
+os.environ["CUDAQ_DEFAULT_SIMULATOR"] = "density-matrix-cpu"
+
+import pytest
+
+import cudaq
+
+def test_env_var_update():
+    """Tests that if the environment variable does not take effect on-the-fly"""
+
+    os.environ["CUDAQ_DEFAULT_SIMULATOR"] = "qpp-cpu"
+    assert("qpp-cpu" != cudaq.get_target().name)
+
+    cudaq.set_target("qpp-cpu")
+    assert("qpp-cpu" == cudaq.get_target().name)
+
+    kernel = cudaq.make_kernel()
+    qubits = kernel.qalloc(2)
+    kernel.h(qubits[0])
+    kernel.cx(qubits[0], qubits[1])
+    kernel.mz(qubits)
+
+    result = cudaq.sample(kernel)
+    result.dump()
+    assert '00' in result
+    assert '11' in result
+
+    cudaq.reset_target()
+    assert ("density-matrix-cpu" == cudaq.get_target().name)
+
+os.environ.pop("CUDAQ_DEFAULT_SIMULATOR")
+
+# leave for gdb debugging
+if __name__ == "__main__":
+    loc = os.path.abspath(__file__)
+    pytest.main([loc, "-rP"])
diff --git a/python/utils/LinkedLibraryHolder.cpp b/python/utils/LinkedLibraryHolder.cpp
@@ -30,6 +30,30 @@ constexpr static const char NVQIR_SIMULATION_BACKEND[] =
     "NVQIR_SIMULATION_BACKEND=";
 constexpr static const char TARGET_DESCRIPTION[] = "TARGET_DESCRIPTION=";
 
+int countGPUs() {
+  char buffer[1024];
+  std::string output;
+  FILE *fp1, *fp2;
+
+  fp1 = popen("nvidia-smi", "r");
+  if (!fp1) {
+    cudaq::info("nvidia-smi: command not found");
+    return -1;
+  }
+  pclose(fp1);
+
+  fp2 = popen("nvidia-smi -L | wc -l", "r");
+  if (!fp2) {
+    cudaq::info("nvidia-smi: command not working");
+    return -1;
+  }
+  while (fgets(buffer, sizeof buffer, fp2)) {
+    output += buffer;
+  }
+  pclose(fp2);
+  return std::stoi(output);
+}
+
 std::size_t RuntimeTarget::num_qpus() {
   auto &platform = cudaq::get_platform();
   return platform.num_qpus();
@@ -38,15 +62,16 @@ std::size_t RuntimeTarget::num_qpus() {
 /// @brief Search the targets folder in the install for available targets.
 void findAvailableTargets(
     const std::filesystem::path &targetPath,
-    std::unordered_map<std::string, RuntimeTarget> &targets) {
+    std::unordered_map<std::string, RuntimeTarget> &targets,
+    std::unordered_map<std::string, RuntimeTarget> &simulationTargets) {
 
   // Loop over all target files
   for (const auto &configFile :
        std::filesystem::directory_iterator{targetPath}) {
     auto path = configFile.path();
     // They must have a .config suffix
     if (path.extension().string() == ".config") {
-
+      bool isSimulationTarget = false;
       // Extract the target name from the file name
       auto fileName = path.filename().string();
       auto targetName = std::regex_replace(fileName, std::regex(".config"), "");
@@ -67,6 +92,7 @@ void findAvailableTargets(
                 std::regex_replace(platformName, std::regex("-"), "_");
 
           } else if (line.find(NVQIR_SIMULATION_BACKEND) != std::string::npos) {
+            isSimulationTarget = true;
             cudaq::trim(line);
             simulatorName = cudaq::split(line, '=')[1];
             // Post-process the string
@@ -91,6 +117,14 @@ void findAvailableTargets(
       // Add the target.
       targets.emplace(targetName, RuntimeTarget{targetName, simulatorName,
                                                 platformName, description});
+      if (isSimulationTarget) {
+        cudaq::info("Found Simulation target: {} -> (sim={}, platform={})",
+                    targetName, simulatorName, platformName);
+        simulationTargets.emplace(targetName,
+                                  RuntimeTarget{targetName, simulatorName,
+                                                platformName, description});
+        isSimulationTarget = false;
+      }
     }
   }
 }
@@ -116,7 +150,7 @@ LinkedLibraryHolder::LinkedLibraryHolder() {
 
   // Populate the map of available targets.
   auto targetPath = cudaqLibPath.parent_path() / "targets";
-  findAvailableTargets(targetPath, targets);
+  findAvailableTargets(targetPath, targets, simulationTargets);
 
   cudaq::info("Init: Library Path is {}.", cudaqLibPath.string());
 
@@ -212,16 +246,32 @@ LinkedLibraryHolder::LinkedLibraryHolder() {
     }
   }
 
-  targets.emplace("qpp-cpu",
-                  RuntimeTarget{"qpp-cpu", "qpp", "default",
-                                "QPP-based CPU-only simulated QPU."});
+  // Set the default target
+  // If environment variable set with a valid value, use it
+  // Otherwise, if GPU(s) available, set default to 'nvidia', else to 'qpp-cpu'
+  defaultTarget = "qpp-cpu";
+  if (countGPUs() > 0) {
+    defaultTarget = "nvidia";
+  }
+  auto env = std::getenv("CUDAQ_DEFAULT_SIMULATOR");
+  if (env) {
+    cudaq::info("'CUDAQ_DEFAULT_SIMULATOR' = {}", env);
+    auto iter = simulationTargets.find(env);
+    if (iter != simulationTargets.end()) {
+      cudaq::info("Valid target");
+      defaultTarget = iter->second.name;
+    }
+  }
+
+  // Initialize current target to default, may be overridden by command line
+  // argument or set_target() API
+  currentTarget = defaultTarget;
 
   if (disallowTargetModification)
     return;
 
-  // We'll always start off with the default platform and the QPP simulator
-  __nvqir__setCircuitSimulator(getSimulator("qpp"));
-  setQuantumPlatformInternal(getPlatform("default"));
+  // We'll always start off with the default target
+  resetTarget();
 }
 
 LinkedLibraryHolder::~LinkedLibraryHolder() {
@@ -253,7 +303,7 @@ LinkedLibraryHolder::getPlatform(const std::string &platformName) {
       std::string("getQuantumPlatform_") + platformName);
 }
 
-void LinkedLibraryHolder::resetTarget() { setTarget("qpp-cpu"); }
+void LinkedLibraryHolder::resetTarget() { setTarget(defaultTarget); }
 
 RuntimeTarget LinkedLibraryHolder::getTarget(const std::string &name) const {
   auto iter = targets.find(name);

diff --git a/python/utils/LinkedLibraryHolder.h b/python/utils/LinkedLibraryHolder.h
@@ -64,8 +64,14 @@ class LinkedLibraryHolder {
   /// @brief Map of available targets.
   std::unordered_map<std::string, RuntimeTarget> targets;
 
+  /// @brief Map of simulation targets
+  std::unordered_map<std::string, RuntimeTarget> simulationTargets;
+
+  /// @brief Store the name of the default target
+  std::string defaultTarget;
+
   /// @brief Store the name of the current target
-  std::string currentTarget = "qpp-cpu";
+  std::string currentTarget;
 
 public:
   LinkedLibraryHolder();

diff --git a/test/NVQPP/qpp_cpu_target.cpp b/test/NVQPP/qpp_cpu_target.cpp
@@ -6,7 +6,10 @@
  * the terms of the Apache License 2.0 which accompanies this distribution.    *
  ******************************************************************************/
 
-// RUN: nvq++ --target qpp-cpu %s -o %basename_t.x && ./%basename_t.x
+// RUN: nvq++ --target qpp-cpu %s -o %basename_t.x && CUDAQ_LOG_LEVEL=info ./%basename_t.x | FileCheck --check-prefix=CHECK-QPP %s
+// RUN: CUDAQ_DEFAULT_SIMULATOR="density-matrix-cpu" nvq++ %s -o %basename_t.x && CUDAQ_LOG_LEVEL=info ./%basename_t.x | FileCheck --check-prefix=CHECK-DM %s
+// RUN: CUDAQ_DEFAULT_SIMULATOR="foo" nvq++ %s -o %basename_t.x && CUDAQ_LOG_LEVEL=info ./%basename_t.x | FileCheck %s
+// RUN: CUDAQ_DEFAULT_SIMULATOR="qpp-cpu" nvq++ --target quantinuum --emulate %s -o %basename_t.x && CUDAQ_LOG_LEVEL=info ./%basename_t.x | FileCheck --check-prefix=CHECK-QPP %s
 
 #include <cudaq.h>
 
@@ -26,3 +29,11 @@ int main() {
   counts.dump();
   return 0;
 }
+
+// CHECK-QPP: [info] [NVQIR.cpp:{{[0-9]+}}] Creating the qpp backend.
+// CHECK-QPP: [info] [DefaultExecutionManager.cpp:{{[0-9]+}}] [DefaultExecutionManager] Creating the qpp backend.
+
+// CHECK-DM: [info] [NVQIR.cpp:{{[0-9]+}}] Creating the dm backend.
+// CHECK-DM: [info] [DefaultExecutionManager.cpp:{{[0-9]+}}] [DefaultExecutionManager] Creating the dm backend.
+
+// CHECK-NOT: foo
diff --git a/tools/nvqpp/nvq++.in b/tools/nvqpp/nvq++.in
@@ -39,10 +39,28 @@ function add_pass_to_pipeline {
 }
 
 function list_targets {
-	ls -I *.cpp -1 ${install_dir}/targets/ | sed -e 's/\.config$//'
+	ls -I *.cpp -1 ${install_dir}/targets/ | grep ".config$" | sed -e 's/\.config$//'
 	exit 0
 }
 
+function list_simulators {
+	for file in $(grep -L "PLATFORM_QPU=" $(ls ${install_dir}/targets/*.config)) ; do
+		if ! grep -q "LIBRARY_MODE_EXECUTION_MANAGER=" $file ; then 
+			echo $(basename $file | cut -d "." -f 1)
+		fi
+	done
+}
+
+function get_simulation_backend {
+	config_file="${install_dir}/targets/$1.config"
+	if [ -f "$config_file" ]; then
+    	line=$(grep "NVQIR_SIMULATION_BACKEND=" "$config_file")
+		if [ $? -eq 0 ]; then
+			echo ${line#*=} | tr -d '"'
+		fi
+	fi
+}
+
 function query_gpu {
 	if [ -x "$(command -v nvidia-smi)" ]; then
 		# Make sure nvidia-smi works.
@@ -260,6 +278,18 @@ gpu_found=$(query_gpu)
 if ${gpu_found} && [ -f "${install_dir}/lib/libnvqir-custatevec-fp32.so" ]; then
 	NVQIR_SIMULATION_BACKEND="custatevec-fp32"
 fi
+# Check environment variable - overrides the default
+if [[ ! -z "${CUDAQ_DEFAULT_SIMULATOR}" ]]; then
+	available_simulators=( $(list_simulators) )
+	for s in "${available_simulators[@]}"
+	do
+		if [ "${CUDAQ_DEFAULT_SIMULATOR}" = "$s" ]; then
+			TARGET_CONFIG="${CUDAQ_DEFAULT_SIMULATOR}"
+			NVQIR_SIMULATION_BACKEND=$(get_simulation_backend "$s")
+			break
+		fi
+	done
+fi
 
 # We default to LIBRARY_MODE, physical 
 # Quantum Targets can override this to turn on