UCT/MM: Fix the FIFO room calculation for tail > head

Signed-off-by: Arun Chandran <[email protected]>

Author: arun-chandran-edarath

src/uct/sm/mm/base/mm_ep.c

@@ -24,21 +24,6 @@ typedef enum {
     UCT_MM_SEND_AM_SHORT_IOV
 } uct_mm_send_op_t;
 
-
-/* Check if the resources on the remote peer are available for sending to it.
- * i.e. check if the remote receive FIFO has room in it.
- * return 1 if can send.
- * return 0 if can't send.
- * Ignore the event arm bit after the subtraction to accommodate
- * a) A head ARMED with UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED bit
- * b) head wrapping around after 0x7fff ffff ffff ffff and
- *    tail going beyond 0x7fff ffff ffff ffff, in this case the subtraction
- *    will wrap around, this scenario is highly unlikely.
- */
-#define UCT_MM_EP_IS_ABLE_TO_SEND(_head, _tail, _fifo_size) \
-    ucs_likely((((_head) - (_tail)) & ~UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED) \
-                 < (uint64_t)(_fifo_size))
-
 static UCS_F_NOINLINE ucs_status_t
 uct_mm_ep_attach_remote_seg(uct_mm_ep_t *ep, uct_mm_seg_id_t seg_id,
                             size_t length, void **address_p)

src/uct/sm/mm/base/mm_ep.h

@@ -2,6 +2,7 @@
 * Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED.
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2019. ALL RIGHTS RESERVED.
 * Copyright (C) ARM Ltd. 2016.  ALL RIGHTS RESERVED.
+* Copyright (C) Advanced Micro Devices, Inc. 2025. ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
 */
 
@@ -17,6 +18,25 @@
 KHASH_INIT(uct_mm_remote_seg, uintptr_t, uct_mm_remote_seg_t, 1,
            kh_int64_hash_func, kh_int64_hash_equal)
 
