Fix slice transfer for flows with alpha channel

lib/fabrics/ofi/CMakeLists.txt

@@ -62,7 +62,7 @@ target_sources(mxl-fabrics-objects
             src/internal/Region.cpp
             src/internal/RegisteredRegion.cpp
             src/internal/RemoteRegion.cpp
-            src/internal/GrainSlices.cpp
+            src/internal/SliceRange.cpp
             src/internal/Target.cpp
             src/internal/TargetInfo.cpp
     )

lib/fabrics/ofi/src/internal/DataLayout.cpp

@@ -4,16 +4,44 @@
 
 #include "DataLayout.hpp"
 #include <cassert>
+#include <algorithm>
 #include <array>
+#include <numeric>
 #include "mxl/flowinfo.h"
+#include "Exception.hpp"
 
 namespace mxl::lib::fabrics::ofi
 {
-    DataLayout DataLayout::fromDiscrete(std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN> const& sliceSizes) noexcept
+    DataLayout DataLayout::fromDiscrete(std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN> const& sliceSizes, std::uint16_t totalSlices) noexcept
     {
-        return DataLayout{DataLayout::Discrete{.sliceSizes = sliceSizes}};
+        return DataLayout{
+            DataLayout::Discrete{.sliceSizes = sliceSizes, .totalSlices = totalSlices}
+        };
     };
 
+    std::size_t DataLayout::Discrete::totalLength() const noexcept
+    {
+        return std::accumulate(sliceSizes.begin(), sliceSizes.end(), std::size_t{0});
+    }
+
+    std::size_t DataLayout::Discrete::activePlaneCount() const noexcept
+    {
+        return std::ranges::count_if(sliceSizes, [](std::uint32_t sliceSize) { return sliceSize > 0; });
+    }
+
+    std::uint32_t DataLayout::Discrete::planePayloadOffset(std::size_t planeIndex, std::uint32_t grainPayloadOffset) const
+    {
+        if (sliceSizes.size() <= planeIndex)
+        {
+            throw Exception::invalidState("Invalid plane index {} not in range {}-{}", planeIndex, 0, sliceSizes.size());
+        }
+
+        return std::accumulate(sliceSizes.begin(),
+            sliceSizes.begin() + planeIndex,
+            grainPayloadOffset,
+            [this](std::uint32_t lhs, std::uint32_t rhs) { return lhs + (rhs * static_cast<std::uint32_t>(totalSlices)); });
+    }
+
     DataLayout DataLayout::fromContinuous(std::size_t sampleSize, std::size_t channelCount, std::size_t bufferLength) noexcept
     {
         return DataLayout{

lib/fabrics/ofi/src/internal/DataLayout.hpp

@@ -22,7 +22,24 @@ namespace mxl::lib::fabrics::ofi
          */
         struct Discrete
         {
-            std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN> sliceSizes; /**< Number of slices per plane. \see MXL_MAX_PLANES_PER_GRAIN */
+            std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN>
+                sliceSizes;            /**< Size in bytes of a single slice for each plane. \see MXL_MAX_PLANES_PER_GRAIN */
+            std::uint16_t totalSlices; /**< Total number of slices (e.g. video lines) per grain. */
+
+            /** \brief Return the total length of all active planes together */
+            [[nodiscard]]
+            std::size_t totalLength() const noexcept;
+
+            /** \brief Return the number of active planes (consecutive non-zero sliceSizes entries). */
+            [[nodiscard]]
+            std::size_t activePlaneCount() const noexcept;
+
+            /** \brief Return the byte offset from the grain start to where a given plane's payload begins.
+             * \param planeIndex The zero-based plane index.
+             * \param grainPayloadOffset The byte offset of the first plane's payload from the grain start (i.e. the grain header size).
+             */
+            [[nodiscard]]
+            std::uint32_t planePayloadOffset(std::size_t planeIndex, std::uint32_t grainPayloadOffset) const;
         };
 
         /** \brief Continuous layout variant of DataLayout.
@@ -37,10 +54,12 @@ namespace mxl::lib::fabrics::ofi
     public:
         /** \brief Create a DataLayout representing video data.
          * \param sliceSizes The slice sizes of each planes in the video data layout. \see MXL_MAX_PLANES_PER_GRAIN
+         * \param totalSlices Total number of slices (e.g. video lines) per grain.
          * \return A DataLayout representing the specified video layout.
          */
         [[nodiscard]]
-        static DataLayout fromDiscrete(std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN> const& sliceSizes) noexcept; // NOLINT
+        static DataLayout fromDiscrete(std::array<std::uint32_t, MXL_MAX_PLANES_PER_GRAIN> const& sliceSizes,
+            std::uint16_t totalSlices) noexcept; // NOLINT
 
         /** \brief Create a DataLayout representing audio data.
          * \param sampleSize The size of each audio sample in bytes.

lib/fabrics/ofi/src/internal/Protocol.hpp

@@ -10,8 +10,8 @@
 #include <rdma/fabric.h>
 #include "DataLayout.hpp"
 #include "Endpoint.hpp"
-#include "GrainSlices.hpp"
 #include "Region.hpp"
+#include "SliceRange.hpp"
 #include "Target.hpp"
 #include "TargetInfo.hpp"
 

lib/fabrics/ofi/src/internal/ProtocolEgressRMA.cpp

@@ -28,16 +28,53 @@ namespace mxl::lib::fabrics::ofi
     void RMAGrainEgressProtocol::transferGrain(Endpoint const& ep, std::uint64_t localIndex, std::uint64_t remoteIndex, std::uint32_t payloadOffset,
         SliceRange const& sliceRange, ::fi_addr_t destAddr)
     {
-        auto const localSize = sliceRange.transferSize(payloadOffset, _layout.sliceSizes[0]);
-        auto const localOffset = sliceRange.transferOffset(payloadOffset, _layout.sliceSizes[0]);
-        auto const remoteSize = sliceRange.transferSize(payloadOffset, _layout.sliceSizes[0]);
-        auto const remoteOffset = sliceRange.transferOffset(payloadOffset, _layout.sliceSizes[0]);
-
-        auto const localRegion = _localRegions[localIndex % _localRegions.size()].sub(localOffset, localSize);
-        auto const remoteRegion = _remoteInfo.remoteRegions[remoteIndex % _remoteInfo.remoteRegions.size()].sub(remoteOffset, remoteSize);
+        auto const localGrain = _localRegions[localIndex % _localRegions.size()];
+        auto const remoteGrain = _remoteInfo.remoteRegions[remoteIndex % _remoteInfo.remoteRegions.size()];
         auto const remoteSlot = remoteIndex % _remoteInfo.remoteRegions.size();
 
-        _pending += ep.write(_token, localRegion, remoteRegion, destAddr, ImmDataGrain{remoteSlot, sliceRange.end()}.data());
+        // Fast path: a full grain is transferred, can do one write no matter how many planes
+        if ((sliceRange.start() == 0) && (sliceRange.end() == _layout.totalSlices))
+        {
+            _pending += ep.write(_token, localGrain, remoteGrain, destAddr, std::make_optional(ImmDataGrain{remoteSlot, _layout.totalSlices}.data()));
+            return;
+        }
+
+        auto const planeCount = _layout.activePlaneCount();
+        for (std::size_t plane = 0; plane < planeCount; ++plane)
+        {
+            auto const sliceSize = _layout.sliceSizes[plane];
+            auto offset = std::uint32_t{0};
+            auto size = std::uint32_t{0};
+
+            if (plane == 0)
+            {
+                // need to include header with slice 0-x
+                if (sliceRange.start() == 0)
+                {
+                    offset = sliceRange.transferOffset(0, sliceSize);
+                    size = sliceRange.transferSize(sliceSize) + payloadOffset;
+                }
+                else
+                {
+                    offset = sliceRange.transferOffset(payloadOffset, sliceSize);
+                    size = sliceRange.transferSize(sliceSize);
+                }
+            }
+            else
+            {
+                auto const planeBase = _layout.planePayloadOffset(plane, payloadOffset);
+                offset = sliceRange.transferOffset(planeBase, sliceSize);
+                size = sliceRange.transferSize(sliceSize);
+            }
+
+            auto const localRegion = localGrain.sub(offset, size);
+            auto const remoteRegion = remoteGrain.sub(offset, size);
+
+            auto const isLastPlane = (plane == planeCount - 1);
+            auto const immData = isLastPlane ? std::make_optional(ImmDataGrain{remoteSlot, sliceRange.end()}.data()) : std::nullopt;
+
+            _pending += ep.write(_token, localRegion, remoteRegion, destAddr, immData);
+        }
     }
 
     void RMAGrainEgressProtocol::transferSamples(Endpoint const&, std::uint64_t, std::size_t, ::fi_addr_t)

lib/fabrics/ofi/src/internal/ProtocolIngressRMA.cpp

@@ -60,7 +60,7 @@ namespace mxl::lib::fabrics::ofi
         auto immData = completionData->data();
         if (!immData)
         {
-            throw Exception::invalidState("Received a completion without immediate data.");
+            return {};
         }
 
         auto [slot, slice] = ImmDataGrain{static_cast<std::uint32_t>(*immData)}.unpack();

lib/fabrics/ofi/src/internal/RCInitiator.cpp

@@ -16,9 +16,9 @@
 #include "Exception.hpp"
 #include "FabricInfo.hpp"
 #include "FabricInfoHelpers.hpp"
-#include "GrainSlices.hpp"
 #include "Protocol.hpp"
 #include "Region.hpp"
+#include "SliceRange.hpp"
 #include "VariantUtils.hpp"
 
 namespace mxl::lib::fabrics::ofi

lib/fabrics/ofi/src/internal/RDMInitiator.hpp

@@ -11,9 +11,9 @@
 #include <rdma/fabric.h>
 #include "mxl/fabrics.h"
 #include "Endpoint.hpp"
-#include "GrainSlices.hpp"
 #include "Initiator.hpp"
 #include "Protocol.hpp"
+#include "SliceRange.hpp"
 #include "TargetInfo.hpp"
 
 namespace mxl::lib::fabrics::ofi

lib/fabrics/ofi/src/internal/Region.cpp

@@ -178,8 +178,9 @@ namespace mxl::lib::fabrics::ofi
                 regions.emplace_back(grainInfoBaseAddr, grainInfoSize + grainPayloadSize, nullptr, nullptr, Region::Location::host());
             }
 
+            auto const totalSlices = discreteFlow.grainAt(0)->header.info.totalSlices;
             return {std::move(regions),
-                DataLayout::fromDiscrete(std::to_array(discreteFlow.flowInfo()->config.discrete.sliceSizes)),
+                DataLayout::fromDiscrete(std::to_array(discreteFlow.flowInfo()->config.discrete.sliceSizes), totalSlices),
                 discreteFlow.flowInfo()->config.common.maxSyncBatchSizeHint};
         }
         else if (mxlIsContinuousDataFormat(static_cast<int>(flow.flowInfo()->config.common.format)))

lib/fabrics/ofi/src/internal/GrainSlices.cpp → lib/fabrics/ofi/src/internal/SliceRange.cpp

@@ -2,7 +2,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-#include "GrainSlices.hpp"
+#include "SliceRange.hpp"
 #include "Exception.hpp"
 
 namespace mxl::lib::fabrics::ofi
@@ -17,28 +17,14 @@ namespace mxl::lib::fabrics::ofi
         return SliceRange{start, end};
     }
 
-    std::uint32_t SliceRange::transferSize(std::uint32_t payloadOffset, std::uint32_t sliceSize) const noexcept
+    std::uint32_t SliceRange::transferSize(std::uint32_t sliceSize) const noexcept
     {
-        auto size = (_end - _start) * sliceSize;
-
-        if (_start == 0)
-        {
-            size += payloadOffset;
-        }
-
-        return size;
+        return ((_end - _start) * sliceSize);
     }
 
     std::uint32_t SliceRange::transferOffset(std::uint32_t payloadOffset, std::uint32_t sliceSize) const noexcept
     {
-        if (_start == 0)
-        {
-            return 0;
-        }
-        else
-        {
-            return payloadOffset + (_start * sliceSize);
-        }
+        return payloadOffset + (_start * sliceSize);
     }
 
     std::uint16_t SliceRange::start() const noexcept

lib/fabrics/ofi/src/internal/GrainSlices.hpp → lib/fabrics/ofi/src/internal/SliceRange.hpp

@@ -40,7 +40,7 @@ namespace mxl::lib::fabrics::ofi
          * \note When start is 0, the size includes the header, it adds the payload offset.
          */
         [[nodiscard]]
-        std::uint32_t transferSize(std::uint32_t payloadOffset, std::uint32_t sliceSize) const noexcept;
+        std::uint32_t transferSize(std::uint32_t sliceSize) const noexcept;
 
         /** \brief Get the offset within the payload for the start of the range.
          * \note When start is 0, the offset is 0, because we include the header in the transfer.

lib/tests/fabrics/CMakeLists.txt

@@ -26,6 +26,7 @@ target_sources(mxl-fabrics-tests
         PRIVATE
             test_basics.cpp
             test_interfaces.cpp
+            test_sliced.cpp
     )
 
 target_link_libraries(mxl-fabrics-tests

lib/tests/fabrics/ofi/CMakeLists.txt

@@ -31,6 +31,8 @@ target_sources(mxl-fabrics-ofi-tests
             test_Provider.cpp
             test_ProviderConfig.cpp
             test_Region.cpp
+            test_DataLayout.cpp
+            test_SliceRange.cpp
     )
 
 target_link_libraries(mxl-fabrics-ofi-tests

lib/tests/fabrics/ofi/Util.hpp

@@ -77,7 +77,7 @@ namespace mxl::lib::fabrics::ofi
 
     inline MxlRegions getEmptyVideoMxlRegions()
     {
-        return MxlRegions({}, DataLayout::fromDiscrete({8, 0, 0, 0}));
+        return MxlRegions({}, DataLayout::fromDiscrete({8, 0, 0, 0}, 1));
     }
 
     inline std::pair<MxlRegions, InnerRegions> getHostRegionGroups()
@@ -97,12 +97,12 @@ namespace mxl::lib::fabrics::ofi
             regions.emplace_back(*innerRegion.data(), innerRegion.size(), nullptr, nullptr);
         }
 
-        auto mxlRegions = MxlRegions(regions, DataLayout::fromDiscrete({8, 0, 0, 0}));
+        auto mxlRegions = MxlRegions(regions, DataLayout::fromDiscrete({8, 0, 0, 0}, 1));
         return {mxlRegions, innerRegions};
     }
 
     inline MxlRegions getMxlRegions(std::vector<std::vector<std::uint8_t>> const& innerRegions,
-        DataLayout dataLayout = DataLayout::fromDiscrete({8, 0, 0, 0}))
+        DataLayout dataLayout = DataLayout::fromDiscrete({8, 0, 0, 0}, 1))
     {
         auto regions = std::vector<Region>{};
         regions.reserve(innerRegions.size());