From 8bdfee73970eac30c2c776956b5aab3fc87b604e Mon Sep 17 00:00:00 2001
From: David Olsen <dolsen@nvidia.com>
Date: Mon, 6 Apr 2020 16:08:40 -0700
Subject: [PATCH] Change DeviceCount() to use the new host/device code pattern

The function DeviceCount was still using "#if CUB_PTX_ARCH == 0" to separate
host code and device code.  Change it to use the new pattern of
"if (CUB_IS_HOST_CODE)".
---
 cub/util_device.cuh | 31 ++++++++++++++++++++-----------
 1 file changed, 20 insertions(+), 11 deletions(-)

diff --git a/cub/util_device.cuh b/cub/util_device.cuh
index 8ac8972ba5..4840019bc6 100644
--- a/cub/util_device.cuh
+++ b/cub/util_device.cuh
@@ -211,18 +211,27 @@ struct ValueCache
  */
 CUB_RUNTIME_FUNCTION __forceinline__ int DeviceCount()
 {
-#if __cplusplus >= 201103L && (CUB_PTX_ARCH == 0) // Host code and C++11.
-
-    // C++11 guarantees that initialization of static locals is thread safe.
-    static ValueCache<int, DeviceCountUncached> cache;
-
-    return cache.value;
-
-#else // Device code or host code before C++11.
-
-    return DeviceCountUncached();
+    int result = -1;
+    if (CUB_IS_HOST_CODE) {
+        #if CUB_INCLUDE_HOST_CODE
+            #if __cplusplus >= 201103L
+                // Host code and C++11.
+                // C++11 guarantees that initialization of static locals is thread safe.
+                static ValueCache<int, DeviceCountUncached> cache;
 
-#endif
+                result = cache.value;
+            #else
+                // Host code and C++98.
+                result = DeviceCountUncached();
+            #endif
+        #endif
+    } else {
+        #if CUB_INCLUDE_DEVICE_CODE
+            // Device code.
+            result = DeviceCountUncached();
+        #endif
+    }
+    return result;
 }
 
 #if __cplusplus >= 201103L // C++11 and later.