From 6f989d921b12f0941915760d8fb972cdba43aaae Mon Sep 17 00:00:00 2001
From: Bernhard Manfred Gruber <bernhardmgruber@gmail.com>
Date: Tue, 5 Sep 2023 20:26:32 +0200
Subject: [PATCH] Refactor mapIdx and mapIdxPitchBytes

---
 include/alpaka/idx/MapIdx.hpp | 281 ++++++++--------------------------
 1 file changed, 63 insertions(+), 218 deletions(-)

diff --git a/include/alpaka/idx/MapIdx.hpp b/include/alpaka/idx/MapIdx.hpp
index b808b73d4656..af0ccea41701 100644
--- a/include/alpaka/idx/MapIdx.hpp
+++ b/include/alpaka/idx/MapIdx.hpp
@@ -12,237 +12,82 @@
 
 namespace alpaka
 {
-    namespace detail
-    {
-        //! Maps a linear index to a N dimensional index.
-        template<std::size_t TidxDimOut, std::size_t TidxDimIn, typename TSfinae = void>
-        struct MapIdx;
-        //! Maps a N dimensional index to the same N dimensional index.
-        template<std::size_t TidxDim>
-        struct MapIdx<TidxDim, TidxDim>
-        {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param extent Spatial size to map the index to.
-            // \return A N dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdx(
-                Vec<DimInt<TidxDim>, TElem> const& idx,
-                [[maybe_unused]] Vec<DimInt<TidxDim>, TElem> const& extent) -> Vec<DimInt<TidxDim>, TElem>
-            {
-                return idx;
-            }
-        };
-        //! Maps a 1 dimensional index to a N dimensional index.
-        template<std::size_t TidxDimOut>
-        struct MapIdx<TidxDimOut, 1u, std::enable_if_t<(TidxDimOut > 1u)>>
-        {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param extent Spatial size to map the index to
-            // \return A N dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdx(
-                Vec<DimInt<1u>, TElem> const& idx,
-                Vec<DimInt<TidxDimOut>, TElem> const& extent) -> Vec<DimInt<TidxDimOut>, TElem>
-            {
-                auto idxNd = Vec<DimInt<TidxDimOut>, TElem>::all(0u);
-
-                constexpr std::size_t lastIdx(TidxDimOut - 1u);
-
-                // fast-dim
-                idxNd[lastIdx] = static_cast<TElem>(idx[0u] % extent[lastIdx]);
-
-                // in-between
-                TElem hyperPlanesBefore = extent[lastIdx];
-                for(std::size_t r(1u); r < lastIdx; ++r)
-                {
-                    std::size_t const d = lastIdx - r;
-                    idxNd[d] = static_cast<TElem>(idx[0u] / hyperPlanesBefore % extent[d]);
-                    hyperPlanesBefore *= extent[d];
-                }
-
-                // slow-dim
-                idxNd[0u] = static_cast<TElem>(idx[0u] / hyperPlanesBefore);
-
-                return idxNd;
-            }
-        };
-        //! Maps a 1 dimensional index to a N dimensional index.
-        template<std::size_t TidxDimIn>
-        struct MapIdx<1u, TidxDimIn, std::enable_if_t<(TidxDimIn > 1u)>>
-        {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param extent Spatial size to map the index to.
-            // \return A 1 dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdx(
-                Vec<DimInt<TidxDimIn>, TElem> const& idx,
-                Vec<DimInt<TidxDimIn>, TElem> const& extent) -> Vec<DimInt<1u>, TElem>
-            {
-                TElem idx1d(idx[0u]);
-                for(std::size_t d(1u); d < TidxDimIn; ++d)
-                {
-                    idx1d = static_cast<TElem>(idx1d * extent[d] + idx[d]);
-                }
-                return {idx1d};
-            }
-        };
-
-        template<std::size_t TidxDimOut>
-        struct MapIdx<TidxDimOut, 0u>
-        {
-            template<typename TElem, std::size_t TidxDimExtents>
-            ALPAKA_FN_HOST_ACC static auto mapIdx(
-                Vec<DimInt<0u>, TElem> const&,
-                Vec<DimInt<TidxDimExtents>, TElem> const&) -> Vec<DimInt<TidxDimOut>, TElem>
-            {
-                return Vec<DimInt<TidxDimOut>, TElem>::zeros();
-            }
-        };
-
-        template<std::size_t TidxDimIn>
-        struct MapIdx<0u, TidxDimIn>
-        {
-            template<typename TElem, std::size_t TidxDimExtents>
-            ALPAKA_FN_HOST_ACC static auto mapIdx(
-                Vec<DimInt<TidxDimIn>, TElem> const&,
-                Vec<DimInt<TidxDimExtents>, TElem> const&) -> Vec<DimInt<0u>, TElem>
-            {
-                return {};
-            }
-        };
-    } // namespace detail
-
-    //! Maps a N dimensional index to a N dimensional position.
+    //! Maps an N-dimensional index to an N-dimensional position. At least one dimension must always be 1 or zero.
     //!
