Merge pull request #2 from jaeyoo/cuquantum_support

Cuquantum support

Author: jaeyoo

WORKSPACE

@@ -34,9 +34,9 @@ cc_library(
 # TODO: After merging this patch later into qsim mainstream, remove this and uncomment the above.
 http_archive(
     name = "qsim",
-    sha256 = "97b26e1a8fe13cfa465611f2618ade00f539480c6a7849f020fbee3582686bbf",
-    strip_prefix = "qsim-0.15.0-dev20230317",
-    urls = ["https://github.com/jaeyoo/qsim/archive/refs/tags/v0.15.0+dev20230317.tar.gz"],
+    sha256 = "",
+    strip_prefix = "qsim-0.15.0-dev20230327_v3",
+    urls = ["https://github.com/jaeyoo/qsim/archive/refs/tags/v0.15.0+dev20230327_v3.tar.gz"],
 )
 
 http_archive(
@@ -81,3 +81,21 @@ bind(
     actual = "@six_archive//:six",
 )
 
+new_local_repository(
+    name = "cuquantum_libs",
+    path = "/usr/local/google/home/jaeyoo/workspace/cuquantum-linux-x86_64-22.11.0.13-archive",
+    build_file_content = """
+cc_library(
+    name = "custatevec_headers",
+    srcs = ["include/custatevec.h"],
+    visibility = ["//visibility:public"],
+)
+
+cc_library(
+    name = "custatevec",
+    srcs = ["lib/libcustatevec.so"],
+    visibility = ["//visibility:public"],
+)
+""",
+)
+

tensorflow_quantum/core/ops/BUILD

@@ -652,11 +652,25 @@ py_library(
     ],
 )
 
+py_library(
+    name = "tfq_simulate_ops_cuquantum_py",
+    srcs = ["tfq_simulate_ops_cuquantum.py"],
+    data = [
+        ":_tfq_simulate_ops_cuquantum.so",
+    ],
+    srcs_version = "PY3",
+    deps = [
+        # tensorflow framework for wrappers
+        ":load_module",
+    ],
+)
+
 py_test(
-    name = "tfq_simulate_ops_cuda_test",
-    srcs = ["tfq_simulate_ops_cuda_test.py"],
+    name = "tfq_simulate_ops_gpu_test",
+    srcs = ["tfq_simulate_ops_gpu_test.py"],
     deps = [
         ":tfq_simulate_ops_cuda_py",
+        ":tfq_simulate_ops_cuquantum_py",
         ":tfq_simulate_ops_py",
         "//tensorflow_quantum/python:util",
     ],
@@ -726,13 +740,94 @@ cc_binary(
         "//tensorflow_quantum/core/proto:program_cc_proto",
         "//tensorflow_quantum/core/src:circuit_parser_qsim",
         "//tensorflow_quantum/core/src:util_qsim",
+        "@eigen//:eigen3",
+        # "@local_cuda//:cuda_headers"
+        # tensorflow core framework
+        # tensorflow core lib
+        # tensorflow core protos
+    ] + if_cuda_is_configured([
+        ":cuda",
+        "@local_config_cuda//cuda:cuda_headers",
         "@qsim//lib:qsim_cuda_lib",
+    ]),
+    # alwayslink=1,
+)
+
+cc_binary(
+    name = "_tfq_simulate_ops_cuquantum.so",
+    srcs = [
+        "tfq_simulate_expectation_op_cuquantum.cu.cc",
+    ],
+    linkshared = 1,
+    features = select({
+        ":windows": ["windows_export_all_symbols"],
+        "//conditions:default": [],
+    }),
+    copts = select({
+        ":windows": [
+            "/D__CLANG_SUPPORT_DYN_ANNOTATION__",
+            "/D_USE_MATH_DEFINES",
+            "/DEIGEN_MPL2_ONLY",
+            "/DEIGEN_MAX_ALIGN_BYTES=64",
+            "/DEIGEN_HAS_TYPE_TRAITS=0",
+            "/DTF_USE_SNAPPY",
+            "/showIncludes",
+            "/MD",
+            "/O2",
+            "/DNDEBUG",
+            "/w",
+            "-DWIN32_LEAN_AND_MEAN",
+            "-DNOGDI",
+            "/d2ReducedOptimizeHugeFunctions",
+            "/arch:AVX",
+            "/std:c++17",
+            "-DTENSORFLOW_MONOLITHIC_BUILD",
+            "/DPLATFORM_WINDOWS",
+            "/DEIGEN_HAS_C99_MATH",
+            "/DTENSORFLOW_USE_EIGEN_THREADPOOL",
+            "/DEIGEN_AVOID_STL_ARRAY",
+            "/Iexternal/gemmlowp",
+            "/wd4018",
+            "/wd4577",
+            "/DNOGDI",
+            "/UTF_COMPILE_LIBRARY",
+        ],
+        "//conditions:default": [
+            "-Iexternal/local_cuda/cuda/include",
+            # "--cuda-gpu-arch=sm_86",
+            # "-L/usr/local/cuda/lib64",
+            # "-lcudart_static",
+            # "-ldl",
+            # "-lrt",
+            "-pthread",
+            "-std=c++17",
+            "-D_GLIBCXX_USE_CXX11_ABI=1",
+            "-O3",
+            "-Iexternal/cuda_headers",
+            "-DNV_CUDNN_DISABLE_EXCEPTION",
+            # "-fpermissive",
+        ],
+    }) + if_cuda_is_configured(["-DTENSORFLOW_USE_NVCC=1", "-DGOOGLE_CUDA=1", "-x cuda", "-nvcc_options=relaxed-constexpr", "-nvcc_options=ftz=true"]),
+    deps = [
+        # cirq cc proto
+        "//tensorflow_quantum/core/ops:parse_context",
+        "//tensorflow_quantum/core/ops:tfq_simulate_utils",
+        "//tensorflow_quantum/core/proto:pauli_sum_cc_proto",
+        "//tensorflow_quantum/core/proto:program_cc_proto",
+        "//tensorflow_quantum/core/src:circuit_parser_qsim",
+        "//tensorflow_quantum/core/src:util_qsim",
         "@eigen//:eigen3",
         # "@local_cuda//:cuda_headers"
         # tensorflow core framework
         # tensorflow core lib
         # tensorflow core protos
-    ] + if_cuda_is_configured([":cuda",  "@local_config_cuda//cuda:cuda_headers"]),
+    ] + if_cuda_is_configured([
+        ":cuda",
+        "@cuquantum_libs//:custatevec",
+        "@cuquantum_libs//:custatevec_headers",
+        "@local_config_cuda//cuda:cuda_headers",
+        "@qsim//lib:qsim_cuquantum_lib",
+    ]),
     # alwayslink=1,
 )