WIP

Author: Binyang2014

src/algorithms/algorithm_utils.cc renamed to src/algorithms/utils.cc

@@ -1,10 +1,13 @@
+#include "algorithms/utils.hpp"
+
 #include <algorithm>
+#include <mscclpp/algorithm.hpp>
 #include <mscclpp/core.hpp>
 #include <mscclpp/memory_channel.hpp>
 #include <mscclpp/switch_channel.hpp>
 
 #include "algorithms/allreduce/allreduce_allpair_packet.hpp"
-#include "algorithms/utils.hpp"
+#include "algorithms/allreduce/allreduce_packet.hpp"
 
 namespace mscclpp {
 namespace algorithm {
@@ -148,8 +151,12 @@ std::shared_ptr<mscclpp::DeviceHandle<mscclpp::BaseMemoryChannel>> setupBaseMemo
   return ptr;
 }
 
-std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders() {
-  return std::vector<std::shared_ptr<AlgorithmBuilder>>();
+std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders(uintptr_t scratchBuffer,
+                                                                           size_t scratchBufferSize) {
+  std::vector<std::shared_ptr<AlgorithmBuilder>> builders;
+  builders.push_back(std::make_shared<AllreducePacket>(scratchBuffer, scratchBufferSize));
+  builders.push_back(std::make_shared<AllreduceAllpairPacket>(scratchBuffer, scratchBufferSize));
+  return builders;
 }
 
 std::shared_ptr<AlgorithmBuilder> getDefaultNativeAlgorithmBuilder(std::string algorithmName, uintptr_t scratchBuffer,

src/include/algorithms/allreduce/allreduce_packet.hpp

@@ -22,8 +22,8 @@ class AllreducePacket : public AlgorithmBuilder {
                                                      size_t, DataType);
   AlgorithmCtxKey generateAllreduceContextKey(const void*, void*, size_t, DataType);
 
-  size_t scratchBufferSize_;
   void* scratchBuffer_;
+  size_t scratchBufferSize_;
   const int nSegmentsForScratchBuffer_ = 2;
   const int maxBlockNum_ = 28;
   std::vector<Connection> conns_;

src/include/algorithms/allreduce/common.hpp

@@ -281,14 +281,17 @@ __forceinline__ __device__ __fp8x2_e4m3 min_elements(__fp8x2_e4m3 a, __fp8x2_e4m
 // FP8 E4M3 vectorized min for 4 elements
 __forceinline__ __device__ __fp8x4_e4m3 min_elements(__fp8x4_e4m3 a, __fp8x4_e4m3 b) {
   // Process as two __fp8x2_e4m3 using min_elements for 2 elements
-  __fp8x2_e4m3* a_pair = reinterpret_cast<__fp8x2_e4m3*>(&a);
-  __fp8x2_e4m3* b_pair = reinterpret_cast<__fp8x2_e4m3*>(&b);
+  union {
+    __fp8x4_e4m3 vec4;
+    __fp8x2_e4m3 vec2[2];
+  } ua, ub, uresult;
+  ua.vec4 = a;
+  ub.vec4 = b;
 
-  __fp8x2_e4m3 result[2];
-  result[0] = min_elements(a_pair[0], b_pair[0]);
-  result[1] = min_elements(a_pair[1], b_pair[1]);
+  uresult.vec2[0] = min_elements(ua.vec2[0], ub.vec2[0]);
+  uresult.vec2[1] = min_elements(ua.vec2[1], ub.vec2[1]);
 
-  return *reinterpret_cast<__fp8x4_e4m3*>(result);
+  return uresult.vec4;
 }
 
 // FP8 E5M2 min operation (single element)
@@ -310,14 +313,17 @@ __forceinline__ __device__ __fp8x2_e5m2 min_elements(__fp8x2_e5m2 a, __fp8x2_e5m
 // FP8 E5M2 vectorized min for 4 elements (CUDA only)
 __forceinline__ __device__ __fp8x4_e5m2 min_elements(__fp8x4_e5m2 a, __fp8x4_e5m2 b) {
   // Process as two __fp8x2_e5m2 using min_elements for 2 elements
-  __fp8x2_e5m2* a_pair = reinterpret_cast<__fp8x2_e5m2*>(&a);
-  __fp8x2_e5m2* b_pair = reinterpret_cast<__fp8x2_e5m2*>(&b);
+  union {
+    __fp8x4_e5m2 vec4;
+    __fp8x2_e5m2 vec2[2];
+  } ua, ub, uresult;
+  ua.vec4 = a;
+  ub.vec4 = b;
 
-  __fp8x2_e5m2 result[2];
-  result[0] = min_elements(a_pair[0], b_pair[0]);
-  result[1] = min_elements(a_pair[1], b_pair[1]);
+  uresult.vec2[0] = min_elements(ua.vec2[0], ub.vec2[0]);
+  uresult.vec2[1] = min_elements(ua.vec2[1], ub.vec2[1]);
 
-  return *reinterpret_cast<__fp8x4_e5m2*>(result);
+  return uresult.vec4;
 }
 #endif  // !defined(__HIP_PLATFORM_AMD__)
 #endif  // __FP8_TYPES_EXIST__

src/include/algorithms/utils.hpp

@@ -1,7 +1,6 @@
 #ifndef ALGORITHM_UTILS_HPP
 #define ALGORITHM_UTILS_HPP
 
-#include <mscclpp/algorithm.hpp>
 #include <mscclpp/concurrency_device.hpp>
 #include <mscclpp/core.hpp>
 #include <mscclpp/env.hpp>
@@ -17,6 +16,8 @@
 #endif
 
 namespace mscclpp {
+
+class AlgorithmBuilder;
 namespace algorithm {
 constexpr int NUM_NVLS_CONNECTION = 8;
 constexpr int NUM_SEMAPHORES = 64;
@@ -60,10 +61,11 @@ std::vector<BaseMemoryChannel> setupBaseMemoryChannels(
 std::shared_ptr<DeviceHandle<BaseMemoryChannel>> setupBaseMemoryChannelDeviceHandles(
     const std::vector<BaseMemoryChannel>& baseMemoryChannels);
 
-std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders();
+std::vector<std::shared_ptr<AlgorithmBuilder>> loadNativeAlgorithmBuilders(uintptr_t scratchBuffer,
+                                                                           size_t scratchBufferSize);
 
 std::shared_ptr<AlgorithmBuilder> getDefaultNativeAlgorithmBuilder(std::string algorithmName, uintptr_t scratchBuffer,
-                                                                    size_t scratchBufferSize);
+                                                                   size_t scratchBufferSize);
 }  // namespace algorithm
 }  // namespace mscclpp
 #endif  // ALGORITHM_UTILS_HPP