 {
 private:
     int const old_device;
-
+    bool const needs_reset;
 public:
     __host__ __forceinline__ SwitchDevice(int new_device)
-      : old_device(CurrentDevice())
-    {
-        CubDebug(cudaSetDevice(new_device));
+      : old_device(CurrentDevice()), needs_reset(old_device != new_device)
+    {
+        if (needs_reset)
+            CubDebug(cudaSetDevice(new_device));
     }
 
     __host__ __forceinline__ ~SwitchDevice()
     {
-        CubDebug(cudaSetDevice(old_device));
+        if (needs_reset)
+            CubDebug(cudaSetDevice(old_device));
     }
 };
+
+/**
+ * \brief Returns the number of CUDA devices available or -1 if an error
+ *        occurred.
+ */
+CUB_RUNTIME_FUNCTION __forceinline__ int DeviceCountUncached()
+{
+#if defined(CUB_RUNTIME_ENABLED) // Host code or device code with the CUDA runtime.
+
+    int count = -1;
+    if (CubDebug(cudaGetDeviceCount(&count)))
+        // CUDA makes no guarantees about the state of the output parameter if
+        // `cudaGetDeviceCount` fails; in practice, they don't, but out of
+        // paranoia we'll reset `count` to `-1`.
+        count = -1;
+    return count;
+
+#else // Device code without the CUDA runtime.
+
+    return -1;
+
+#endif
+}
+
+#if __cplusplus >= 201103L // C++11 and later.
+
+/**
+ * \brief Cache for an arbitrary value produced by a nullary function.
+ */
+template <typename T, T(*Function)()>
+struct ValueCache
+{
+    T const value;
+
+    /**
+     * \brief Call the nullary function to produce the value and construct the
+     *        cache.
+     */
+    __host__ __forceinline__ ValueCache() : value(Function()) {}
+};
+
+#endif
+
+/**
+ * \brief Returns the number of CUDA devices available.
+ *
+ * \note This function may cache the result internally.
+ *
+ * \note This function is thread safe.
+ */
+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();
+
+#endif
+}
+
+#if __cplusplus >= 201103L // C++11 and later.
 
 /**
  * \brief Per-device cache for a CUDA attribute value; the attribute is queried

__host__ __device__ __forceinline__
cudaError_t AliasTemporaries(
    void    *d_temp_storage,                    ///< [in] %Device-accessible allocation of temporary storage.  When NULL, the required allocation size is written to \p temp_storage_bytes and no work is done.
    size_t  &temp_storage_bytes,                ///< [in,out] Size in bytes of \t d_temp_storage allocation
    void*   (&allocations)[ALLOCATIONS],        ///< [in,out] Pointers to device allocations needed
    size_t  (&allocation_sizes)[ALLOCATIONS])   ///< [in] Sizes in bytes of device allocations needed
{

/**
 * \brief Retrieves the PTX version that will be used on \p device (major * 100 + minor * 10).
 */
__host__ __forceinline__ cudaError_t PtxVersionUncached(int& ptx_version, int device)
{
    SwitchDevice sd(device);
    return PtxVersionUncached(ptx_version);
}

#if __cplusplus >= 201103L // C++11 and later.
template <typename Tag>
__host__ __forceinline__ PerDeviceAttributeCache& GetPerDeviceAttributeCache()
{
    // C++11 guarantees that initialization of static locals is thread safe.
    static PerDeviceAttributeCache cache;
    return cache;
}

struct PtxVersionCacheTag {};
struct SmVersionCacheTag {};
#endif

/**
 * \brief Retrieves the PTX version that will be used on \p device (major * 100 + minor * 10).
 *
 * \note This function may cache the result internally.
 *
 * \note This function is thread safe.
 */
__host__ __forceinline__ cudaError_t PtxVersion(int& ptx_version, int device)
{
#if __cplusplus >= 201103L // C++11 and later.

    auto const payload = GetPerDeviceAttributeCache<PtxVersionCacheTag>()(
      // If this call fails, then we get the error code back in the payload,
      // which we check with `CubDebug` below.
      [=] (int& pv) { return PtxVersionUncached(pv, device); },
      device);

    if (!CubDebug(payload.error))
        ptx_version = payload.attribute;

    return payload.error;

#else // Pre C++11.

    return PtxVersionUncached(ptx_version, device);

#endif
}
