Made CUDA optional

FluxML · Dec 18, 2022 · d1fb8e6 · d1fb8e6
1 parent af53e2b
commit d1fb8e6
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/build/torch_api.cpp b/build/torch_api.cpp
@@ -2,8 +2,10 @@
 #include<torch/torch.h>
 #include<ATen/autocast_mode.h>
 #include<torch/script.h>
+#ifdef CUDA
 #include<c10/cuda/CUDACachingAllocator.h>
 #include<c10/cuda/CUDAStream.h>
+#endif
 #include<vector>
 // #include<caml/fail.h>
 // #include<julia.h>
@@ -74,13 +76,15 @@ int at_new_tensor(tensor *out__) {
 return 1;
 }
 
+#ifdef CUDA
 int at_empty_cache() {
   PROTECT(
     c10::cuda::CUDACachingAllocator::emptyCache();
     return 0;
   )
 return 1;
 }
+#endif
 
 int at_no_grad(int flag) {
   PROTECT(
@@ -91,6 +95,7 @@ int at_no_grad(int flag) {
 return 1;
 }
 
+#ifdef CUDA
 int at_sync() {
   PROTECT(
     at::cuda::CUDAStream stream = at::cuda::getCurrentCUDAStream();
@@ -100,6 +105,7 @@ int at_sync() {
   // torch::cuda::synchronize();
 return 1;
 }
+#endif
 
 int at_tensor_of_data(tensor *out__, void *vs, int64_t *dims, int ndims, int element_size_in_bytes, int type) {
   PROTECT(
@@ -754,6 +760,7 @@ int ats_free(scalar s) {
 return 1;
 }
 
+#ifdef CUDA
 int atc_cuda_device_count(int *i) {
   PROTECT(
     i[0] = torch::cuda::device_count();
@@ -788,6 +795,7 @@ int atc_set_benchmark_cudnn(int b) {
   )
 return 1;
 }
+#endif
 
 int atm_load(char *filename, module *out__) {
   PROTECT(

diff --git a/build/torch_api.h b/build/torch_api.h
@@ -33,9 +33,13 @@ int flush_error();
 
 int at_manual_seed(int64_t);
 int at_new_tensor(tensor *);
+#ifdef CUDA
 int at_empty_cache();
+#endif
 int at_no_grad(int flag);
+#ifdef CUDA
 int at_sync();
+#endif
 int at_from_blob(tensor *, void *data, int64_t *dims, int ndims, int64_t *strides, int nstrides, int dev);
 int at_tensor_of_data(tensor *, void *vs, int64_t *dims, int ndims, int element_size_in_bytes, int type);
 int at_copy_data(tensor tensor, void *vs, int64_t numel, int element_size_in_bytes);
@@ -124,10 +128,12 @@ int ats_int(scalar *, int64_t);
 int ats_float(scalar *, double);
 int ats_free(scalar);
 
+#ifdef CUDA
 int atc_cuda_device_count(int *);
 int atc_cuda_is_available(int *);
 int atc_cudnn_is_available(int *);
 int atc_set_benchmark_cudnn(int b);
+#endif
 
 int atm_load(char *, module *);
 int atm_forward(tensor *, module, tensor *tensors, int ntensors);