[KMCompiler][TLERaw] Adjust ag-gemm implementation for multi-node execution (#918)

Co-authored-by: iwanna-lxy <37344393+i3wanna2@users.noreply.github.com>

Author: lizhangyu258

python/triton/experimental/tle/raw/cuda/runtime.py

@@ -234,6 +234,7 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None:
         self.lowered_region_dialect: Final[str] = "llvm"
         self.arg_dialect: Final[str] = "llvm"
         self.source_file: Final[str] = str(file)
+        self.opt_level: Final[str] = kwargs.get("opt_level", "-O3")
 
         if self.library == "nvshmem":
             from triton.experimental.tle.raw.nvshmem.utils import enable_nvshmem_device_bc
@@ -259,7 +260,7 @@ def register_pending_source(self, *, hint: str = "") -> str:
             extern_func_name=self.extern_func_name,
             source=self.code,
             hint=hint,
-            extra={"source_file": self.source_file},
+            extra={"source_file": self.source_file, "opt_level": self.opt_level},
         )
 
     def create_region_by_llvm(self, builder, llvm: str, handles, alias_indices, hint: str = "",
@@ -285,7 +286,7 @@ def make_llvm(self, mlir_context) -> str:
                 "--cuda-device-only",
                 _get_cuda_gpu_arch(),
                 "-emit-llvm",
-                "-O2",
+                self.opt_level,
                 "-S",
                 "-",
                 "-o",
@@ -314,5 +315,6 @@ def read_text(self):
         file=_CudaSourceFile(),
         extern_func_name=entry.get("extern_func_name"),
         deferred=True,
+        opt_level=entry.get("opt_level"),
     )
     return cuda_fn.make_llvm(context)

python/triton/experimental/tle/raw/nvshmem/utils.py

@@ -336,6 +336,18 @@ def set_signal_cuda_ptr(signal_ptr, signal, stream):
     CUDA_CHECK(err)
 
 
+def copy_on_stream(dst_ptr, src_ptr, nbytes, stream):
+    _, cudart = _get_cuda_modules()
+    (err, ) = cudart.cudaMemcpyAsync(
+        dst_ptr,
+        src_ptr,
+        nbytes,
+        cudart.cudaMemcpyKind.cudaMemcpyDefault,
+        stream.cuda_stream,
+    )
+    CUDA_CHECK(err)
+
+
 def print_perf(
     name: str,
     value: float,

python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-device.cu

@@ -1,10 +1,13 @@
-#include <cuda_fp16.h>
 #include <stddef.h>
 #include <stdint.h>
 
 extern "C" __device__ uint64_t nvshmem_signal_wait_until(uint64_t *sig_addr,
                                                          int cmp,
                                                          uint64_t cmp_value);
+extern "C" __device__ void
+nvshmemx_putmem_signal_block(void *dest, const void *source, size_t bytes,
+                             uint64_t *sig_addr, uint64_t signal, int sig_op,
+                             int pe);
 
 enum {
   NVSHMEM_CMP_GE = 5,
@@ -16,7 +19,7 @@ ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready,
                     int rank) {
   if (threadIdx.x == 0) {
     __threadfence_system();
-    ready[(size_t)rank] = 1;
+    ready[rank] = 1;
   }
   __syncthreads();
 }
@@ -29,3 +32,13 @@ ag_wait_ready(__attribute__((address_space(1))) uint64_t *ready,
   }
   __syncthreads();
 }
+
+extern "C" __device__ __attribute__((always_inline)) void
+ag_putmem_signal_block(__attribute__((address_space(1))) void *dest,
+                       __attribute__((address_space(1))) const void *source,
+                       size_t bytes,
+                       __attribute__((address_space(1))) uint64_t *sig_addr,
+                       uint64_t signal, int peer) {
+  nvshmemx_putmem_signal_block(dest, source, bytes, sig_addr, signal,
+                               NVSHMEM_SIGNAL_SET, peer);
+}

python/tutorials/tle/raw/nvshmem/02-allgather-gemm/ag-gemm-host.cu

@@ -1,10 +1,8 @@
-#include <cuda_fp16.h>
 #include <cuda_runtime.h>
 #include <nvshmem.h>
 #include <nvshmemx.h>
 #include <stdint.h>
 #include <stdio.h>
-#include <stdlib.h>
 
 #define CUDA_CHECK(stmt)                                                       \
   do {                                                                         \
@@ -16,10 +14,12 @@
     }                                                                          \
   } while (0)
 
-extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace,
-                                        uint64_t **ready, int *mype, int *npes,
-                                        int *mype_in_node, int *npes_in_node) {
-  if (elements_per_rank <= 0 || workspace == nullptr || ready == nullptr) {
+extern "C" int ag_gemm_workspace_create(int elements_per_rank, int element_size,
+                                        void **workspace, uint64_t **ready,
+                                        int *mype, int *npes, int *mype_in_node,
+                                        int *npes_in_node) {
+  if (elements_per_rank <= 0 || element_size <= 0 || workspace == nullptr ||
+      ready == nullptr) {
     return -1;
   }
 
@@ -29,7 +29,7 @@ extern "C" int ag_gemm_workspace_create(int elements_per_rank, void **workspace,
   *npes_in_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE);
   CUDA_CHECK(cudaSetDevice(*mype_in_node));
 
-  size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * sizeof(__half);
+  size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * element_size;
   *workspace = nvshmem_malloc(workspace_bytes);
   *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes), sizeof(uint64_t));
   if (*workspace == nullptr || *ready == nullptr) {
@@ -61,3 +61,13 @@ extern "C" void *ag_gemm_peer_workspace_ptr(void *workspace, int peer) {
 extern "C" uint64_t *ag_gemm_peer_ready_ptr(uint64_t *ready, int peer) {
   return (uint64_t *)nvshmem_ptr(ready, peer);
 }
+
+extern "C" void ag_gemm_barrier_all_on_stream(cudaStream_t stream) {
+  nvshmemx_barrier_all_on_stream(stream);
+}
+
+extern "C" void ag_gemm_signal_wait_until_on_stream(uint64_t *signal,
+                                                    uint64_t value,
+                                                    cudaStream_t stream) {
+  nvshmemx_signal_wait_until_on_stream(signal, NVSHMEM_CMP_GE, value, stream);
+}