Add huge page support for symmetric heap to improve CXI ATU performance

src/shmem_comm.h

@@ -425,13 +425,18 @@ static inline
 void shmem_internal_copy_self(void *dest, const void *source, size_t nelems)
 {
 #ifdef USE_FI_HMEM
-    // "completion" set to 1 to wait for completion of put operation initiated
-    // by shmem_internal_put_nb, even if "completion" not incremented in call 
-    // to shmem_internal_put_nb.
-    long completion = 1;
+    /* put_nb routes through inject, bounce-buffer, or put_large depending on
+     * size.  The inject path has no counter event, so put_wait (watermark-
+     * based) is not sufficient — it would return immediately with completion=0
+     * and leave the GPU write unordered.  put_quiet drains all pending puts
+     * and provides the NIC-level ordering fence needed to guarantee dest is
+     * visible at the target GPU before returning.
+     * bounce-buffer and put_large also set *completion, but put_quiet subsumes
+     * that wait, so no separate put_wait call is needed. */
+    long completion = 0;
     shmem_internal_put_nb(SHMEM_CTX_DEFAULT, dest, source, nelems,
                           shmem_internal_my_pe, &completion);
-    shmem_internal_put_wait(SHMEM_CTX_DEFAULT, &completion);
+    shmem_transport_put_quiet((shmem_transport_ctx_t *)SHMEM_CTX_DEFAULT);
 #else
     memcpy(dest, source, nelems);
 #endif

src/symmetric_heap_c.c

@@ -162,15 +162,16 @@ shmem_internal_get_next(intptr_t incr)
 
 /* alloc VM space starting @ '_end' + 1GB */
 #define ONEGIG (1024UL*1024UL*1024UL)
-static void *mmap_alloc(size_t bytes)
+static void *mmap_alloc(size_t bytes, size_t *mapped_bytes)
 {
     char *file_name = NULL;
     int fd = 0;
     char *directory = NULL;
-    void *requested_base =
-        (void*) (((unsigned long) shmem_internal_data_base +
-                  shmem_internal_data_length + 2 * ONEGIG) & ~(ONEGIG - 1));
+    void *requested_base = (void*) (((unsigned long) shmem_internal_data_base + shmem_internal_data_length + 2 * ONEGIG) & ~(ONEGIG - 1));
     void *ret;
+    size_t hugetlbfs_bytes = 0;  /* Rounded size for hugetlbfs, 0 if not used */
+
+    *mapped_bytes = bytes;  /* default: actual mapped size equals requested size */
 
 #ifdef __linux__
     /* huge page support only on Linux for now, default is to use 2MB large pages */
@@ -191,42 +192,109 @@ static void *mmap_alloc(size_t bytes)
                     sprintf(file_name, "%s/%s.%d", directory, basename, getpid());
                     fd = open(file_name, O_CREAT | O_RDWR, 0755);
                     if (fd < 0) {
-                        RAISE_WARN_STR("file open failed, cannot use huge pages");
+                        DEBUG_MSG("file open failed, will fall back to anonymous MAP_HUGETLB");
+                        free(directory);
+                        free(file_name);
+                        directory = NULL;
+                        file_name = NULL;
                         fd = 0;
                     } else {
-                        /* have to round up by the pagesize being used */
-                        bytes = CEILING(bytes, shmem_internal_params.SYMMETRIC_HEAP_PAGE_SIZE);
+                        /* Round up by the pagesize for hugetlbfs file */
+                        hugetlbfs_bytes = CEILING(bytes, shmem_internal_params.SYMMETRIC_HEAP_PAGE_SIZE);
                     }
                 }
             }
         }
     }
-#endif /* __linux__ */
 
