Skip to content

Commit f97c496

Browse files
jaeyoojaeyoo
committed
Add cuQuantum test (but failed at numeric diff test)
1 parent 1999f26 commit f97c496

File tree

5 files changed

+215
-15
lines changed

5 files changed

+215
-15
lines changed

‎WORKSPACE

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,9 @@ cc_library(
3434
# TODO: After merging this patch later into qsim mainstream, remove this and uncomment the above.
3535
http_archive(
3636
name = "qsim",
37-
sha256 = "97b26e1a8fe13cfa465611f2618ade00f539480c6a7849f020fbee3582686bbf",
38-
strip_prefix = "qsim-0.15.0-dev20230317",
39-
urls = ["https://github.com/jaeyoo/qsim/archive/refs/tags/v0.15.0+dev20230317.tar.gz"],
37+
sha256 = "",
38+
strip_prefix = "qsim-0.15.0-dev20230327_v3",
39+
urls = ["https://github.com/jaeyoo/qsim/archive/refs/tags/v0.15.0+dev20230327_v3.tar.gz"],
4040
)
4141

4242
http_archive(
@@ -81,3 +81,21 @@ bind(
8181
actual = "@six_archive//:six",
8282
)
8383

84+
new_local_repository(
85+
name = "cuquantum_libs",
86+
path = "/usr/local/google/home/jaeyoo/workspace/cuquantum-linux-x86_64-22.11.0.13-archive",
87+
build_file_content = """
88+
cc_library(
89+
name = "custatevec_headers",
90+
srcs = ["include/custatevec.h"],
91+
visibility = ["//visibility:public"],
92+
)
93+
94+
cc_library(
95+
name = "custatevec",
96+
srcs = ["lib/libcustatevec.so"],
97+
visibility = ["//visibility:public"],
98+
)
99+
""",
100+
)
101+

‎tensorflow_quantum/core/ops/BUILD

Lines changed: 98 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -652,11 +652,25 @@ py_library(
652652
],
653653
)
654654

655+
py_library(
656+
name = "tfq_simulate_ops_cuquantum_py",
657+
srcs = ["tfq_simulate_ops_cuquantum.py"],
658+
data = [
659+
":_tfq_simulate_ops_cuquantum.so",
660+
],
661+
srcs_version = "PY3",
662+
deps = [
663+
# tensorflow framework for wrappers
664+
":load_module",
665+
],
666+
)
667+
655668
py_test(
656-
name = "tfq_simulate_ops_cuda_test",
657-
srcs = ["tfq_simulate_ops_cuda_test.py"],
669+
name = "tfq_simulate_ops_gpu_test",
670+
srcs = ["tfq_simulate_ops_gpu_test.py"],
658671
deps = [
659672
":tfq_simulate_ops_cuda_py",
673+
":tfq_simulate_ops_cuquantum_py",
660674
":tfq_simulate_ops_py",
661675
"//tensorflow_quantum/python:util",
662676
],
@@ -726,13 +740,94 @@ cc_binary(
726740
"//tensorflow_quantum/core/proto:program_cc_proto",
727741
"//tensorflow_quantum/core/src:circuit_parser_qsim",
728742
"//tensorflow_quantum/core/src:util_qsim",
743+
"@eigen//:eigen3",
744+
# "@local_cuda//:cuda_headers"
745+
# tensorflow core framework
746+
# tensorflow core lib
747+
# tensorflow core protos
748+
] + if_cuda_is_configured([
749+
":cuda",
750+
"@local_config_cuda//cuda:cuda_headers",
729751
"@qsim//lib:qsim_cuda_lib",
752+
]),
753+
# alwayslink=1,
754+
)
755+
756+
cc_binary(
757+
name = "_tfq_simulate_ops_cuquantum.so",
758+
srcs = [
759+
"tfq_simulate_expectation_op_cuquantum.cu.cc",
760+
],
761+
linkshared = 1,
762+
features = select({
763+
":windows": ["windows_export_all_symbols"],
764+
"//conditions:default": [],
765+
}),
766+
copts = select({
767+
":windows": [
768+
"/D__CLANG_SUPPORT_DYN_ANNOTATION__",
769+
"/D_USE_MATH_DEFINES",
770+
"/DEIGEN_MPL2_ONLY",
771+
"/DEIGEN_MAX_ALIGN_BYTES=64",
772+
"/DEIGEN_HAS_TYPE_TRAITS=0",
773+
"/DTF_USE_SNAPPY",
774+
"/showIncludes",
775+
"/MD",
776+
"/O2",
777+
"/DNDEBUG",
778+
"/w",
779+
"-DWIN32_LEAN_AND_MEAN",
780+
"-DNOGDI",
781+
"/d2ReducedOptimizeHugeFunctions",
782+
"/arch:AVX",
783+
"/std:c++17",
784+
"-DTENSORFLOW_MONOLITHIC_BUILD",
785+
"/DPLATFORM_WINDOWS",
786+
"/DEIGEN_HAS_C99_MATH",
787+
"/DTENSORFLOW_USE_EIGEN_THREADPOOL",
788+
"/DEIGEN_AVOID_STL_ARRAY",
789+
"/Iexternal/gemmlowp",
790+
"/wd4018",
791+
"/wd4577",
792+
"/DNOGDI",
793+
"/UTF_COMPILE_LIBRARY",
794+
],
795+
"//conditions:default": [
796+
"-Iexternal/local_cuda/cuda/include",
797+
# "--cuda-gpu-arch=sm_86",
798+
# "-L/usr/local/cuda/lib64",
799+
# "-lcudart_static",
800+
# "-ldl",
801+
# "-lrt",
802+
"-pthread",
803+
"-std=c++17",
804+
"-D_GLIBCXX_USE_CXX11_ABI=1",
805+
"-O3",
806+
"-Iexternal/cuda_headers",
807+
"-DNV_CUDNN_DISABLE_EXCEPTION",
808+
# "-fpermissive",
809+
],
810+
}) + if_cuda_is_configured(["-DTENSORFLOW_USE_NVCC=1", "-DGOOGLE_CUDA=1", "-x cuda", "-nvcc_options=relaxed-constexpr", "-nvcc_options=ftz=true"]),
811+
deps = [
812+
# cirq cc proto
813+
"//tensorflow_quantum/core/ops:parse_context",
814+
"//tensorflow_quantum/core/ops:tfq_simulate_utils",
815+
"//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
816+
"//tensorflow_quantum/core/proto:program_cc_proto",
817+
"//tensorflow_quantum/core/src:circuit_parser_qsim",
818+
"//tensorflow_quantum/core/src:util_qsim",
730819
"@eigen//:eigen3",
731820
# "@local_cuda//:cuda_headers"
732821
# tensorflow core framework
733822
# tensorflow core lib
734823
# tensorflow core protos
735-
] + if_cuda_is_configured([":cuda", "@local_config_cuda//cuda:cuda_headers"]),
824+
] + if_cuda_is_configured([
825+
":cuda",
826+
"@cuquantum_libs//:custatevec",
827+
"@cuquantum_libs//:custatevec_headers",
828+
"@local_config_cuda//cuda:cuda_headers",
829+
"@qsim//lib:qsim_cuquantum_lib",
830+
]),
736831
# alwayslink=1,
737832
)
738833

