Add hipcub::AliasTemporaries and some macros

CHANGELOG.md

@@ -12,6 +12,8 @@ Full documentation for hipCUB is available at [https://rocm.docs.amd.com/project
 * Added `--emulation` option for `rtest.py`
   * Unit tests can be run with `[--emulation|-e|--test|-t]=<test_name>`
 * Added `DeviceSelect::FlaggedIf` and its inplace overload.
+* Added CUB macros missing from hipCUB
+* Added `hipcub::AliasTemporaries` function for CUB parity.
 
 ### Changed
 * Changed the subset of tests that are run for smoke tests such that the smoke test will complete with faster run-time and to never exceed 2GB of vram usage. Use `python rtest.py [--emulation|-e|--test|-t]=smoke` to run these tests.

hipcub/include/hipcub/backend/cub/hipcub.hpp

@@ -112,5 +112,7 @@
 // These functions must be wrapped so they return
 // hipError_t instead of cudaError_t
 #include "util_allocator.hpp"
+#include "util_device.hpp"
+#include "util_temporary_storage.hpp"
 
 #endif // HIPCUB_CUB_HIPCUB_HPP_

hipcub/include/hipcub/backend/cub/util_device.hpp

@@ -0,0 +1,39 @@
+/******************************************************************************
+ * Copyright (c) 2011, Duane Merrill.  All rights reserved.
+ * Copyright (c) 2011-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Modifications Copyright (c) 2024, Advanced Micro Devices, Inc.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#ifndef HIPCUB_CUB_UTIL_DEVICE_HPP_
+#define HIPCUB_CUB_UTIL_DEVICE_HPP_
+
+#include "../../config.hpp"
+
+#include "util_temporary_storage.hpp"
+
+#include "cub/util_device.cuh"
+
+#endif // HIPCUB_CUB_UTIL_DEVICE_HPP_

hipcub/include/hipcub/backend/cub/util_macro.hpp

@@ -0,0 +1,71 @@
+/******************************************************************************
+ * Copyright (c) 2011, Duane Merrill.  All rights reserved.
+ * Copyright (c) 2011-2018, NVIDIA CORPORATION.  All rights reserved.
+ * Modifications Copyright (c) 2024, Advanced Micro Devices, Inc.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#ifndef HIPCUB_CUB_MACRO_HPP_
+#define HIPCUB_CUB_MACRO_HPP_
+
+#include "../../config.hpp"
+
+#include "cub/util_macro.cuh"
+
+BEGIN_HIPCUB_NAMESPACE
+
+#ifndef HIPCUB_MAX
+    /// Select maximum(a, b)
+    #define HIPCUB_MAX(a, b) CUB_MAX(a, b)
+#endif
+
+#ifndef HIPCUB_MIN
+    /// Select minimum(a, b)
+    #define HIPCUB_MIN(a, b) CUB_MIN(a, b)
+#endif
+
+#ifndef HIPCUB_QUOTIENT_FLOOR
+    /// Quotient of x/y rounded down to nearest integer
+    #define HIPCUB_QUOTIENT_FLOOR(x, y) CUB_QUOTIENT_FLOOR(x, y)
+#endif
+
+#ifndef HIPCUB_QUOTIENT_CEILING
+    /// Quotient of x/y rounded up to nearest integer
+    #define HIPCUB_QUOTIENT_CEILING(x, y) CUB_QUOTIENT_CEILING(x, y)
+#endif
+
+#ifndef HIPCUB_ROUND_UP_NEAREST
+    /// x rounded up to the nearest multiple of y
+    #define HIPCUB_ROUND_UP_NEAREST(x, y) CUB_ROUND_UP_NEAREST(x, y)
+#endif
+
+#ifndef HIPCUB_ROUND_DOWN_NEAREST
+    /// x rounded down to the nearest multiple of y
+    #define HIPCUB_ROUND_DOWN_NEAREST(x, y) CUB_ROUND_DOWN_NEAREST(x, y)
+#endif
+
+END_HIPCUB_NAMESPACE
+
+#endif // HIPCUB_CUB_MACRO_HPP_

hipcub/include/hipcub/backend/cub/util_temporary_storage.hpp

@@ -0,0 +1,71 @@
+/******************************************************************************
+ * Copyright (c) 2011, Duane Merrill.  All rights reserved.
+ * Copyright (c) 2011-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Modifications Copyright (c) 2024, Advanced Micro Devices, Inc.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#ifndef HIPCUB_CUB_UTIL_TEMPORARY_STORAGE_HPP_
+#define HIPCUB_CUB_UTIL_TEMPORARY_STORAGE_HPP_
+
+#include "../../config.hpp"
+
+#include <cub/util_temporary_storage.cuh>
+
+BEGIN_HIPCUB_NAMESPACE
+
+/// \brief Alias temporaries to externally-allocated device storage (or simply return the amount of storage needed).
+/// \tparam ALLOCATIONS - The number of allocations that are needed.
+/// \param 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.
+/// \param temp_storage_bytes - [in,out] Size in bytes of \t d_temp_storage allocation.
+/// \param allocations - [out] Pointers to device allocations needed.
+/// \param allocation_sizes - [in] Sizes in bytes of device allocations needed.
+template<int ALLOCATIONS>
+HIPCUB_HOST_DEVICE
+HIPCUB_FORCEINLINE hipError_t AliasTemporaries(void*   d_temp_storage,
+                                               size_t& temp_storage_bytes,
+                                               void* (&allocations)[ALLOCATIONS],
+                                               size_t (&allocation_sizes)[ALLOCATIONS])
+{
+    cudaError_t error = ::cub::AliasTemporaries(d_temp_storage,
+                                                temp_storage_bytes,
+                                                allocations,
+                                                allocation_sizes);
+
+    if(cudaSuccess == error)
+    {
+        return hipSuccess;
+    }
+    else if(cudaErrorInvalidValue == error)
+    {
+        return hipErrorInvalidValue;
+    }
+
+    return hipErrorUnknown;
+}
+
+END_HIPCUB_NAMESPACE
+
+#endif // HIPCUB_CUB_UTIL_TEMPORARY_STORAGE_HPP_

hipcub/include/hipcub/backend/rocprim/hipcub.hpp

@@ -102,7 +102,10 @@
 // Util
 #include "tuple.hpp"
 #include "util_allocator.hpp"
+#include "util_device.hpp"
+#include "util_macro.hpp"
 #include "util_ptx.hpp"
+#include "util_temporary_storage.hpp"
 #include "util_type.hpp"
 
 #endif // HIPCUB_ROCPRIM_HIPCUB_HPP_

hipcub/include/hipcub/backend/rocprim/util_device.hpp

@@ -0,0 +1,37 @@
+/******************************************************************************
+ * Copyright (c) 2011, Duane Merrill.  All rights reserved.
+ * Copyright (c) 2011-2024, NVIDIA CORPORATION.  All rights reserved.
+ * Modifications Copyright (c) 2024, Advanced Micro Devices, Inc.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#ifndef HIPCUB_ROCPRIM_UTIL_DEVICE_HPP_
+#define HIPCUB_ROCPRIM_UTIL_DEVICE_HPP_
+
+#include "../../config.hpp"
+
+#include "util_temporary_storage.hpp"
+
+#endif // HIPCUB_ROCPRIM_UTIL_DEVICE_HPP_

hipcub/include/hipcub/backend/rocprim/util_macro.hpp

@@ -0,0 +1,98 @@
+/******************************************************************************
+ * Copyright (c) 2011, Duane Merrill.  All rights reserved.
+ * Copyright (c) 2011-2018, NVIDIA CORPORATION.  All rights reserved.
+ * Modifications Copyright (c) 2024, Advanced Micro Devices, Inc.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#ifndef HIPCUB_ROCPRIM_MACRO_HPP_
+#define HIPCUB_ROCPRIM_MACRO_HPP_
+
+#include "../../config.hpp"
+
+#include <utility>
+
+BEGIN_HIPCUB_NAMESPACE
+
+/**
+ * \addtogroup UtilModule
+ * @{
+ */
+
+#ifndef DOXYGEN_SHOULD_SKIP_THIS
+    #define HIPCUB_PREVENT_MACRO_SUBSTITUTION
+template<typename T, typename U>
+constexpr __host__ __device__
+auto min HIPCUB_PREVENT_MACRO_SUBSTITUTION(T&& t, U&& u)
+    -> decltype(t < u ? std::forward<T>(t) : std::forward<U>(u))
+{
+    return t < u ? std::forward<T>(t) : std::forward<U>(u);
+}
+
+template<typename T, typename U>
+constexpr __host__ __device__
+auto max HIPCUB_PREVENT_MACRO_SUBSTITUTION(T&& t, U&& u)
+    -> decltype(t < u ? std::forward<U>(u) : std::forward<T>(t))
+{
+    return t < u ? std::forward<U>(u) : std::forward<T>(t);
+}
+    #undef HIPCUB_PREVENT_MACRO_SUBSTITUTION
+#endif
+
+#ifndef HIPCUB_MAX
+    /// Select maximum(a, b)
+    #define HIPCUB_MAX(a, b) (((b) > (a)) ? (b) : (a))
+#endif
+
+#ifndef HIPCUB_MIN
+    /// Select minimum(a, b)
+    #define HIPCUB_MIN(a, b) (((b) < (a)) ? (b) : (a))
+#endif
+
+#ifndef HIPCUB_QUOTIENT_FLOOR
+    /// Quotient of x/y rounded down to nearest integer
+    #define HIPCUB_QUOTIENT_FLOOR(x, y) ((x) / (y))
+#endif
+
+#ifndef HIPCUB_QUOTIENT_CEILING
+    /// Quotient of x/y rounded up to nearest integer
+    #define HIPCUB_QUOTIENT_CEILING(x, y) (((x) + (y)-1) / (y))
+#endif
+
+#ifndef HIPCUB_ROUND_UP_NEAREST
+    /// x rounded up to the nearest multiple of y
+    #define HIPCUB_ROUND_UP_NEAREST(x, y) (HIPCUB_QUOTIENT_CEILING(x, y) * y)
+#endif
+
+#ifndef HIPCUB_ROUND_DOWN_NEAREST
+    /// x rounded down to the nearest multiple of y
+    #define HIPCUB_ROUND_DOWN_NEAREST(x, y) (((x) / (y)) * y)
+#endif
+
+/** @} */ // end group UtilModule
+
+END_HIPCUB_NAMESPACE
+
+#endif // HIPCUB_ROCPRIM_MACRO_HPP_