Stub CUDA runtime for CPU-only release builds

Author: Trinkle23897

third_party/cuda/cuda.bzl

@@ -16,12 +16,61 @@
 
 _CUDA_DIR = "CUDA_DIR"
 
+_STUB_CUDA_RUNTIME_API_H = """
+#ifndef CUDA_RUNTIME_API_H_
+#define CUDA_RUNTIME_API_H_
+
+#include <stddef.h>
+#include <stdlib.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef void* cudaStream_t;
+typedef int cudaError_t;
+
+enum cudaMemcpyKind {
+    cudaMemcpyHostToHost = 0,
+    cudaMemcpyHostToDevice = 1,
+    cudaMemcpyDeviceToHost = 2,
+    cudaMemcpyDeviceToDevice = 3,
+};
+
+static inline cudaError_t cudaMemcpyAsync(
+    void* dst, const void* src, size_t count, enum cudaMemcpyKind kind,
+    cudaStream_t stream) {
+    (void)dst;
+    (void)src;
+    (void)count;
+    (void)kind;
+    (void)stream;
+    abort();
+    return 1;
+}
+
+static inline cudaError_t cudaStreamSynchronize(cudaStream_t stream) {
+    (void)stream;
+    abort();
+    return 1;
+}
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif  /* CUDA_RUNTIME_API_H_ */
+"""
+
 def _impl(rctx):
     cuda_dir = rctx.os.environ.get(_CUDA_DIR, default = "/usr/local/cuda")
-    rctx.symlink("{}/include".format(cuda_dir), "include")
-    rctx.symlink("{}/lib64".format(cuda_dir), "lib64")
+    cuda_include = rctx.path("{}/include/cuda_runtime_api.h".format(cuda_dir))
+    cudart_static = rctx.path("{}/lib64/libcudart_static.a".format(cuda_dir))
     rctx.file("WORKSPACE")
-    rctx.file("BUILD", content = """
+    if cuda_include.exists and cudart_static.exists:
+        rctx.symlink("{}/include".format(cuda_dir), "include")
+        rctx.symlink("{}/lib64".format(cuda_dir), "lib64")
+        rctx.file("BUILD", content = """
 package(default_visibility = ["//visibility:public"])
 
 cc_library(
@@ -33,6 +82,17 @@ cc_library(
     ]),
     strip_include_prefix = "include",
 )
+""")
+    else:
+        rctx.file("include/cuda_runtime_api.h", _STUB_CUDA_RUNTIME_API_H)
+        rctx.file("BUILD", content = """
+package(default_visibility = ["//visibility:public"])
+
+cc_library(
+    name = "cudart_static",
+    hdrs = ["include/cuda_runtime_api.h"],
+    strip_include_prefix = "include",
+)
 """)
 
 cuda_configure = repository_rule(