-    //! \tparam TidxDimOut Dimension of the index vector to map to.
-    //! \tparam TidxDimIn Dimension of the index vector to map from.
-    //! \tparam TElem Type of the elements of the index vector to map from.
+    //! \tparam TDimOut Dimension of the index vector to map to.
+    //! \param in The index vector to map from.
+    //! \param extent The extents of the input or output space, whichever has more than 1 dimensions.
     ALPAKA_NO_HOST_ACC_WARNING template<
-        std::size_t TidxDimOut,
-        std::size_t TidxDimIn,
-        std::size_t TidxDimExtents,
+        std::size_t TDimOut,
+        std::size_t TDimIn,
+        std::size_t TDimExtents,
         typename TElem>
-    ALPAKA_FN_HOST_ACC auto mapIdx(
-        Vec<DimInt<TidxDimIn>, TElem> const& idx,
-        Vec<DimInt<TidxDimExtents>, TElem> const& extent) -> Vec<DimInt<TidxDimOut>, TElem>
+    ALPAKA_FN_HOST_ACC auto mapIdx(Vec<DimInt<TDimIn>, TElem> const& in, Vec<DimInt<TDimExtents>, TElem> const& extent)
+        -> Vec<DimInt<TDimOut>, TElem>
     {
-        return detail::MapIdx<TidxDimOut, TidxDimIn>::mapIdx(idx, extent);
-    }
-
-    namespace detail
-    {
-        //! Maps a linear index to a N dimensional index assuming a buffer without padding.
-        template<std::size_t TidxDimOut, std::size_t TidxDimIn, typename TSfinae = void>
-        struct MapIdxPitchBytes;
-
-        //! Maps a N dimensional index to the same N dimensional index assuming a buffer without padding.
-        template<std::size_t TidxDim>
-        struct MapIdxPitchBytes<TidxDim, TidxDim>
-        {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param pitch Spatial pitch (in elems) to map the index to
-            // \return N dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdxPitchBytes(
-                Vec<DimInt<TidxDim>, TElem> const& idx,
-                [[maybe_unused]] Vec<DimInt<TidxDim>, TElem> const& pitch) -> Vec<DimInt<TidxDim>, TElem>
-            {
-                return idx;
-            }
-        };
-        //! Maps a 1 dimensional index to a N dimensional index assuming a buffer without padding.
-        template<std::size_t TidxDimOut>
-        struct MapIdxPitchBytes<TidxDimOut, 1u, std::enable_if_t<(TidxDimOut > 1u)>>
+        if constexpr(TDimOut == 0 || TDimIn == 0)
+            return Vec<DimInt<TDimOut>, TElem>::zeros();
+        else if constexpr(TDimOut == TDimIn)
+            return in;
+        else if constexpr(TDimOut == 1)
         {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param pitch Spatial pitch (in elems) to map the index to
-            // \return N dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdxPitchBytes(
-                Vec<DimInt<1u>, TElem> const& idx,
-                Vec<DimInt<TidxDimOut>, TElem> const& pitch) -> Vec<DimInt<TidxDimOut>, TElem>
-            {
-                auto idxNd = Vec<DimInt<TidxDimOut>, TElem>::zeros();
-
-                constexpr std::size_t lastIdx = TidxDimOut - 1u;
-
-                TElem tmp = idx[0u];
-                for(std::size_t d(0u); d < lastIdx; ++d)
-                {
-                    idxNd[d] = static_cast<TElem>(tmp / pitch[d]);
-                    tmp %= pitch[d];
-                }
-                idxNd[lastIdx] = tmp;
-
-                return idxNd;
-            }
-        };
-        //! Maps a N dimensional index to a 1 dimensional index assuming a buffer without padding.
-        template<std::size_t TidxDimIn>
-        struct MapIdxPitchBytes<1u, TidxDimIn, std::enable_if_t<(TidxDimIn > 1u)>>
+            TElem out = in[0];
+            for(std::size_t d = 1; d < TDimIn; ++d)
+                out = static_cast<TElem>(out * extent[d] + in[d]);
+            return {out};
+        }
+        else if constexpr(TDimIn == 1)
         {
-            // \tparam TElem Type of the index values.
-            // \param idx Idx to be mapped.
-            // \param pitch Spatial pitch (in elems) to map the index to
-            // \return A 1 dimensional vector.
-            ALPAKA_NO_HOST_ACC_WARNING
-            template<typename TElem>
-            ALPAKA_FN_HOST_ACC static auto mapIdxPitchBytes(
-                Vec<DimInt<TidxDimIn>, TElem> const& idx,
-                Vec<DimInt<TidxDimIn>, TElem> const& pitch) -> Vec<DimInt<1u>, TElem>
+            auto flat = in.front();
+            Vec<DimInt<TDimOut>, TElem> out;
+            for(std::size_t d = TDimOut - 1u; d > 0; d--)
             {
-                using DimMinusOne = DimInt<TidxDimIn - 1>;
-                return {idx.back() + (subVecBegin<DimMinusOne>(pitch) * subVecBegin<DimMinusOne>(idx)).sum()};
+                out[d] = static_cast<TElem>(flat % extent[d]);
+                flat /= extent[d];
             }
-        };
+            out.front() = static_cast<TElem>(flat);
+            return out;
+        }
+        else
+            static_assert(!sizeof(TElem), "Not implemented");
+    }
 
