From 9604114da019d0934f5869c9e04d86e077faaac4 Mon Sep 17 00:00:00 2001
From: ct-clmsn <ct.clmsn@gmail.com>
Date: Wed, 20 Dec 2023 23:19:51 -0500
Subject: [PATCH] added baseline makefile support; fixed several compilation
 warnings

---
 Makefile      | 47 +++++++++++++++++++++++++++++++++++++++++++++++
 ggml-oshmem.c | 12 ++++++------
 ggml-oshmem.h |  2 +-
 3 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 8273f8400..0de6b579e 100644
--- a/Makefile
+++ b/Makefile
@@ -345,6 +345,53 @@ ifdef LLAMA_MPI
 	OBJS        += ggml-mpi.o
 endif # LLAMA_MPI
 
+ifdef LLAMA_OPENSHMEM
+
+        OPENSHMEM_FOUND:=0
+	PKG:=sandia-openshmem
+	REQPKG:=$(shell pkg-config --exists $(PKG) && echo '$(PKG)')
+	ifneq ($(REQPKG),)
+                OPENSHMEM_FOUND:=1
+		OPENSHMEM_CFLAGS:=$(shell pkg-config --cflags-only-other sandia-openshmem)
+		OPENSHMEM_LDFLAGS:=$(shell pkg-config --libs sandia-openshmem)
+	else
+		$(warning '$(PKG)' not found)
+	endif
+
+        ifneq($(OPENSHMEM_FOUND),1)
+		PKG:=osss-ucx
+		REQPKG:=$(shell pkg-config --exists $(PKG) && echo '$(PKG)')
+		ifneq ($(REQPKG),)
+        	        OPENSHMEM_FOUND:=1
+			OPENSHMEM_CFLAGS:=$(shell pkg-config --cflags-only-other osss-ucx)
+			OPENSHMEM_LDFLAGS:=$(shell pkg-config --libs osss-ucx)
+		else
+			$(warning '$(PKG)' not found)
+		endif
+	endif
+
+        ifneq($(OPENSHMEM_FOUND),1)
+		PKG:=oshmem
+		REQPKG:=$(shell pkg-config --exists $(PKG) && echo '$(PKG)')
+		ifneq ($(REQPKG),)
+        	        OPENSHMEM_FOUND:=1
+			OPENSHMEM_CFLAGS:=$(shell oshmem_info --path libdir)
+			OPENSHMEM_LDFLAGS:=$(shell oshmem_info --path incdir)
+		else
+			$(warning '$(PKG)' not found)
+		endif
+	endif
+
+        ifneq($(OPENSHMEM_FOUND),1)
+		$(error '$(PKG)' not found)
+	endif
+
+	MK_CPPFLAGS += -DGGML_USE_OPENSHMEM
+	MK_CFLAGS   += -Wno-cast-qual $(OPENSHMEM_CFLAGS)
+	MK_LDFLAGS += -Wno-cast-qual $(OPENSHMEM_LDFLAGS)
+	OBJS        += ggml-oshmem.o
+endif # LLAMA_OPENSHMEM
+
 ifdef LLAMA_OPENBLAS
 	MK_CPPFLAGS += -DGGML_USE_OPENBLAS $(shell pkg-config --cflags-only-I openblas)
 	MK_CFLAGS   += $(shell pkg-config --cflags-only-other openblas)
diff --git a/ggml-oshmem.c b/ggml-oshmem.c
index 6acc3b5d4..ff8477141 100644
--- a/ggml-oshmem.c
+++ b/ggml-oshmem.c
@@ -20,7 +20,7 @@ struct ggml_openshmem_context {
     int64_t symmetric_buffer_size;
     int64_t symmetric_comm_structure_size;
     uint8_t * symmetric_comm_structure;
-    long * recv_signal;
+    uint64_t * recv_signal;
 };
 
 void ggml_openshmem_backend_init(void) {
@@ -56,7 +56,7 @@ struct ggml_openshmem_context * ggml_openshmem_init(void) {
     /*
      * uint8_t signal_byte[shmem_npes()];
      */
-    ctx->recv_signal = (long*)shmem_calloc(1, ctx->n_pes*sizeof(long));
+    ctx->recv_signal = (uint64_t*)shmem_calloc(1, ctx->n_pes*sizeof(uint64_t));
 
     return ctx;
 }
@@ -115,9 +115,9 @@ static void ggml_openshmem_tensor_send(struct ggml_openshmem_context * ctx, stru
         ((int64_t*)dst_symmetric_comm_offset)+sizeof(int64_t);
     uint8_t * dst_symmetric_comm_buffer =
         ((uint8_t*)dst_symmetric_comm_length)+sizeof(int64_t);
-    long * dst_recv_signal =
+    uint64_t * dst_recv_signal =
         ctx->recv_signal+dst_pe;
-    long * my_recv_signal =
+    uint64_t * my_recv_signal =
         ctx->recv_signal+ctx->pe;
 
     const int64_t nelements = ggml_nelements(t);
@@ -203,9 +203,9 @@ static void ggml_openshmem_tensor_recv(struct ggml_openshmem_context * ctx, stru
         ((int64_t*)src_symmetric_comm_offset)+sizeof(int64_t);
     uint8_t * src_symmetric_comm_buffer =
         ((uint8_t*)src_symmetric_comm_length)+sizeof(int64_t);
-    long * src_recv_signal =
+    uint64_t * src_recv_signal =
         ctx->recv_signal+src_pe;
-    long* my_recv_signal =
+    uint64_t * my_recv_signal =
         ctx->recv_signal+ctx->pe;
 
     int64_t total_loop_count = 0;
diff --git a/ggml-oshmem.h b/ggml-oshmem.h
index fb953fb0f..ea88585ad 100644
--- a/ggml-oshmem.h
+++ b/ggml-oshmem.h
@@ -18,7 +18,7 @@ void ggml_openshmem_backend_free(void);
 struct ggml_openshmem_context * ggml_openshmem_init(void);
 void ggml_openshmem_free(struct ggml_openshmem_context * ctx);
 
-int ggml_openshmem_rank(struct ggml_openshmem_context * ctx);
+int ggml_openshmem_pe(struct ggml_openshmem_context * ctx);
 
 void ggml_openshmem_eval_init(
         struct ggml_openshmem_context * ctx_openshmem,