Skip to content

Commit 3d8b598

Browse files
sacpissacpis
committed
Adding unittests for cudm_op_conversion
Signed-off-by: Sachin Pisal <[email protected]>
1 parent bc43e40 commit 3d8b598

File tree

9 files changed

+229
-20
lines changed

9 files changed

+229
-20
lines changed

‎runtime/cudaq/cudm_helpers.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
#include <vector>
1818

1919
namespace cudaq {
20+
std::vector<std::complex<double>> flatten_matrix(const matrix_2 &matrix);
21+
2022
void scale_state(cudensitymatHandle_t handle, cudensitymatState_t state,
2123
double scale_factor, cudaStream_t stream);
2224

‎runtime/cudaq/cudm_op_conversion.h

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,13 @@ class cudm_op_conversion {
3939

4040
// Addition of two operator terms
4141
std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,
42-
double>
42+
std::complex<double>>
4343
add(const std::variant<cudensitymatOperatorTerm_t,
44-
cudensitymatWrappedScalarCallback_t, double> &op1,
44+
cudensitymatWrappedScalarCallback_t,
45+
std::complex<double>> &op1,
4546
const std::variant<cudensitymatOperatorTerm_t,
46-
cudensitymatWrappedScalarCallback_t, double> &op2);
47+
cudensitymatWrappedScalarCallback_t,
48+
std::complex<double>> &op2);
4749

4850
// Evaluate an operator and convert it to cudensitymatOperatorTerm_t
4951
std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,

‎runtime/cudaq/dynamics/cudm_helpers.cpp

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
#include "cudaq/cudm_helpers.h"
1010
#include "cudaq/cudm_error_handling.h"
1111

12+
using namespace cudaq;
13+
1214
namespace cudaq {
1315
// Function to flatten a matrix into a 1D array (column major)
1416
std::vector<std::complex<double>> flatten_matrix(const matrix_2 &matrix) {
@@ -51,6 +53,9 @@ cudensitymatElementaryOperator_t create_elementary_operator(
5153
throw std::invalid_argument("subspace_extents cannot be empty.");
5254
}
5355

56+
std::cout << "Subspace extents size: " << subspace_extents.size()
57+
<< ", Matrix size: " << flat_matrix.size() << std::endl;
58+
5459
cudensitymatElementaryOperator_t cudm_elem_op = nullptr;
5560

5661
// FIXME: leak (need to track this buffer somewhere and delete **after** the
@@ -84,8 +89,20 @@ void append_elementary_operator_to_term(
8489
throw std::invalid_argument("elem_op cannot be null.");
8590
}
8691

92+
std::cout << "Appending Elementary Operator to Term"
93+
<< " - Degrees: " << degrees.size() << " - Term: " << term
94+
<< std::endl;
95+
96+
for (int degree : degrees) {
97+
if (degree < 0) {
98+
throw std::out_of_range("Degree cannot be negative!");
99+
}
100+
}
101+
87102
std::vector<cudensitymatElementaryOperator_t> elem_ops = {elem_op};
88103

104+
std::cout << "elem_ops.data(): " << elem_ops.data() << std::endl;
105+
89106
std::vector<int32_t> modeActionDuality(degrees.size(), 0);
90107
assert(elem_ops.size() == degrees.size());
91108
HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(
@@ -111,8 +128,19 @@ void scale_state(cudensitymatHandle_t handle, cudensitymatState_t state,
111128
throw std::invalid_argument("Invalid state provided to scale_state.");
112129
}
113130

131+
cudaEvent_t start, stop;
132+
cudaEventCreate(&start);
133+
cudaEventCreate(&stop);
134+
cudaEventRecord(start, stream);
135+
114136
HANDLE_CUDM_ERROR(
115137
cudensitymatStateComputeScaling(handle, state, &scale_factor, stream));
138+
139+
cudaEventRecord(stop, stream);
140+
cudaEventSynchronize(stop);
141+
float milliseconds = 0;
142+
cudaEventElapsedTime(&milliseconds, start, stop);
143+
std::cout << "Time taken: " << milliseconds << " ms" << std::endl;
116144
}
117145

118146
cudensitymatOperator_t

‎runtime/cudaq/dynamics/cudm_op_conversion.cpp

Lines changed: 55 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,16 @@ namespace cudaq {
1818
cudm_op_conversion::cudm_op_conversion(const cudensitymatHandle_t handle,
1919
const std::map<int, int> &dimensions,
2020
std::shared_ptr<Schedule> schedule)
21-
: handle_(handle), dimensions_(dimensions), schedule_(schedule) {}
21+
: handle_(handle), dimensions_(dimensions), schedule_(schedule) {
22+
std::cout << "Handle: " << handle_ << std::endl;
23+
if (handle_ == nullptr) {
24+
throw std::runtime_error("Handle cannot be null.");
25+
}
26+
27+
if (dimensions_.empty()) {
28+
throw std::invalid_argument("Dimensions map must not be empty.");
29+
}
30+
}
2231

2332
cudensitymatOperatorTerm_t cudm_op_conversion::_scalar_to_op(
2433
const cudensitymatWrappedScalarCallback_t &scalar) {
@@ -91,19 +100,37 @@ cudm_op_conversion::mul(
91100
}
92101

93102
std::variant<cudensitymatOperatorTerm_t, cudensitymatWrappedScalarCallback_t,
94-
double>
95-
cudm_op_conversion::add(
96-
const std::variant<cudensitymatOperatorTerm_t,
97-
cudensitymatWrappedScalarCallback_t, double> &op1,
98-
const std::variant<cudensitymatOperatorTerm_t,
99-
cudensitymatWrappedScalarCallback_t, double> &op2) {
100-
if (std::holds_alternative<double>(op1) ||
101-
std::holds_alternative<double>(op2)) {
102-
return std::get<double>(op1) + std::get<double>(op2);
103+
std::complex<double>>
104+
cudm_op_conversion::add(const std::variant<cudensitymatOperatorTerm_t,
105+
cudensitymatWrappedScalarCallback_t,
106+
std::complex<double>> &op1,
107+
const std::variant<cudensitymatOperatorTerm_t,
108+
cudensitymatWrappedScalarCallback_t,
109+
std::complex<double>> &op2) {
110+
if (std::holds_alternative<std::complex<double>>(op1) &&
111+
std::holds_alternative<std::complex<double>>(op2)) {
112+
return std::get<std::complex<double>>(op1) +
113+
std::get<std::complex<double>>(op2);
103114
}
104115

116+
if (std::holds_alternative<std::complex<double>>(op1)) {
117+
return _callback_mult_op(
118+
_wrap_callback(scalar_operator(std::get<std::complex<double>>(op1))),
119+
std::get<cudensitymatOperatorTerm_t>(op2));
120+
}
121+
122+
if (std::holds_alternative<std::complex<double>>(op2)) {
123+
return _callback_mult_op(
124+
_wrap_callback(scalar_operator(std::get<std::complex<double>>(op2))),
125+
std::get<cudensitymatOperatorTerm_t>(op1));
126+
}
127+
128+
// FIXME: Need to check later
129+
int32_t num_space_modes =
130+
std::max(static_cast<int32_t>(dimensions_.size()), 1);
131+
105132
cudensitymatOperatorTerm_t result;
106-
HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle_, dimensions_.size(),
133+
HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle_, num_space_modes,
107134
nullptr, &result));
108135

109136
HANDLE_CUDM_ERROR(cudensitymatOperatorAppendTerm(
@@ -136,18 +163,31 @@ cudm_op_conversion::evaluate(
136163
if (std::holds_alternative<matrix_operator>(op)) {
137164
const matrix_operator &mat_op = std::get<matrix_operator>(op);
138165

166+
std::vector<int64_t> space_mode_extents;
167+
for (const auto &dim : dimensions_) {
168+
space_mode_extents.push_back(dim.second);
169+
}
170+
139171
cudensitymatOperatorTerm_t opterm;
140172
HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(
141-
handle_, dimensions_.size(), nullptr, &opterm));
173+
handle_, dimensions_.size(), space_mode_extents.data(), &opterm));
142174

143175
cudensitymatElementaryOperator_t elem_op;
144176
cudensitymatWrappedTensorCallback_t callback =
145177
_wrap_callback_tensor(mat_op);
146178

179+
auto flat_matrix = flatten_matrix(mat_op.to_matrix(dimensions_, {}));
180+
181+
void *tensor_data = create_array_gpu(flat_matrix);
182+
if (!tensor_data) {
183+
throw std::runtime_error(
184+
"Failed to allocate GPU memory for tensor_data.");
185+
}
186+
147187
HANDLE_CUDM_ERROR(cudensitymatCreateElementaryOperator(
148-
handle_, mat_op.degrees.size(), nullptr,
149-
CUDENSITYMAT_OPERATOR_SPARSITY_NONE, 0, nullptr, CUDA_C_64F, nullptr,
150-
callback, &elem_op));
188+
handle_, mat_op.degrees.size(), space_mode_extents.data(),
189+
CUDENSITYMAT_OPERATOR_SPARSITY_NONE, 0, nullptr, CUDA_C_64F,
190+
tensor_data, callback, &elem_op));
151191

152192
HANDLE_CUDM_ERROR(cudensitymatOperatorTermAppendElementaryProduct(
153193
handle_, opterm, 1, &elem_op, mat_op.degrees.data(), nullptr,

‎runtime/cudaq/dynamics/cudm_state.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ cudm_state::cudm_state(cudensitymatHandle_t handle,
2828

2929
// Allocate device memory
3030
size_t dataSize = rawData_.size() * sizeof(std::complex<double>);
31-
cudaMalloc(reinterpret_cast<void **>(&gpuData_), dataSize);
31+
HANDLE_CUDA_ERROR(cudaMalloc(reinterpret_cast<void **>(&gpuData_), dataSize));
3232

3333
// Copy data from host to device
3434
HANDLE_CUDA_ERROR(

‎unittests/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,7 @@ if (CUDA_FOUND)
295295
dynamics/test_cudm_helpers.cpp
296296
dynamics/test_cudm_state.cpp
297297
dynamics/test_cudm_time_stepper.cpp
298+
dynamics/test_cudm_op_conversion.cpp
298299
)
299300
add_executable(test_dynamics main.cpp ${CUDAQ_DYNAMICS_TEST_SOURCES})
300301
if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND NOT APPLE)

‎unittests/dynamics/test_cudm_helpers.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ TEST_F(CuDensityMatHelpersTestFixture, ScaleState) {
5050

5151
ASSERT_TRUE(state.is_initialized());
5252

53-
EXPECT_NO_THROW(cudaq::scale_state(handle, state.get_impl(), 2.0, stream));
53+
EXPECT_NO_THROW(cudaq::scale_state(handle, state.get_impl(), {2.0}, stream));
5454
}
5555

5656
// Test for compute_lindblad_op

‎unittests/dynamics/test_cudm_op_conversion.cpp

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*******************************************************************************
2+
* Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates. *
3+
* All rights reserved. *
4+
* *
5+
* This source code and the accompanying materials are made available under *
6+
* the terms of the Apache License 2.0 which accompanies this distribution. *
7+
******************************************************************************/
8+
9+
#include <gtest/gtest.h>
10+
#include "cudaq/cudm_op_conversion.h"
11+
#include "cudaq/operators.h"
12+
#include "test_mocks.h"
13+
#include <stdexcept>
14+
#include <complex>
15+
#include <memory>
16+
17+
using namespace cudaq;
18+
19+
class CuDmOpConversion : public ::testing::Test {
20+
protected:
21+
cudensitymatHandle_t handle;
22+
std::map<int, int> dimensions;
23+
std::shared_ptr<cudaq::Schedule> schedule;
24+
std::unique_ptr<cudm_op_conversion> converter;
25+
std::vector<int64_t> space_mode_extents;
26+
27+
void SetUp() override {
28+
handle = mock_handle();
29+
dimensions = {{0, 2}, {1, 2}};
30+
for (const auto &dim : dimensions) {
31+
space_mode_extents.push_back(dim.second);
32+
}
33+
schedule = std::shared_ptr<Schedule>();
34+
converter = std::make_unique<cudm_op_conversion>(handle, dimensions, schedule);
35+
}
36+
};
37+
38+
TEST_F(CuDmOpConversion, ConstructorValid) {
39+
EXPECT_NO_THROW(cudm_op_conversion converter(handle, dimensions, schedule));
40+
}
41+
42+
TEST_F(CuDmOpConversion, ConstructorEmptyDimensions) {
43+
std::map<int, int> empty_dimensions;
44+
EXPECT_THROW(cudm_op_conversion converter(handle, empty_dimensions, schedule), std::invalid_argument);
45+
}
46+
47+
TEST_F(CuDmOpConversion, ConstructorInvalidHandle) {
48+
cudensitymatHandle_t invalid_handle = nullptr;
49+
EXPECT_THROW(cudm_op_conversion converter(invalid_handle, dimensions, schedule), std::runtime_error);
50+
}
51+
52+
TEST_F(CuDmOpConversion, EvaluateScalarConstant) {
53+
scalar_operator scalar_op(2.5);
54+
auto result = converter->evaluate(scalar_op);
55+
56+
ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
57+
EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(2.5, 0.0));
58+
}
59+
60+
TEST_F(CuDmOpConversion, EvaluateScalarCallback) {
61+
scalar_operator scalar_op([](std::map<std::string, std::complex<double>>) {
62+
return std::complex<double>(1.0, -1.0);
63+
});
64+
auto result = converter->evaluate(scalar_op);
65+
66+
ASSERT_TRUE(std::holds_alternative<cudensitymatWrappedScalarCallback_t>(result));
67+
}
68+
69+
// TEST_F(CuDmOpConversion, EvaluateMatrixOperator) {
70+
// matrix_operator mat_op("H", {0});
71+
// auto result = converter->evaluate(mat_op);
72+
73+
// ASSERT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(result));
74+
// }
75+
76+
TEST_F(CuDmOpConversion, EvaluateProductOperator) {
77+
auto op0 = cudaq::matrix_operator::annihilate(0);
78+
auto op1 = cudaq::matrix_operator::create(0);
79+
product_operator<matrix_operator> product_op = op0 * op1;
80+
EXPECT_THROW(converter->evaluate(product_op), std::runtime_error);
81+
}
82+
83+
TEST_F(CuDmOpConversion, AddOperators) {
84+
scalar_operator scalar_op1(2.0);
85+
scalar_operator scalar_op2(3.0);
86+
87+
auto result = converter->add(converter->evaluate(scalar_op1), converter->evaluate(scalar_op2));
88+
89+
ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
90+
EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(5.0, 0.0));
91+
}
92+
93+
TEST_F(CuDmOpConversion, AddComplexScalars) {
94+
std::complex<double> scalar_1(2.0, 1.0);
95+
std::complex<double> scalar_2(3.0, -1.0);
96+
97+
auto result = converter->add(scalar_1, scalar_2);
98+
99+
ASSERT_TRUE(std::holds_alternative<std::complex<double>>(result));
100+
EXPECT_EQ(std::get<std::complex<double>>(result), std::complex<double>(5.0, 0.0));
101+
}
102+
103+
// TEST_F(CuDmOpConversion, AddScalarAndOperator) {
104+
// scalar_operator scalar_op(1.0);
105+
// matrix_operator mat_op("X", {0});
106+
107+
// auto scalar_result = converter->evaluate(scalar_op);
108+
// auto op_result = converter->evaluate(mat_op);
109+
110+
// auto final_result = converter->add(scalar_result, op_result);
111+
112+
// ASSERT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(final_result));
113+
// }
114+
115+
TEST_F(CuDmOpConversion, TensorProductOfScalars) {
116+
auto result = converter->tensor(2.0, 3.0);
117+
EXPECT_TRUE(std::holds_alternative<double>(result));
118+
EXPECT_EQ(std::get<double>(result), 6.0);
119+
}
120+
121+
// TEST_F(CuDmOpConversion, TensorProductScalarAndOperator) {
122+
// cudensitymatOperatorTerm_t op_term;
123+
// HANDLE_CUDM_ERROR(cudensitymatCreateOperatorTerm(handle, dimensions.size(), space_mode_extents.data(), &op_term));
124+
125+
// auto result = converter->tensor(2.0, op_term);
126+
// EXPECT_TRUE(std::holds_alternative<cudensitymatOperatorTerm_t>(result));
127+
128+
// HANDLE_CUDM_ERROR(cudensitymatDestroyOperatorTerm(op_term));
129+
// }

‎unittests/dynamics/test_mocks.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,13 @@
1414
#include <cudensitymat.h>
1515
#include <iostream>
1616

17+
// Mock cudensitymatHandle_t
18+
inline cudensitymatHandle_t mock_handle() {
19+
cudensitymatHandle_t handle;
20+
HANDLE_CUDM_ERROR(cudensitymatCreate(&handle));
21+
return handle;
22+
}
23+
1724
// Mock Liouvillian operator creation
1825
inline cudensitymatOperator_t mock_liouvillian(cudensitymatHandle_t handle) {
1926
cudensitymatOperator_t liouvillian = nullptr;

0 commit comments

Comments
 (0)