add intranode P2P async transfer and remove chunk

Author: lizhangyu258

python/triton/experimental/tle/language/raw/core.py

@@ -70,12 +70,14 @@ def _tle_raw_call(func, args, *, output_indices, hint, smem, _semantic):
         source_id = func.register_pending_source(hint=hint)
         dsl_region_op = func.create_region_deferred(_semantic.builder, source_id, handles, alias_indices, hint)
     else:
-        if func.compiler.lower() == "nvcc" or (func.compiler.lower() == "clang" and func.target.lower() == "bc"):
-            patch_hash_method_for_pointer_type()
-            module = import_from_path(func.extern_file)
-            target_fn = getattr(module, func.extern_func_name)
-            ret = target_fn(*args, _semantic=_semantic)
-            return ret
+        if func.compiler is not None:
+            if func.compiler.lower() == "nvcc" or (func.compiler.lower() == "clang" and func.target is not None
+                                                   and func.target.lower() == "bc"):
+                patch_hash_method_for_pointer_type()
+                module = import_from_path(func.extern_file)
+                target_fn = getattr(module, func.extern_func_name)
+                ret = target_fn(*args, _semantic=_semantic)
+                return ret
 
         handles = [arg.handle for arg in args]
         context = _semantic.builder.get_context()

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

@@ -165,7 +165,8 @@ def __init__(self, fn: Any, file: Path, *args, **kwargs) -> None:
         self.deferred: Final[bool] = kwargs.get("deferred", False)
         self.__triton_builtin__: Final[bool] = True
 
-        if self.compiler.lower() == "nvcc" and knobs.runtime.add_stages_inspection_hook is None:
+        if self.compiler is not None and self.compiler.lower(
+        ) == "nvcc" and knobs.runtime.add_stages_inspection_hook is None:
             nvcc_cuda_hook = partial(make_cubin_inspection_hook, self)
             knobs.runtime.add_stages_inspection_hook = nvcc_cuda_hook
 

python/tutorials/tle/raw/nvshmem/01-simple-shift/simple-shift.py

@@ -9,6 +9,7 @@
 from common.utils import (
     install_cumodule_hook,
     load_library,
+    tensor_from_pointer,
     prepare_clang_bitcode,
 )
 
@@ -30,18 +31,6 @@ def simple_shift_kernel(destination_ptr, ):
     tle_raw.call(simple_shift, [destination_ptr])
 
 
-def tensor_from_pointer(pointer, shape, dtype, device):
-    num_elements = 1
-    for extent in shape:
-        num_elements *= extent
-    storage = torch._C._construct_storage_from_data_pointer(
-        pointer.value,
-        device,
-        num_elements * dtype.itemsize,
-    )
-    return torch.empty(0, dtype=dtype, device=device).set_(storage).view(shape)
-
-
 def simpe_shift():
     common_path = Path(__file__).parents[1] / "common" / "common-host.so"
     host_path = Path(__file__).with_name("simple-shift-host.so")
@@ -79,6 +68,7 @@ def simpe_shift():
         common_lib,
         mype_in_node.value,
         Path(__file__).with_name("simple-shift-device.bc"),
+        Path(__file__).with_name("simple-shift-device.cu"),
         simple_shift,
     )
 

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

@@ -17,43 +17,33 @@ enum {
 };
 
 extern "C" __device__ __attribute__((always_inline)) void
-ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready, int rank,
-                    int num_chunks) {
-  int chunk_id = (int)blockIdx.x;
-  if (chunk_id >= num_chunks) {
-    return;
-  }
+ag_mark_local_ready(__attribute__((address_space(1))) uint64_t *ready,
+                    int rank) {
   if (threadIdx.x == 0) {
     __threadfence_system();
-    ready[(size_t)rank * num_chunks + chunk_id] = 1;
+    ready[(size_t)rank] = 1;
   }
   __syncthreads();
 }
 
-// One Triton program publishes one chunk of this rank's A slice to one peer.
 extern "C" __device__ __attribute__((always_inline)) void
-ag_publish_local_chunk(__attribute__((address_space(1))) __half *workspace,
-                       __attribute__((address_space(1))) uint64_t *ready,
-                       int elements_per_rank, int elements_per_chunk,
-                       int num_chunks, int rank, int world_size) {
-  int block_id = (int)blockIdx.x;
-  int peer_offset = block_id / num_chunks + 1;
-  int chunk_id = block_id % num_chunks;
+ag_publish_local_rank(__attribute__((address_space(1))) __half *workspace,
+                      __attribute__((address_space(1))) uint64_t *ready,
+                      int elements_per_rank, int rank, int world_size) {
+  int peer_offset = (int)blockIdx.x + 1;
   if (peer_offset >= world_size) {
     return;
   }
 
   int peer = (rank + peer_offset) % world_size;
-  __half *local_chunk = workspace + (size_t)rank * elements_per_rank +
-                        (size_t)chunk_id * elements_per_chunk;
-  uint64_t *chunk_ready = ready + (size_t)rank * num_chunks + chunk_id;
+  __half *local_rank = workspace + (size_t)rank * elements_per_rank;
+  uint64_t *rank_ready = ready + (size_t)rank;
 
-  nvshmemx_putmem_signal_nbi_block(local_chunk, local_chunk,
-                                   (size_t)elements_per_chunk * sizeof(__half),
-                                   chunk_ready, 1, NVSHMEM_SIGNAL_SET, peer);
+  nvshmemx_putmem_signal_nbi_block(local_rank, local_rank,
+                                   (size_t)elements_per_rank * sizeof(__half),
+                                   rank_ready, 1, NVSHMEM_SIGNAL_SET, peer);
 }
 
-// The GEMM program waits for the source chunk whose A rows it will consume.
 extern "C" __device__ __attribute__((always_inline)) void
 ag_wait_ready(__attribute__((address_space(1))) uint64_t *ready,
               int signal_index) {

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

@@ -16,12 +16,10 @@
     }                                                                          \
   } while (0)
 
-extern "C" int ag_gemm_workspace_create(int elements_per_rank, int num_chunks,
-                                        void **workspace, uint64_t **ready,
-                                        int *mype, int *npes, int *mype_in_node,
-                                        int *npes_in_node) {
-  if (elements_per_rank <= 0 || num_chunks <= 0 || workspace == nullptr ||
-      ready == nullptr) {
+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) {
     return -1;
   }
 
@@ -33,8 +31,7 @@ extern "C" int ag_gemm_workspace_create(int elements_per_rank, int num_chunks,
 
   size_t workspace_bytes = (size_t)(*npes) * elements_per_rank * sizeof(__half);
   *workspace = nvshmem_malloc(workspace_bytes);
-  *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes) * num_chunks,
-                                      sizeof(uint64_t));
+  *ready = (uint64_t *)nvshmem_calloc((size_t)(*npes), sizeof(uint64_t));
   if (*workspace == nullptr || *ready == nullptr) {
     if (*ready != nullptr) {
       nvshmem_free(*ready);
@@ -56,3 +53,11 @@ extern "C" int ag_gemm_workspace_destroy(void *workspace, void *ready) {
   nvshmem_free(workspace);
   return 0;
 }
+
+extern "C" void *ag_gemm_peer_workspace_ptr(void *workspace, int peer) {
+  return nvshmem_ptr(workspace, peer);
+}
+
+extern "C" uint64_t *ag_gemm_peer_ready_ptr(uint64_t *ready, int peer) {
+  return (uint64_t *)nvshmem_ptr(ready, peer);
+}