+    DEBUG_MSG("mmap_alloc: bytes=%zu, hugetlbfs_bytes=%zu, fd=%d",
+              bytes, hugetlbfs_bytes, fd);
+
+    if (fd) {
+        /* Map the hugetlbfs file directly; MAP_ANON must not be used here
+         * because MAP_ANONYMOUS causes the kernel to ignore the fd, which
+         * would silently fall back to regular pages. */
+        if (ftruncate(fd, hugetlbfs_bytes) == -1) {
+            DEBUG_MSG("ftruncate on hugetlbfs file failed (%s), falling back to THP", strerror(errno));
+            unlink(file_name);
+            close(fd);
+            free(directory);
+            free(file_name);
+            directory = NULL;
+            file_name = NULL;
+            fd = 0;
+            /* Prefer requested_base to preserve virtual address symmetry (required for RVA);
+             * only use NULL as a last resort if requested_base is unavailable. */
+            ret = mmap(requested_base, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+            if (ret == MAP_FAILED)
+                ret = mmap(NULL, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+            if (ret != MAP_FAILED) {
+                if (madvise(ret, bytes, MADV_HUGEPAGE) != 0) {
+                    DEBUG_MSG("madvise(MADV_HUGEPAGE) failed (%s), using regular pages", strerror(errno));
+                }
+            }
+        } else {
+            ret = mmap(requested_base, hugetlbfs_bytes, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_HUGETLB, fd, 0);
+            if (ret != MAP_FAILED) {
+                DEBUG_MSG("Allocated symmetric heap via hugetlbfs file: %zu bytes", hugetlbfs_bytes);
+                *mapped_bytes = hugetlbfs_bytes;
+            }
+            unlink(file_name);
+            close(fd);
+            free(directory);
+            free(file_name);
+            directory = NULL;
+            file_name = NULL;
+            fd = 0;
+        }
+    } else if (shmem_internal_params.SYMMETRIC_HEAP_USE_HUGE_PAGES) {
+        /* Try anonymous MAP_HUGETLB first (works with nr_overcommit_hugepages).
+         * Explicitly request 2MB pages via MAP_HUGE_SHIFT (21 << MAP_HUGE_SHIFT = 2^21 = 2MB). */
+        ret = mmap(requested_base, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE | MAP_HUGETLB | (21 << MAP_HUGE_SHIFT), -1, 0);
+        if (ret == MAP_FAILED) {
+            DEBUG_MSG("mmap(MAP_HUGETLB) failed (%s), falling back to THP via madvise", strerror(errno));
+            /* Prefer requested_base to preserve virtual address symmetry (required for RVA);
+             * only use NULL as a last resort. MAP_HUGETLB failure means huge pages are
+             * unavailable, not that the virtual address range is blocked. */
+            ret = mmap(requested_base, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+            if (ret == MAP_FAILED)
+                ret = mmap(NULL, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+            if (ret != MAP_FAILED) {
+                if (madvise(ret, bytes, MADV_HUGEPAGE) != 0) {
+                    RAISE_WARN_MSG("madvise(MADV_HUGEPAGE) failed (%s), using regular pages\n", strerror(errno));
+                }
+            }
+        } else {
+            DEBUG_MSG("Allocated symmetric heap via anonymous MAP_HUGETLB (2MB pages): %zu bytes", bytes);
+        }
+    } else {
+        ret = mmap(requested_base, bytes, PROT_READ | PROT_WRITE, MAP_ANON | MAP_PRIVATE, -1, 0);
+    }
+#else
     ret = mmap(requested_base,
                bytes,
                PROT_READ | PROT_WRITE,
                MAP_ANON | MAP_PRIVATE,
                fd,
                0);
+#endif /* __linux__ */
     if (ret == MAP_FAILED) {
         RAISE_WARN_MSG("Unable to allocate sym. heap, size %zuB: %s\n"
                        RAISE_PE_PREFIX
                        "Try reducing SHMEM_SYMMETRIC_SIZE or number of PEs per node\n",
                        bytes, strerror(errno), shmem_internal_my_pe);
         ret = NULL;
     }
-    if (fd) {
-        if (file_name)
-            unlink(file_name);
-        close(fd);
-    }
+    /* Cleanup any remaining allocations (will be NULL if already freed above) */
     if (directory) {
         free(directory);
     }
     if (file_name) {
         free(file_name);
     }
+    if (fd > 0) {
+        close(fd);
+    }
     return ret;
 }
 
@@ -239,9 +307,13 @@ shmem_internal_symmetric_init(void)
                                  SHMEM_INTERNAL_HEAP_OVERHEAD;
 
     if (!shmem_internal_params.SYMMETRIC_HEAP_USE_MALLOC) {
+        size_t mapped_length = shmem_internal_heap_length;
         shmem_internal_heap_base =
             shmem_internal_heap_curr =
-            mmap_alloc(shmem_internal_heap_length);
+            mmap_alloc(shmem_internal_heap_length, &mapped_length);
+        /* Use the actual mapped size for munmap and transport registration.
+         * On the hugetlbfs path this may be rounded up to a huge-page boundary. */
+        shmem_internal_heap_length = mapped_length;
     } else {
         shmem_internal_heap_base =
             shmem_internal_heap_curr =