TFQ simulate expectation CUDA op

Author: jaeyoo

WORKSPACE

@@ -24,11 +24,19 @@ cc_library(
         ],
 )
 
+# http_archive(
+#     name = "qsim",
+#     sha256 = "b9c1eba09a885a938b5e73dfc2e02f5231cf3b01d899415caa24769346a731d5",
+#     strip_prefix = "qsim-0.13.3",
+#     urls = ["https://github.com/quantumlib/qsim/archive/refs/tags/v0.13.3.zip"],
+# )
+
+# TODO: After merging this patch later into qsim mainstream, remove this and uncomment the above.
 http_archive(
     name = "qsim",
-    sha256 = "b9c1eba09a885a938b5e73dfc2e02f5231cf3b01d899415caa24769346a731d5",
-    strip_prefix = "qsim-0.13.3",
-    urls = ["https://github.com/quantumlib/qsim/archive/refs/tags/v0.13.3.zip"],
+    sha256 = "97b26e1a8fe13cfa465611f2618ade00f539480c6a7849f020fbee3582686bbf",
+    strip_prefix = "qsim-0.15.0-dev20230317",
+    urls = ["https://github.com/jaeyoo/qsim/archive/refs/tags/v0.15.0+dev20230317.tar.gz"],
 )
 
 http_archive(

configure.sh

@@ -62,11 +62,11 @@ while [[ "$TF_NEED_CUDA" == "" ]]; do
 done
 
 while [[ "$TF_CUDA_VERSION" == "" ]]; do
-  read -p "Are you building against TensorFlow 2.1(including RCs) or newer?[Y/n] " INPUT
+  read -p "Are you building against TensorFlow 2.11(including RCs) or newer?[Y/n] " INPUT
   case $INPUT in
-    [Yy]* ) echo "Build against TensorFlow 2.1 or newer."; TF_CUDA_VERSION=10.1;;
-    [Nn]* ) echo "Build against TensorFlow <2.1."; TF_CUDA_VERSION=10.0;;
-    "" ) echo "Build against TensorFlow 2.1 or newer."; TF_CUDA_VERSION=10.1;;
+    [Yy]* ) echo "Build against TensorFlow 2.11 or newer."; TF_CUDA_VERSION=11;;
+    [Nn]* ) echo "Build against TensorFlow <2.11."; TF_CUDA_VERSION=10.0;;
+    "" ) echo "Build against TensorFlow 2.11 or newer."; TF_CUDA_VERSION=11;;
     * ) echo "Invalid selection: " $INPUT;;
   esac
 done
@@ -94,10 +94,11 @@ fi
 TF_CFLAGS=( $(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_compile_flags()))') )
 TF_LFLAGS="$(python -c 'import tensorflow as tf; print(" ".join(tf.sysconfig.get_link_flags()))')"
 
+# This config refers to building CUDA op kernels with nvcc.
 write_to_bazelrc "build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true"
-if [[ "$PIP_MANYLINUX2010" == "0" ]]; then
-  write_to_bazelrc "build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain"
-fi
+write_to_bazelrc "build:cuda --@local_config_cuda//:enable_cuda"
+write_to_bazelrc "build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain"
+
 
 write_to_bazelrc "build --experimental_repo_remote_exec"
 write_to_bazelrc "build --spawn_strategy=standalone"
@@ -143,7 +144,7 @@ fi
 # TODO(yifeif): do not hardcode path
 if [[ "$TF_NEED_CUDA" == "1" ]]; then
   write_action_env_to_bazelrc "TF_CUDA_VERSION" ${TF_CUDA_VERSION}
-  write_action_env_to_bazelrc "TF_CUDNN_VERSION" "7"
+  write_action_env_to_bazelrc "TF_CUDNN_VERSION" "8"
   if is_windows; then
     write_action_env_to_bazelrc "CUDNN_INSTALL_PATH" "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${TF_CUDA_VERSION}"
     write_action_env_to_bazelrc "CUDA_TOOLKIT_PATH" "C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v${TF_CUDA_VERSION}"

tensorflow_quantum/core/ops/BUILD

@@ -1,4 +1,6 @@
 # load op_wrapper
+load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_gpu_kernel_library", "tf_gen_op_wrapper_py")
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda")
 
 package(default_visibility = ["//visibility:public"])
 
@@ -12,6 +14,23 @@ config_setting(
     constraint_values = ["@bazel_tools//platforms:windows"],
 )
 
+cc_library(
+    name = "cuda",
+    data = [
+        "@local_config_cuda//cuda:cudart",
+    ],
+    linkopts = select({
+        ":windows": [],
+        "//conditions:default": [
+            "-Wl,-rpath,../local_config_cuda/cuda/lib64",
+            "-Wl,-rpath,../local_config_cuda/cuda/extras/CUPTI/lib64",
+        ],
+    }),
+    deps = [
+        "@local_config_cuda//cuda:cudart",
+    ],
+)
+
 py_library(
     name = "ops",
     srcs = ["__init__.py"],
@@ -24,6 +43,7 @@ py_library(
         ":tfq_adj_grad_op_py",
         ":tfq_ps_util_ops_py",
         ":tfq_simulate_ops_py",
+        ":tfq_simulate_ops_cuda_py",
         ":tfq_unitary_op_py",
         ":tfq_utility_ops_py",
         # test addons
@@ -619,6 +639,103 @@ py_test(
     ],
 )
 
+py_library(
+    name = "tfq_simulate_ops_cuda_py",
+    srcs = ["tfq_simulate_ops_cuda.py"],
+    data = [
+        ":_tfq_simulate_ops_cuda.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"],
+    deps = [
+        ":tfq_simulate_ops_cuda_py",
+        ":tfq_simulate_ops_py",
+        "//tensorflow_quantum/python:util",
+    ],
+    srcs_version = "PY3",
+)
+
+cc_binary(
+    name = "_tfq_simulate_ops_cuda.so",
+    srcs = [
+        "tfq_simulate_expectation_op_cuda.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",
+        "@qsim//lib:qsim_cuda_lib",
+        "@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"]),
+    # alwayslink=1,
+)
+
 py_library(
     name = "load_module",
     srcs = ["load_module.py"],