support new path to CUDAGeneratorImpl.h (NVIDIA#1267)

Pedrexus · Jan 13, 2022 · b2fdf9c · b2fdf9c
1 parent b88c507
commit b2fdf9c
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 17 deletions.
diff --git a/apex/contrib/csrc/fmha/src/fmha.h b/apex/contrib/csrc/fmha/src/fmha.h
@@ -30,7 +30,12 @@
 #include <cuda.h>
 #include <vector>
 
+#ifdef OLD_GENERATOR_PATH
 #include <ATen/CUDAGeneratorImpl.h>
+#else
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 
 #include <fmha_utils.h>

diff --git a/apex/contrib/csrc/multihead_attn/dropout.cuh b/apex/contrib/csrc/multihead_attn/dropout.cuh
@@ -1,10 +1,10 @@
 #pragma once
 #include <ATen/ATen.h>
 
-#ifdef OLD_GENERATOR
-#include <ATen/CUDAGenerator.h>
-#else
+#ifdef OLD_GENERATOR_PATH
 #include <ATen/CUDAGeneratorImpl.h>
+#else
+#include <ATen/cuda/CUDAGeneratorImpl.h>
 #endif
 
 #include <ATen/cuda/CUDAContext.h>
@@ -181,15 +181,10 @@ void apex_fused_dropout_cuda(scalar_t const *inputs, scalar_t *outputs,
   std::pair<uint64_t, uint64_t> rng_engine_inputs;
   {
     // See Note [Acquire lock when using random generators]
-#ifdef OLD_GENERATOR
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
-#else
     std::lock_guard<std::mutex> lock(gen.mutex());
     rng_engine_inputs =
         at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(
             counter_offset);
-#endif
   }
 
   apex_fused_dropout_kernel<scalar_t, accscalar_t, IndexType>
@@ -222,15 +217,10 @@ void apex_dropout_add_cuda(scalar_t const *inputs, scalar_t const *add_inputs,
   std::pair<uint64_t, uint64_t> rng_engine_inputs;
   {
     // See Note [Acquire lock when using random generators]
-#ifdef OLD_GENERATOR
-    std::lock_guard<std::mutex> lock(gen->mutex_);
-    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
-#else
     std::lock_guard<std::mutex> lock(gen.mutex());
     rng_engine_inputs =
         at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(
             counter_offset);
-#endif
   }
 
   apex_dropout_add_kernel<scalar_t, accscalar_t, IndexType>

diff --git a/apex/contrib/csrc/multihead_attn/softmax.cuh b/apex/contrib/csrc/multihead_attn/softmax.cuh
@@ -1,9 +1,14 @@
 #pragma once
 #include "philox.cuh"
-#include <ATen/CUDAGeneratorImpl.h>
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <curand_kernel.h>
 
+#ifdef OLD_GENERATOR_PATH
+#include <ATen/CUDAGeneratorImpl.h>
+#else
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <assert.h>
 #include <cfloat>
 #include <cmath>

diff --git a/apex/contrib/csrc/transducer/transducer_joint_kernel.cu b/apex/contrib/csrc/transducer/transducer_joint_kernel.cu
@@ -4,7 +4,13 @@
 
 #include <torch/extension.h>
 #include <ATen/AccumulateType.h>
+
+#ifdef OLD_GENERATOR_PATH
 #include <ATen/CUDAGeneratorImpl.h>
+#else
+#include <ATen/cuda/CUDAGeneratorImpl.h>
+#endif
+
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/CUDAGraphsUtils.cuh>
 #include <c10/macros/Macros.h>

diff --git a/setup.py b/setup.py
@@ -375,11 +375,12 @@ def append_nvcc_threads(nvcc_extra_args):
         )
     )
 
-# Check, if ATen/CUDAGenerator.h is found, otherwise use the new ATen/CUDAGeneratorImpl.h, due to breaking change in https://github.com/pytorch/pytorch/pull/36026
+# Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
+# See https://github.com/pytorch/pytorch/pull/70650
 generator_flag = []
 torch_dir = torch.__path__[0]
-if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGenerator.h")):
-    generator_flag = ["-DOLD_GENERATOR"]
+if os.path.exists(os.path.join(torch_dir, "include", "ATen", "CUDAGeneratorImpl.h")):
+    generator_flag = ["-DOLD_GENERATOR_PATH"]
 
 if "--fast_layer_norm" in sys.argv:
     sys.argv.remove("--fast_layer_norm")