+/*
+ * Check if the remote receive FIFO has room.
+ * Returns 1 if can send, 0 otherwise.
+ *
+ * Logic (ignore UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED on head, compare signed delta):
+ *   - Compute s = (int64_t)(((uint64_t)_head & ~UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED) -
+ *                           (uint64_t)_tail)
+ *   - Room available iff s < (int64_t)_fifo_size
+ *
+ * Practical note (head counter runtime): We assume the head counter
+ * increments once every 1 ns. On this timescale, the signed63 midpoint (2^62)
+ * is ~4.61e18 ticks (~146 years). Over 5 years, head would advance by
+ * ~1.5768e17 ticks (~3.4% of that midpoint), which is far from any wraparound
+ * edge case.
+ */
+#define UCT_MM_EP_IS_ABLE_TO_SEND(_head, _tail, _fifo_size) \
+    (((int64_t)(((_head) & ~UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED) - \
+                (_tail))) < (int64_t)(_fifo_size))
+
 
 /**
  * MM transport endpoint

test/gtest/Makefile.am

@@ -3,7 +3,7 @@
 # Copyright (C) UT-Battelle, LLC. 2015. ALL RIGHTS RESERVED.
 # Copyright (C) The University of Tennessee and the University of Tennessee Research Foundation. 2016. ALL RIGHTS RESERVED.
 # Copyright (C) Los Alamos National Security, LLC. 2018 ALL RIGHTS RESERVED.
-# Copyright (C) Advanced Micro Devices, Inc. 2019. ALL RIGHTS RESERVED.
+# Copyright (C) Advanced Micro Devices, Inc. 2025. ALL RIGHTS RESERVED.
 # Copyright (C) ARM Ltd. 2020.  ALL RIGHTS RESERVED.
 # Copyright (C) NextSilicon Ltd. 2021.  ALL RIGHTS RESERVED.
 #
@@ -118,6 +118,7 @@ gtest_SOURCES = \
 	uct/test_many2one_am.cc \
 	uct/test_md.cc \
 	uct/test_mm.cc \
+	uct/sm/mm/test_mm_fifo_room.cc \
 	uct/test_mem.cc \
 	uct/test_p2p_am.cc \
 	uct/test_p2p_err.cc \

test/gtest/uct/sm/mm/test_mm_fifo_room.cc

@@ -0,0 +1,88 @@
+/**
+* Copyright (C) Advanced Micro Devices, Inc. 2025. ALL RIGHTS RESERVED.
+* See file LICENSE for terms.
+*/
+#include <common/test.h>
+#include <cstdint>
+#include "uct/sm/mm/base/mm_iface.h"
+#include "uct/sm/mm/base/mm_ep.h"
+
+namespace {
+
+static constexpr uint64_t EA  = UCT_MM_IFACE_FIFO_HEAD_EVENT_ARMED;
+
+/* Practical upper bound for head/tail in real deployments: 2^62.
+ * Rationale: head advances roughly once per nanosecond in the worst case.
+ * The signed-63 midpoint is 2^62 (~4.61e18). At 1 tick/ns, reaching 2^62
+ * takes ~146 years. Therefore tests cap counters at <= 2^62 to reflect
+ * realistic long-running processes while still exercising wrap-related logic.
+ */
+static constexpr uint64_t LIM       = 1ull << 62;
+/* 2^31 constant to guard against regressions to 32-bit comparisons */
+static constexpr uint64_t POW2_31   = 1ull << 31;
+
+struct case_item {
+    const char *name;
+    uint64_t head;
+    uint64_t tail;
+    unsigned fifo;
+    bool expect;
+};
+
+static const case_item k_cases[] = {
+    /* head < tail */
+    {"lt:<fifo",        512,                     600,           256, true},
+    {"lt:==fifo",       512,                     768,           256, true},
+    {"lt:>fifo",        512,                     900,           256, true},
+    {"EA lt:<fifo",     EA | 512,                600,           256, true},
+    {"EA lt:==fifo",    EA | 512,                768,           256, true},
+    {"EA lt:>fifo",     EA | 512,                900,           256, true},
+
+    /* head > tail */
+    {"gt:<fifo",        100,                     0,             256, true},
+    {"gt:==fifo",       256,                     0,             256, false},
+    {"gt:>fifo",        300,                     0,             256, false},
+    {"EA gt:<fifo",     EA | 100,                0,             256, true},
+    {"EA gt:==fifo",    EA | 256,                0,             256, false},
+    {"EA gt:>fifo",     EA | 300,                0,             256, false},
+
+    /* Large deltas around 2^31 to catch regressions to 32-bit compare */
+    {"gt:d=2^31-1@t0",  POW2_31 - 1ull,          0,             256, false},
+    {"gt:d=2^31@t0",    POW2_31,                 0,             256, false},
+    {"gt:d=2^31+1@t0",  POW2_31 + 1ull,          0,             256, false},
+    {"EA gt:d=2^31@t0", EA | POW2_31,            0,             256, false},
+
+    /* head == tail */
+    {"eq:zero",         0,                       0,             256, true},
+    {"eq:EA",           EA,                      0,             256, true},
+
+    /* Around 2^62 boundaries (head < tail deltas) */
+    {"lt:2^62-1",       512,                     512 + LIM - 1, 256, true},
+    {"lt:2^62",         512,                     512 + LIM,     256, true},
+    {"lt:2^62+1",       512,                     512 + LIM + 1, 256, true},
+
+    /* Special tail at MSB (robustness) */
+    {"tailEA:+255",     0xff,                    EA,            256, true},
+    {"tailEA:+256",     0x100,                   EA,            256, true},
+
+    /* Practical cap at 2^62 */
+    {"cap:eq",          LIM,                     LIM,           256, true},
+    {"cap:eq EA",       EA | LIM,                LIM,           256, true}
+};
+
+}
+
+class test_mm_fifo_room : public ucs::test {
+protected:
+    void check_case(const case_item &c) {
+        bool got = UCT_MM_EP_IS_ABLE_TO_SEND(c.head, c.tail, c.fifo);
+        EXPECT_EQ(c.expect, got) << c.name;
+    }
+};
+
+UCS_TEST_F(test_mm_fifo_room, predicate_matrix) {
+    for (const auto &c : k_cases) {
+        check_case(c);
+    }
+
+}