Merge pull request #186 from NVIDIA/bugfix/nvtarget_nvrtc_device_anno…

…tations Fix NVRTC processing of <nv/target>
NVIDIA · Jul 16, 2021 · de88841 · de88841
2 parents 8bde5f9 + 51debf7
commit de88841
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/.upstream-tests/test/cuda/test_platform.pass.cpp b/.upstream-tests/test/cuda/test_platform.pass.cpp
@@ -8,10 +8,13 @@
 //===----------------------------------------------------------------------===//
 
 #include <nv/target>
-#include <stdio.h>
+
+#if !defined(__CUDACC_RTC__)
 #include <assert.h>
+#include <stdio.h>
+#endif
 
-#if defined(__NVCC__)
+#if defined(__NVCC__) || defined(__CUDACC_RTC__)
 #  define TEST_NVCC
 #elif defined(__NVCOMPILER)
 #  define TEST_NVCXX

diff --git a/include/nv/target b/include/nv/target
@@ -17,7 +17,12 @@
 #  define _NV_COMPILER_NVCC
 #elif defined(__NVCOMPILER) && __cplusplus >= 201103L
 #  define _NV_COMPILER_NVCXX
+#endif
+
+#if defined(__CUDACC_RTC__)
+#  define _NV_FUNCTION_ANNOTATION __device__
 #else
+#  define _NV_FUNCTION_ANNOTATION
 #endif
 
 #if defined(_NV_COMPILER_NVCXX)
@@ -61,6 +66,7 @@ namespace nv {
       // Store a set of targets as a set of bits
       struct _NV_BITSET_ATTRIBUTE target_description {
         base_int_t targets;
+        _NV_FUNCTION_ANNOTATION
         constexpr target_description(base_int_t a) : targets(a) { }
       };
 
@@ -72,9 +78,11 @@ namespace nv {
         sm_70 = 70, sm_72 = 72, sm_75 = 75,
         sm_80 = 80, sm_86 = 86,
       };
+      _NV_FUNCTION_ANNOTATION
       constexpr base_int_t toint(sm_selector a) {
         return static_cast<base_int_t>(a);
       }
+      _NV_FUNCTION_ANNOTATION
       constexpr base_int_t bitexact(sm_selector a) {
         return toint(a) == 35 ? sm_35_bit :
                toint(a) == 37 ? sm_37_bit :
@@ -90,6 +98,7 @@ namespace nv {
                toint(a) == 80 ? sm_80_bit :
                toint(a) == 86 ? sm_86_bit : 0;
       }
+      _NV_FUNCTION_ANNOTATION
       constexpr base_int_t bitrounddown(sm_selector a) {
         return toint(a) >= 86 ? sm_86_bit :
                toint(a) >= 80 ? sm_80_bit :
@@ -108,26 +117,31 @@ namespace nv {
 
       // Public API for NVIDIA GPUs
 
+      _NV_FUNCTION_ANNOTATION
       constexpr target_description is_exactly(sm_selector a) {
         return target_description(bitexact(a));
       }
 
+      _NV_FUNCTION_ANNOTATION
       constexpr target_description provides(sm_selector a) {
         return target_description(~(bitrounddown(a) - 1) & all_devices);
       }
 
       // Boolean operations on target sets
 
+      _NV_FUNCTION_ANNOTATION
       constexpr target_description operator&&(target_description a,
                                               target_description b) {
         return target_description(a.targets & b.targets);
       }
 
+      _NV_FUNCTION_ANNOTATION
       constexpr target_description operator||(target_description a,
                                               target_description b) {
         return target_description(a.targets | b.targets);
       }
 
+      _NV_FUNCTION_ANNOTATION
       constexpr target_description operator!(target_description a) {
         return target_description(~a.targets & (all_devices | all_hosts));
       }