Adding unittests for cudm_op_conversion

Signed-off-by: Sachin Pisal <[email protected]>

runtime/cudaq/cudm_helpers.h

@@ -17,6 +17,8 @@
 #include <vector>
 
 namespace cudaq {
+std::vector<std::complex<double>> flatten_matrix(const matrix_2 &matrix);
+
 void scale_state(cudensitymatHandle_t handle, cudensitymatState_t state,
                  double scale_factor, cudaStream_t stream);
 

runtime/cudaq/cudm_op_conversion.h

@@ -39,11 +39,13 @@ class cudm_op_conversion {
 
   // Addition of two operator terms
   std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,
-               double>
+               std::complex<double>>
   add(const std::variant<cudensitymatOperatorTerm_t,
-                         cudensitymatWrappedScalarCallback_t, double> &op1,
+                         cudensitymatWrappedScalarCallback_t,
+                         std::complex<double>> &op1,
       const std::variant<cudensitymatOperatorTerm_t,
-                         cudensitymatWrappedScalarCallback_t, double> &op2);
+                         cudensitymatWrappedScalarCallback_t,
+                         std::complex<double>> &op2);
 
   // Evaluate an operator and convert it to cudensitymatOperatorTerm_t
   std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,

runtime/cudaq/dynamics/cudm_helpers.cpp

@@ -9,6 +9,8 @@
 #include "cudaq/cudm_helpers.h"
 #include "cudaq/cudm_error_handling.h"
 
+using namespace cudaq;
+
 namespace cudaq {
 // Function to flatten a matrix into a 1D array (column major)
 std::vector<std::complex<double>> flatten_matrix(const matrix_2 &matrix) {
@@ -51,6 +53,9 @@ cudensitymatElementaryOperator_t create_elementary_operator(
     throw std::invalid_argument("subspace_extents cannot be empty.");
   }
 
+  std::cout << "Subspace extents size: " << subspace_extents.size()
+            << ", Matrix size: " << flat_matrix.size() << std::endl;
+
   cudensitymatElementaryOperator_t cudm_elem_op = nullptr;
 
   // FIXME: leak (need to track this buffer somewhere and delete **after** the
@@ -84,8 +89,20 @@ void append_elementary_operator_to_term(
     throw std::invalid_argument("elem_op cannot be null.");
   }
 
+  std::cout << "Appending Elementary Operator to Term"
+            << " - Degrees: " << degrees.size() << " - Term: " << term
+            << std::endl;
+
+  for (int degree : degrees) {
+    if (degree < 0) {
+      throw std::out_of_range("Degree cannot be negative!");
+    }
+  }
+
   std::vector<cudensitymatElementaryOperator_t> elem_ops = {elem_op};
 
+  std::cout << "elem_ops.data(): " << elem_ops.data() << std::endl;
+
   std::vector<int32_t> modeActionDuality(degrees.size(), 0);
   assert(elem_ops.size() == degrees.size());
   HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(
@@ -111,8 +128,19 @@ void scale_state(cudensitymatHandle_t handle, cudensitymatState_t state,
     throw std::invalid_argument("Invalid state provided to scale_state.");
   }
 
+  cudaEvent_t start, stop;
+  cudaEventCreate(&start);
+  cudaEventCreate(&stop);
+  cudaEventRecord(start, stream);
+
   HANDLE_CUDM_ERROR(
       cudensitymatStateComputeScaling(handle, state, &scale_factor, stream));
+
+  cudaEventRecord(stop, stream);
+  cudaEventSynchronize(stop);
+  float milliseconds = 0;
+  cudaEventElapsedTime(&milliseconds, start, stop);
+  std::cout << "Time taken: " << milliseconds << " ms" << std::endl;
 }
 
 cudensitymatOperator_t

runtime/cudaq/dynamics/cudm_op_conversion.cpp

@@ -18,7 +18,16 @@ namespace cudaq {
 cudm_op_conversion::cudm_op_conversion(const cudensitymatHandle_t handle,
                                        const std::map<int, int> &dimensions,
                                        std::shared_ptr<Schedule> schedule)
-    : handle_(handle), dimensions_(dimensions), schedule_(schedule) {}
+    : handle_(handle), dimensions_(dimensions), schedule_(schedule) {
+  std::cout << "Handle: " << handle_ << std::endl;
+  if (handle_ == nullptr) {
+    throw std::runtime_error("Handle cannot be null.");
+  }
+
+  if (dimensions_.empty()) {
+    throw std::invalid_argument("Dimensions map must not be empty.");
+  }
+}
 
 cudensitymatOperatorTerm_t cudm_op_conversion::_scalar_to_op(
     const cudensitymatWrappedScalarCallback_t &scalar) {
@@ -91,19 +100,37 @@ cudm_op_conversion::mul(
 }
 
 std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,
-             double>
-cudm_op_conversion::add(
-    const std::variant<cudensitymatOperatorTerm_t,
-                       cudensitymatWrappedScalarCallback_t, double> &op1,
-    const std::variant<cudensitymatOperatorTerm_t,
-                       cudensitymatWrappedScalarCallback_t, double> &op2) {
-  if (std::holds_alternative<double>(op1) ||
-      std::holds_alternative<double>(op2)) {
-    return std::get<double>(op1) + std::get<double>(op2);
+             std::complex<double>>
+cudm_op_conversion::add(const std::variant<cudensitymatOperatorTerm_t,
+                                           cudensitymatWrappedScalarCallback_t,
+                                           std::complex<double>> &op1,
+                        const std::variant<cudensitymatOperatorTerm_t,
+                                           cudensitymatWrappedScalarCallback_t,
+                                           std::complex<double>> &op2) {
+  if (std::holds_alternative<std::complex<double>>(op1) &&
+      std::holds_alternative<std::complex<double>>(op2)) {
+    return std::get<std::complex<double>>(op1) +
+           std::get<std::complex<double>>(op2);
   }
 
+  if (std::holds_alternative<std::complex<double>>(op1)) {
+    return _callback_mult_op(
+        _wrap_callback(scalar_operator(std::get<std::complex<double>>(op1))),
+        std::get<cudensitymatOperatorTerm_t>(op2));
+  }
+
+  if (std::holds_alternative<std::complex<double>>(op2)) {
+    return _callback_mult_op(
+        _wrap_callback(scalar_operator(std::get<std::complex<double>>(op2))),
+        std::get<cudensitymatOperatorTerm_t>(op1));
+  }
+
+  // FIXME: Need to check later
+  int32_t num_space_modes =
+      std::max(static_cast<int32_t>(dimensions_.size()), 1);
+
   cudensitymatOperatorTerm_t result;
-  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle_, dimensions_.size(),
+  HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle_, num_space_modes,
                                                    nullptr, &result));
 
   HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(
@@ -136,18 +163,31 @@ cudm_op_conversion::evaluate(
   if (std::holds_alternative<matrix_operator>(op)) {
     const matrix_operator &mat_op = std::get<matrix_operator>(op);
 
+    std::vector<int64_t> space_mode_extents;
+    for (const auto &dim : dimensions_) {
+      space_mode_extents.push_back(dim.second);
+    }
+
     cudensitymatOperatorTerm_t opterm;
     HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(
-        handle_, dimensions_.size(), nullptr, &opterm));
+        handle_, dimensions_.size(), space_mode_extents.data(), &opterm));
 
     cudensitymatElementaryOperator_t elem_op;
     cudensitymatWrappedTensorCallback_t callback =
         _wrap_callback_tensor(mat_op);
 
+    auto flat_matrix = flatten_matrix(mat_op.to_matrix(dimensions_, {}));
+
+    void *tensor_data = create_array_gpu(flat_matrix);
+    if (!tensor_data) {
+      throw std::runtime_error(
+          "Failed to allocate GPU memory for tensor_data.");
+    }
+
     HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(
-        handle_, mat_op.degrees.size(), nullptr,
-        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, 0, nullptr, CUDA_C_64F, nullptr,
-        callback, &elem_op));
+        handle_, mat_op.degrees.size(), space_mode_extents.data(),
+        CUDENSITYMAT_OPERATOR_SPARSITY_NONE, 0, nullptr, CUDA_C_64F,
+        tensor_data, callback, &elem_op));
 
     HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(
         handle_, opterm, 1, &elem_op, mat_op.degrees.data(), nullptr,

runtime/cudaq/dynamics/cudm_state.cpp

@@ -28,7 +28,7 @@ cudm_state::cudm_state(cudensitymatHandle_t handle,
 
   // Allocate device memory
   size_t dataSize = rawData_.size() * sizeof(std::complex<double>);
-  cudaMalloc(reinterpret_cast<void **>(&gpuData_), dataSize);
+  HANDLE_CUDA_ERROR(cudaMalloc(reinterpret_cast<void **>(&gpuData_), dataSize));
 
   // Copy data from host to device
   HANDLE_CUDA_ERROR(

unittests/CMakeLists.txt

@@ -295,6 +295,7 @@ if (CUDA_FOUND)
     dynamics/test_cudm_helpers.cpp
     dynamics/test_cudm_state.cpp
     dynamics/test_cudm_time_stepper.cpp
+    dynamics/test_cudm_op_conversion.cpp
   )
   add_executable(test_dynamics main.cpp ${CUDAQ_DYNAMICS_TEST_SOURCES})
   if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT APPLE)

unittests/dynamics/test_cudm_helpers.cpp

@@ -50,7 +50,7 @@ TEST_F(CuDensityMatHelpersTestFixture, ScaleState) {
 
   ASSERT_TRUE(state.is_initialized());
 
-  EXPECT_NO_THROW(cudaq::scale_state(handle, state.get_impl(), 2.0, stream));
+  EXPECT_NO_THROW(cudaq::scale_state(handle, state.get_impl(), {2.0}, stream));
 }
 
 // Test for compute_lindblad_op

unittests/dynamics/test_cudm_op_conversion.cpp

@@ -0,0 +1,129 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2025 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.    *
+ ******************************************************************************/
+
+#include <gtest/gtest.h>
+#include "cudaq/cudm_op_conversion.h"
+#include "cudaq/operators.h"
+#include "test_mocks.h"
+#include <stdexcept>
+#include <complex>
+#include <memory>
+
+using namespace cudaq;
+
+class CuDmOpConversion : public ::testing::Test {
+protected:
+    cudensitymatHandle_t handle;
+    std::map<int, int> dimensions;
+    std::shared_ptr<cudaq::Schedule> schedule;
+    std::unique_ptr<cudm_op_conversion> converter;
+    std::vector<int64_t> space_mode_extents;
+
+    void SetUp() override {
+        handle = mock_handle();
+        dimensions = {{0, 2}, {1, 2}};
+        for (const auto &dim : dimensions) {
+            space_mode_extents.push_back(dim.second);
+        }
+        schedule = std::shared_ptr<Schedule>();
+        converter = std::make_unique<cudm_op_conversion>(handle, dimensions, schedule);
+    }
+};
+
+TEST_F(CuDmOpConversion, ConstructorValid) {
+    EXPECT_NO_THROW(cudm_op_conversion converter(handle, dimensions, schedule));
+}
+
+TEST_F(CuDmOpConversion, ConstructorEmptyDimensions) {
+    std::map<int, int> empty_dimensions;
+    EXPECT_THROW(cudm_op_conversion converter(handle, empty_dimensions, schedule), std::invalid_argument);
+}
+
+TEST_F(CuDmOpConversion, ConstructorInvalidHandle) {
+    cudensitymatHandle_t invalid_handle = nullptr;
+    EXPECT_THROW(cudm_op_conversion converter(invalid_handle, dimensions, schedule), std::runtime_error);
+}
+
+TEST_F(CuDmOpConversion, EvaluateScalarConstant) {
+    scalar_operator scalar_op(2.5);
+    auto result = converter->evaluate(scalar_op);
+
+    ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
+    EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(2.5, 0.0));
+}
+
+TEST_F(CuDmOpConversion, EvaluateScalarCallback) {
+    scalar_operator scalar_op([](std::map<std::string, std::complex<double>>) {
+        return std::complex<double>(1.0, -1.0);
+    });
+    auto result = converter->evaluate(scalar_op);
+
+    ASSERT_TRUE(std::holds_alternative<cudensitymatWrappedScalarCallback_t>(result));
+}
+
+// TEST_F(CuDmOpConversion, EvaluateMatrixOperator) {
+//     matrix_operator mat_op("H", {0});
+//     auto result = converter->evaluate(mat_op);
+
+//     ASSERT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(result));
+// }
+
+TEST_F(CuDmOpConversion, EvaluateProductOperator) {
+    auto op0 = cudaq::matrix_operator::annihilate(0);
+    auto op1 = cudaq::matrix_operator::create(0);
+    product_operator<matrix_operator> product_op = op0 * op1;
+    EXPECT_THROW(converter->evaluate(product_op), std::runtime_error);
+}
+
+TEST_F(CuDmOpConversion, AddOperators) {
+    scalar_operator scalar_op1(2.0);
+    scalar_operator scalar_op2(3.0);
+
+    auto result = converter->add(converter->evaluate(scalar_op1), converter->evaluate(scalar_op2));
+
+    ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
+    EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(5.0, 0.0));
+}
+
+TEST_F(CuDmOpConversion, AddComplexScalars) {
+    std::complex<double> scalar_1(2.0, 1.0);
+    std::complex<double> scalar_2(3.0, -1.0);
+
+    auto result = converter->add(scalar_1, scalar_2);
+
+    ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
+    EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(5.0, 0.0));
+}
+
+// TEST_F(CuDmOpConversion, AddScalarAndOperator) {
+//     scalar_operator scalar_op(1.0);
+//     matrix_operator mat_op("X", {0});
+
+//     auto scalar_result = converter->evaluate(scalar_op);
+//     auto op_result = converter->evaluate(mat_op);
+
+//     auto final_result = converter->add(scalar_result, op_result);
+
+//     ASSERT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(final_result));
+// }
+
+TEST_F(CuDmOpConversion, TensorProductOfScalars) {
+    auto result = converter->tensor(2.0, 3.0);
+    EXPECT_TRUE(std::holds_alternative<double>(result));
+    EXPECT_EQ(std::get<double>(result), 6.0);
+}
+
+// TEST_F(CuDmOpConversion, TensorProductScalarAndOperator) {
+//     cudensitymatOperatorTerm_t op_term;
+//     HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle, dimensions.size(), space_mode_extents.data(), &op_term));
+
+//     auto result = converter->tensor(2.0, op_term);
+//     EXPECT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(result));
+
+//     HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(op_term));
+// }

unittests/dynamics/test_mocks.h

@@ -14,6 +14,13 @@
 #include <cudensitymat.h>
 #include <iostream>
 
+// Mock cudensitymatHandle_t
+inline cudensitymatHandle_t mock_handle() {
+  cudensitymatHandle_t handle;
+  HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
+  return handle;
+}
+
 // Mock Liouvillian operator creation
 inline cudensitymatOperator_t mock_liouvillian(cudensitymatHandle_t handle) {
   cudensitymatOperator_t liouvillian = nullptr;