From 18558ebfaf8d38ee239bf9eea8da4c15ae9eb8b1 Mon Sep 17 00:00:00 2001
From: Levon Ter-grigoryan <patrios@google.com>
Date: Wed, 18 Sep 2024 16:27:23 +0200
Subject: [PATCH] Allow default alignment less then 16.

---
 src/enqueue.cc | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/enqueue.cc b/src/enqueue.cc
index 4edb42dec..1257f641b 100644
--- a/src/enqueue.cc
+++ b/src/enqueue.cc
@@ -1286,6 +1286,7 @@ static ncclResult_t uploadWork(struct ncclComm* comm, struct ncclKernelPlan* pla
   size_t batchBytes = plan->nWorkBatches*sizeof(struct ncclDevWorkBatch);
   void* fifoBufHost;
   uint32_t fifoCursor, fifoMask;
+  const int alignment = 16;
 
   switch (plan->workStorageType) {
   case ncclDevWorkStorageTypeArgs:
@@ -1302,8 +1303,9 @@ static ncclResult_t uploadWork(struct ncclComm* comm, struct ncclKernelPlan* pla
     plan->kernelArgs->workBuf = comm->workFifoBufDev;
     break;
   case ncclDevWorkStorageTypePersistent:
-    static_assert(16 <= alignof(max_align_t), "We rely on 16-byte alignment.");
-    fifoBufHost = malloc(workBytes);
+    fifoBufHost =
+        aligned_alloc(16,
+                      ((workBytes + alignment - 1) / alignment) * alignment);
     fifoCursor = 0;
     fifoMask = ~0u;
     break;