‎tensorflow_quantum/core/ops/tfq_simulate_expectation_op_cuquantum.cu.cc

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@ limitations under the License.
1414
#include <vector>
1515

1616
#include <chrono>
17-
#include <custatevec.h>
1817

18+
#include "../cuquantum_libs/include/custatevec.h"
1919
#include "../qsim/lib/circuit.h"
2020
#include "../qsim/lib/gate_appl.h"
2121
#include "../qsim/lib/gates_cirq.h"
@@ -48,7 +48,7 @@ typedef qsim::Circuit<QsimGate> QsimCircuit;
4848
class TfqSimulateExpectationOpCuQuantum : public tensorflow::OpKernel {
4949
public:
5050
explicit TfqSimulateExpectationOpCuQuantum(tensorflow::OpKernelConstruction* context)
51-
: OpKernel(context) { }
51+
: OpKernel(context) {}
5252

5353
void Compute(tensorflow::OpKernelContext* context) override {
5454
// TODO (mbbrough): add more dimension checks for other inputs here.
@@ -147,7 +147,7 @@ class TfqSimulateExpectationOpCuQuantum : public tensorflow::OpKernel {
147147
// Launch the cuda kernel.
148148
// Begin simulation.
149149
int largest_nq = 1;
150-
Simulator sim = Simulator(custatevec_handle_);
150+
Simulator sim = Simulator(cublas_handle_, custatevec_handle_);
151151
StateSpace ss = StateSpace(cublas_handle_, custatevec_handle_);
152152
auto sv = ss.Create(largest_nq);
153153
ss.SetStateZero(sv);
@@ -207,7 +207,7 @@ class TfqSimulateExpectationOpCuQuantum : public tensorflow::OpKernel {
207207
int cur_op_index;
208208

209209
// Launch custatevec, begin simulation.
210-
auto sim = Simulator(custatevec_handle_);
210+
auto sim = Simulator(cublas_handle_, custatevec_handle_);
211211
auto ss = StateSpace(cublas_handle_, custatevec_handle_);
212212
auto sv = ss.Create(largest_nq);
213213
auto scratch = ss.Create(largest_nq);
@@ -259,10 +259,10 @@ class TfqSimulateExpectationOpCuQuantum : public tensorflow::OpKernel {
259259
};
260260

261261
REGISTER_KERNEL_BUILDER(
262-
Name("TfqSimulateExpectationOpCuQuantum").Device(tensorflow::DEVICE_CPU),
263-
TfqSimulateExpectationOpCuQuantumOp);
262+
Name("TfqSimulateExpectationCuquantum").Device(tensorflow::DEVICE_CPU),
263+
TfqSimulateExpectationOpCuQuantum);
264264

265-
REGISTER_OP("TfqSimulateExpectationOpCuQuantum")
265+
REGISTER_OP("TfqSimulateExpectationCuquantum")
266266
.Input("programs: string")
267267
.Input("symbol_names: string")
268268
.Input("symbol_values: float")

‎tensorflow_quantum/core/ops/tfq_simulate_ops_cuquantum.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
# Copyright 2023 The TensorFlow Quantum Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
# ==============================================================================
15+
"""Module to register cuQuantum simulation python op."""
16+
import os
17+
import tensorflow as tf
18+
from tensorflow_quantum.core.ops.load_module import load_module
19+
20+
SIM_OP_MODULE = load_module("_tfq_simulate_ops_cuquantum.so")
21+
22+
23+
def tfq_simulate_expectation(programs, symbol_names, symbol_values, pauli_sums):
24+
"""Calculates the expectation value of circuits wrt some operator(s).
25+
Args:
26+
programs: `tf.Tensor` of strings with shape [batch_size] containing
27+
the string representations of the circuits to be executed.
28+
symbol_names: `tf.Tensor` of strings with shape [n_params], which
29+
is used to specify the order in which the values in
30+
`symbol_values` should be placed inside of the circuits in
31+
`programs`.
32+
symbol_values: `tf.Tensor` of real numbers with shape
33+
[batch_size, n_params] specifying parameter values to resolve
34+
into the circuits specificed by programs, following the ordering
35+
dictated by `symbol_names`.
36+
pauli_sums: `tf.Tensor` of strings with shape [batch_size, n_ops]
37+
containing the string representation of the operators that will
38+
be used on all of the circuits in the expectation calculations.
39+
Returns:
40+
`tf.Tensor` with shape [batch_size, n_ops] that holds the
41+
expectation value for each circuit with each op applied to it
42+
(after resolving the corresponding parameters in).
43+
"""
44+
return SIM_OP_MODULE.tfq_simulate_expectation_cuquantum(
45+
programs, symbol_names, tf.cast(symbol_values, tf.float32), pauli_sums)

‎tensorflow_quantum/core/ops/tfq_simulate_ops_cuda_test.py renamed to ‎tensorflow_quantum/core/ops/tfq_simulate_ops_gpu_test.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414
# ==============================================================================
15-
"""Tests that specifically target tfq_simulate_ops_cuda."""
15+
"""Tests that specifically target tfq_simulate_ops_cu*."""
1616
import os
1717
import time
1818
import numpy as np
@@ -22,6 +22,7 @@
2222

2323
from tensorflow_quantum.core.ops import tfq_simulate_ops
2424
from tensorflow_quantum.core.ops import tfq_simulate_ops_cuda
25+
from tensorflow_quantum.core.ops import tfq_simulate_ops_cuquantum
2526
from tensorflow_quantum.python import util
2627

2728
def measure_average_runtime(fn, tag, num_samples=10):
@@ -36,7 +37,7 @@ def measure_average_runtime(fn, tag, num_samples=10):
3637
return avg_time, result
3738

3839

39-
class SimulateExpectationTest(tf.test.TestCase):
40+
class SimulateExpectationGpuTest(tf.test.TestCase):
4041
"""Tests tfq_simulate_expectation."""
4142

4243
def test_simulate_expectation_cpu_vs_cuda(self):
@@ -81,6 +82,47 @@ def test_simulate_expectation_cpu_vs_cuda(self):
8182
# CUDA op should be faster than CPU op.
8283
self.assertGreater(cpu_avg_time, cuda_avg_time)
8384

85+
def test_simulate_expectation_cpu_vs_cuquantum(self):
86+
"""Make sure that cpu & gpu(cuquantum) ops have the same results."""
87+
n_qubits = 20
88+
batch_size = 5
89+
symbol_names = ['alpha']
90+
qubits = cirq.GridQubit.rect(1, n_qubits)
91+
circuit_batch, resolver_batch = \
92+
util.random_symbol_circuit_resolver_batch(
93+
qubits, symbol_names, batch_size)
94+
95+
circuit_batch_tensor = util.convert_to_tensor(circuit_batch)
96+
97+
symbol_values_array = np.array(
98+
[[resolver[symbol]
99+
for symbol in symbol_names]
100+
for resolver in resolver_batch])
101+
102+
pauli_sums = util.random_pauli_sums(qubits, 3, batch_size)
103+
pauli_sums_tensor = util.convert_to_tensor([[x] for x in pauli_sums])
104+
105+
cpu_avg_time, res_cpu = measure_average_runtime(
106+
lambda: tfq_simulate_ops.tfq_simulate_expectation(
107+
circuit_batch_tensor,
108+
symbol_names, symbol_values_array.astype(np.float64),
109+
pauli_sums_tensor),
110+
"CPU"
111+
)
112+
113+
cuda_avg_time, res_cuda = measure_average_runtime(
114+
lambda: tfq_simulate_ops_cuquantum.tfq_simulate_expectation(
115+
circuit_batch_tensor,
116+
symbol_names, symbol_values_array.astype(np.float64),
117+
pauli_sums_tensor),
118+
"cuQuantum"
119+
)
120+
121+
# The result should be the similar within a tolerance.
122+
np.testing.assert_allclose(res_cpu, res_cuda, atol=1e-5)
123+
124+
# cuQuantum op should be faster than CPU op.
125+
self.assertGreater(cpu_avg_time, cuda_avg_time)
84126

85127
if __name__ == "__main__":
86128
tf.test.main()

0 commit comments

Comments
 (0)