Use SoA implementation in Copy specialization directly

alpaka-group · Oct 2, 2023 · 9f46000 · 9f46000
1 parent 8edac59
commit 9f46000
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 25 deletions.
diff --git a/include/llama/Copy.hpp b/include/llama/Copy.hpp
@@ -177,34 +177,15 @@ namespace llama
             throw std::runtime_error{"Destination SoA mapping's total array elements must be evenly divisible by the "
                                      "source AoSoA Lane count."};
 
-        // the same as SoA::blobNrAndOffset but takes a flat array index
-        auto mapSoA = [&](std::size_t flatArrayIndex, auto rc, bool mb) LLAMA_LAMBDA_INLINE
-        {
-            const auto blob = mb * flatRecordCoord<RecordDim, decltype(rc)>;
-            const auto offset = !mb * offsetOf<RecordDim, decltype(rc)> * flatSize
-                + sizeof(GetType<RecordDim, decltype(rc)>) * flatArrayIndex;
-            return NrAndOffset{blob, offset};
-        };
-
         auto mapSrc = [&](std::size_t flatArrayIndex, auto rc) LLAMA_LAMBDA_INLINE
         {
-            if constexpr(srcIsAoSoA)
-                return &srcView.blobs()[0][0] + srcView.mapping().blobNrAndOffset(flatArrayIndex, rc).offset;
-            else
-            {
-                const auto [blob, off] = mapSoA(flatArrayIndex, rc, isSrcMB);
-                return &srcView.blobs()[blob][off];
-            }
+            const auto [blob, off] = srcView.mapping().blobNrAndOffset(flatArrayIndex, rc);
+            return &srcView.blobs()[blob][off];
         };
         auto mapDst = [&](std::size_t flatArrayIndex, auto rc) LLAMA_LAMBDA_INLINE
         {
-            if constexpr(dstIsAoSoA)
-                return &dstView.blobs()[0][0] + dstView.mapping().blobNrAndOffset(flatArrayIndex, rc).offset;
-            else
-            {
-                const auto [blob, off] = mapSoA(flatArrayIndex, rc, isDstMB);
-                return &dstView.blobs()[blob][off];
-            }
+            const auto [blob, off] = dstView.mapping().blobNrAndOffset(flatArrayIndex, rc);
+            return &dstView.blobs()[blob][off];
         };
 
         static constexpr auto l = []

diff --git a/include/llama/mapping/SoA.hpp b/include/llama/mapping/SoA.hpp
@@ -123,10 +123,19 @@ namespace llama::mapping
         template<std::size_t... RecordCoords>
         LLAMA_FN_HOST_ACC_INLINE constexpr auto blobNrAndOffset(
             typename Base::ArrayIndex ai,
+            RecordCoord<RecordCoords...> rc = {}) const -> NrAndOffset<size_type>
+        {
+            return blobNrAndOffset(LinearizeArrayIndexFunctor{}(ai, Base::extents()), rc);
+        }
+
+        // Exposed for aosoaCommonBlockCopy. Should be private ...
+        template<std::size_t... RecordCoords>
+        LLAMA_FN_HOST_ACC_INLINE constexpr auto blobNrAndOffset(
+            size_type flatArrayIndex,
             RecordCoord<RecordCoords...> = {}) const -> NrAndOffset<size_type>
         {
-            const size_type elementOffset = LinearizeArrayIndexFunctor{}(ai, Base::extents())
-                * static_cast<size_type>(sizeof(GetType<TRecordDim, RecordCoord<RecordCoords...>>));
+            const size_type elementOffset
+                = flatArrayIndex * static_cast<size_type>(sizeof(GetType<TRecordDim, RecordCoord<RecordCoords...>>));
             if constexpr(blobs == Blobs::OnePerField)
             {
                 constexpr auto blob = flatRecordCoord<TRecordDim, RecordCoord<RecordCoords...>>;