Simplify forEachArrayIndex

alpaka-group · Aug 10, 2023 · 1ca0480 · 1ca0480
1 parent 32beb4e
commit 1ca0480
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 13 deletions.
diff --git a/include/llama/ArrayExtents.hpp b/include/llama/ArrayExtents.hpp
@@ -239,23 +239,20 @@ namespace llama
 
     template<typename SizeType, std::size_t Dim, typename Func, typename... OuterIndices>
     LLAMA_FN_HOST_ACC_INLINE void forEachArrayIndex(
-        [[maybe_unused]] ArrayIndex<SizeType, Dim> adSize,
+        [[maybe_unused]] const ArrayIndex<SizeType, Dim>& extents,
         Func&& func,
         OuterIndices... outerIndices)
     {
+        constexpr auto fixedIndices = sizeof...(outerIndices);
         LLAMA_BEGIN_SUPPRESS_HOST_DEVICE_WARNING
-        if constexpr(Dim > 0)
+        if constexpr(fixedIndices < Dim)
         {
-            for(SizeType i = 0; i < adSize[0]; i++)
-                forEachArrayIndex(
-                    ArrayIndex<SizeType, Dim - 1>{popFront(adSize)},
-                    std::forward<Func>(func),
-                    outerIndices...,
-                    i);
+            for(SizeType i = 0; i < extents[fixedIndices]; i++)
+                forEachArrayIndex(extents, std::forward<Func>(func), outerIndices..., i);
         }
         else
         {
-            std::forward<Func>(func)(ArrayIndex<SizeType, sizeof...(outerIndices)>{outerIndices...});
+            std::forward<Func>(func)(ArrayIndex<SizeType, fixedIndices>{outerIndices...});
         }
         LLAMA_END_SUPPRESS_HOST_DEVICE_WARNING
     }

diff --git a/include/llama/Copy.hpp b/include/llama/Copy.hpp
@@ -103,10 +103,7 @@ namespace llama
         {
             using SrcSizeType = typename SrcMapping::ArrayExtents::value_type;
             if constexpr(dims > 1)
-                forEachArrayIndex(
-                    ArrayIndex<SrcSizeType, dims - 1>{popFront(extents)},
-                    copyOne,
-                    static_cast<SrcSizeType>(i));
+                forEachArrayIndex(extents, copyOne, static_cast<SrcSizeType>(i));
             else
                 copyOne(ArrayIndex<SrcSizeType, dims>{static_cast<std::size_t>(i)});
         }