-        template<std::size_t TidxDimOut>
-        struct MapIdxPitchBytes<TidxDimOut, 0u>
+    //! Maps an N dimensional index to a N dimensional position based on the pitches of a view without padding or a
+    //! byte view. At least one dimension must always be 1 or zero.
+    //!
+    //! \tparam TDimOut Dimension of the index vector to map to.
+    //! \param in The index vector to map from.
+    //! \param pitches The pitches of the input or output space, whichever has more than 1 dimensions.
+    ALPAKA_NO_HOST_ACC_WARNING
+    template<std::size_t TDimOut, std::size_t TDimIn, std::size_t TidxDimPitch, typename TElem>
+    ALPAKA_FN_HOST_ACC auto mapIdxPitchBytes(
+        Vec<DimInt<TDimIn>, TElem> const& in,
+        Vec<DimInt<TidxDimPitch>, TElem> const& pitches) -> Vec<DimInt<TDimOut>, TElem>
+    {
+        if constexpr(TDimOut == 0 || TDimIn == 0)
+            return Vec<DimInt<TDimOut>, TElem>::zeros();
+        else if constexpr(TDimOut == TDimIn)
+            return in;
+        else if constexpr(TDimOut == 1)
         {
-            template<typename TElem, std::size_t TidxDimExtents>
-            ALPAKA_FN_HOST_ACC static auto mapIdxPitchBytes(
-                Vec<DimInt<0u>, TElem> const&,
-                Vec<DimInt<TidxDimExtents>, TElem> const&) -> Vec<DimInt<TidxDimOut>, TElem>
-            {
-                return Vec<DimInt<TidxDimOut>, TElem>::zeros();
-            }
-        };
-
-        template<std::size_t TidxDimIn>
-        struct MapIdxPitchBytes<0u, TidxDimIn>
+            using DimMinusOne = DimInt<TDimIn - 1>;
+            return {in.back() + (subVecBegin<DimMinusOne>(pitches) * subVecBegin<DimMinusOne>(in)).sum()};
+        }
+        else if constexpr(TDimIn == 1)
         {
-            template<typename TElem, std::size_t TidxDimExtents>
-            ALPAKA_FN_HOST_ACC static auto mapIdxPitchBytes(
-                Vec<DimInt<TidxDimIn>, TElem> const&,
-                Vec<DimInt<TidxDimExtents>, TElem> const&) -> Vec<DimInt<0u>, TElem>
+            auto result = Vec<DimInt<TDimOut>, TElem>::zeros();
+
+            TElem out = in.front();
+            for(std::size_t d = 0; d < TDimOut - 1u; ++d)
             {
-                return {};
+                result[d] = static_cast<TElem>(out / pitches[d]);
+                out %= pitches[d];
             }
-        };
-    } // namespace detail
+            result.back() = out;
 
-    //! Maps an N dimensional index to a N dimensional position based on
-    //! pitch in a buffer without padding or a byte buffer.
-    //!
-    //! \tparam TidxDimOut Dimension of the index vector to map to.
-    //! \tparam TidxDimIn Dimension of the index vector to map from.
-    //! \tparam TElem Type of the elements of the index vector to map from.
-    ALPAKA_NO_HOST_ACC_WARNING
-    template<std::size_t TidxDimOut, std::size_t TidxDimIn, std::size_t TidxDimPitch, typename TElem>
-    ALPAKA_FN_HOST_ACC auto mapIdxPitchBytes(
-        Vec<DimInt<TidxDimIn>, TElem> const& idx,
-        Vec<DimInt<TidxDimPitch>, TElem> const& pitch) -> Vec<DimInt<TidxDimOut>, TElem>
-    {
-        return detail::MapIdxPitchBytes<TidxDimOut, TidxDimIn>::mapIdxPitchBytes(idx, pitch);
+            return result;
+        }
+        else
+            static_assert(!sizeof(TElem), "Not implemented");
     }
 } // namespace alpaka