feat: windows support

.gitignore

@@ -4,6 +4,8 @@
 *.dSYM
 *.data
 *.bin
+*.exe
+
 __pycache__
 
 *-test

Makefile

@@ -1,18 +1,28 @@
 CXX = g++
 CXXFLAGS = -std=c++11 -Werror -O3 -march=native -mtune=native
 
+# Default settings
+INCLUDES =
+LIBPATH =
+LIBS = -lpthread
+
+# Conditional settings for Windows
+ifdef WIN32
+    LIBS = -lws2_32 # or -lpthreadGC2 if needed
+endif
+
 utils: src/utils.cpp
-	$(CXX) $(CXXFLAGS) -c src/utils.cpp -o utils.o
+	$(CXX) $(CXXFLAGS) $(INCLUDES) -c src/utils.cpp -o utils.o
 quants: src/quants.cpp
 	$(CXX) $(CXXFLAGS) -c src/quants.cpp -o quants.o
 funcs: src/funcs.cpp
 	$(CXX) $(CXXFLAGS) -c src/funcs.cpp -o funcs.o
 funcs-test: src/funcs-test.cpp funcs
 	$(CXX) $(CXXFLAGS) src/funcs-test.cpp -o funcs-test funcs.o
 socket: src/socket.cpp
-	$(CXX) $(CXXFLAGS) -c src/socket.cpp -o socket.o
+	$(CXX) $(CXXFLAGS) $(INCLUDES) -c src/socket.cpp -o socket.o
 transformer: src/utils.cpp
-	$(CXX) $(CXXFLAGS) -c src/transformer.cpp -o transformer.o
+	$(CXX) $(CXXFLAGS) $(INCLUDES) -c src/transformer.cpp -o transformer.o
 tasks: src/tasks.cpp
 	$(CXX) $(CXXFLAGS) -c src/tasks.cpp -o tasks.o
 llama2-tasks: src/llama2-tasks.cpp
@@ -22,21 +32,22 @@ grok1-tasks: src/grok1-tasks.cpp
 mixtral-tasks: src/mixtral-tasks.cpp
 	$(CXX) $(CXXFLAGS) -c src/mixtral-tasks.cpp -o mixtral-tasks.o
 tokenizer: src/tokenizer.cpp
-	$(CXX) $(CXXFLAGS) -c src/tokenizer.cpp -o tokenizer.o
+	$(CXX) $(CXXFLAGS) $(INCLUDES) -c src/tokenizer.cpp -o tokenizer.o
 app: src/app.cpp
 	$(CXX) $(CXXFLAGS) -c src/app.cpp -o app.o
 
 dllama: src/apps/dllama/dllama.cpp utils quants funcs socket transformer tasks llama2-tasks grok1-tasks mixtral-tasks tokenizer app
-	$(CXX) $(CXXFLAGS) src/apps/dllama/dllama.cpp -o dllama utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o mixtral-tasks.o tokenizer.o app.o -lpthread
+	$(CXX) $(CXXFLAGS) src/apps/dllama/dllama.cpp -o dllama utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o mixtral-tasks.o tokenizer.o app.o $(LIBPATH) $(LIBS)
 dllama-api: src/apps/dllama-api/dllama-api.cpp utils quants funcs socket transformer tasks llama2-tasks grok1-tasks mixtral-tasks tokenizer app
-	$(CXX) $(CXXFLAGS) src/apps/dllama-api/dllama-api.cpp -o dllama-api utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o mixtral-tasks.o tokenizer.o app.o -lpthread
+	$(CXX) $(CXXFLAGS) src/apps/dllama-api/dllama-api.cpp -o dllama-api utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o mixtral-tasks.o tokenizer.o app.o $(LIBPATH) $(LIBS)
+
 funcs-test: src/funcs-test.cpp funcs utils quants
-	$(CXX) $(CXXFLAGS) src/funcs-test.cpp -o funcs-test funcs.o utils.o quants.o -lpthread
+	$(CXX) $(CXXFLAGS) src/funcs-test.cpp -o funcs-test funcs.o utils.o quants.o $(LIBPATH) $(LIBS)
 quants-test: src/quants.cpp utils quants
-	$(CXX) $(CXXFLAGS) src/quants-test.cpp -o quants-test utils.o quants.o -lpthread
+	$(CXX) $(CXXFLAGS) src/quants-test.cpp -o quants-test utils.o quants.o $(LIBPATH) $(LIBS)
 transformer-test: src/transformer-test.cpp funcs utils quants transformer socket
-	$(CXX) $(CXXFLAGS) src/transformer-test.cpp -o transformer-test funcs.o utils.o quants.o transformer.o socket.o -lpthread
+	$(CXX) $(CXXFLAGS) src/transformer-test.cpp -o transformer-test funcs.o utils.o quants.o transformer.o socket.o $(LIBPATH) $(LIBS)
 llama2-tasks-test: src/llama2-tasks-test.cpp utils quants funcs socket transformer tasks llama2-tasks tokenizer
-	$(CXX) $(CXXFLAGS) src/llama2-tasks-test.cpp -o llama2-tasks-test utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o tokenizer.o -lpthread
+	$(CXX) $(CXXFLAGS) src/llama2-tasks-test.cpp -o llama2-tasks-test utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o tokenizer.o $(LIBPATH) $(LIBS)
 grok1-tasks-test: src/grok1-tasks-test.cpp utils quants funcs socket transformer tasks llama2-tasks grok1-tasks tokenizer
-	$(CXX) $(CXXFLAGS) src/grok1-tasks-test.cpp -o grok1-tasks-test utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o tokenizer.o -lpthread
+	$(CXX) $(CXXFLAGS) src/grok1-tasks-test.cpp -o grok1-tasks-test utils.o quants.o funcs.o socket.o transformer.o tasks.o llama2-tasks.o grok1-tasks.o tokenizer.o $(LIBPATH) $(LIBS)

src/app.cpp

@@ -4,6 +4,7 @@
 #include <cstdio>
 #include <cassert>
 #include <stdexcept>
+#include <ctime>
 #include "app.hpp"
 
 FloatType parseFloatType(char* val) {

src/apps/dllama-api/dllama-api.cpp

@@ -5,10 +5,17 @@
 #include <cassert>
 #include <sstream>
 #include <iostream>
+#include <vector>
+
+#ifdef _WIN32
+#include <winsock2.h>
+#include <ws2tcpip.h>
+#pragma comment(lib, "ws2_32.lib")
+#else
 #include <sys/socket.h>
 #include <netinet/in.h>
 #include <unistd.h>
-#include <vector>
+#endif
 
 #include "../../utils.hpp"
 #include "../../socket.hpp"

src/common/pthread.h

@@ -0,0 +1,40 @@
+#ifndef PTHREAD_WRAPPER
+#define PTHREAD_WRAPPER
+
+#ifdef _WIN32 
+#include <windows.h>
+
+typedef HANDLE dl_thread;
+typedef DWORD thread_ret_t;
+typedef DWORD (WINAPI *thread_func_t)(void *);
+
+static int pthread_create(dl_thread * out, void * unused, thread_func_t func, void * arg) {
+    (void) unused;
+    dl_thread handle = CreateThread(NULL, 0, func, arg, 0, NULL);
+    if (handle == NULL) {
+        return EAGAIN;
+    }
+
+    *out = handle;
+    return 0;
+}
+
+static int pthread_join(dl_thread thread, void * unused) {
+    (void) unused;
+    DWORD ret = WaitForSingleObject(thread, INFINITE);
+    if (ret == WAIT_FAILED) {
+        return -1;
+    }
+    CloseHandle(thread);
+    return 0;
+}
+#else
+#include <pthread.h>
+
+typedef pthread_t dl_thread;
+typedef void* thread_ret_t;
+typedef void* (*thread_func_t)(void *);
+
+#endif
+
+#endif  // PTHREAD_WRAPPER

src/funcs.cpp

@@ -1,7 +1,7 @@
 #include <cmath>
 #include <cassert>
 #include <cstdio>
-#include <pthread.h>
+#include "common/pthread.h"
 #include "quants.hpp"
 #include "funcs.hpp"
 
@@ -145,7 +145,7 @@ void rmsnorm(float* o, const float* x, const float ms, const float* weight, cons
 }
 
 struct MatmulThreadInfo {
-    pthread_t handler;
+    dl_thread handler;
     float* output;
     const void* input;
     const void* weights;

src/socket.cpp

@@ -1,21 +1,39 @@
 #include <cstdio>
 #include <cstdlib>
 #include <cassert>
-#include <arpa/inet.h>
-#include <netinet/tcp.h>
 #include <errno.h>
 #include <string.h>
 #include <fcntl.h>
 #include <ctime>
-#include <unistd.h>
 #include <stdexcept>
 #include <vector>
+#include <string>
+#include <iostream>
 #include "socket.hpp"
 
+#ifdef _WIN32
+#include <winsock2.h>
+#include <ws2tcpip.h> // For inet_addr and other functions
+#include <windows.h>  // For SSIZE_T
+typedef SSIZE_T ssize_t;
+#define close closesocket
+#else
+#include <sys/socket.h>
+#include <netinet/tcp.h>
+#include <arpa/inet.h>
+#include <unistd.h>
+#endif
+
 #define SOCKET_LAST_ERRCODE errno
 #define SOCKET_LAST_ERROR strerror(errno)
 
 static inline void setNonBlocking(int socket, bool enabled) {
+#ifdef _WIN32
+    u_long mode = enabled ? 1 : 0;
+    if (ioctlsocket(socket, FIONBIO, &mode) != 0) {
+        throw std::runtime_error("Error setting socket to non-blocking");
+    }
+#else
     int flags = fcntl(socket, F_GETFL, 0);
     if (enabled) {
         flags |= O_NONBLOCK;
@@ -24,6 +42,7 @@ static inline void setNonBlocking(int socket, bool enabled) {
     }
     if (fcntl(socket, F_SETFL, flags) < 0)
         throw std::runtime_error("Error setting socket to non-blocking");
+#endif
 }
 
 static inline void setNoDelay(int socket) {
@@ -33,35 +52,54 @@ static inline void setNoDelay(int socket) {
 }
 
 static inline void setQuickAck(int socket) {
+#ifndef _WIN32
 #ifdef TCP_QUICKACK
     int value = 1;
     if (setsockopt(socket, IPPROTO_TCP, TCP_QUICKACK, (char*)&value, sizeof(int)) < 0)
         throw std::runtime_error("Error setting quick ack");
 #endif
+#endif
 }
 
 static inline void writeSocket(int socket, const void* data, size_t size) {
     while (size > 0) {
-        int s = send(socket, (char*)data, size, 0);
+        int s = send(socket, (const char*)data, size, 0);
         if (s < 0) {
+#ifdef _WIN32
+            if (WSAGetLastError() == WSAEWOULDBLOCK) {
+                continue;
+            }
+#else
             if (SOCKET_LAST_ERRCODE == EAGAIN) {
                 continue;
             }
-            throw WriteSocketException(SOCKET_LAST_ERRCODE, SOCKET_LAST_ERROR);
+#endif
+            throw std::runtime_error("Error writing to socket");
         } else if (s == 0) {
-            throw WriteSocketException(0, "Socket closed");
+            throw std::runtime_error("Socket closed");
         }
         size -= s;
-        data = (char*)data + s;
+        data = (const char*)data + s;
     }
 }
 
 static inline bool tryReadSocket(int socket, void* data, size_t size, unsigned long maxAttempts) {
     // maxAttempts = 0 means infinite attempts
     size_t s = size;
     while (s > 0) {
-        int r = recv(socket, data, s, 0);
+        int r = recv(socket, (char*)data, s, 0);
         if (r < 0) {
+#ifdef _WIN32
+            if (WSAGetLastError() == WSAEWOULDBLOCK) {
+                if (s == size && maxAttempts > 0) {
+                    maxAttempts--;
+                    if (maxAttempts == 0) {
+                        return false;
+                    }
+                }
+                continue;
+            }
+#else
             if (SOCKET_LAST_ERRCODE == EAGAIN) {
                 if (s == size && maxAttempts > 0) {
                     maxAttempts--;
@@ -71,9 +109,10 @@ static inline bool tryReadSocket(int socket, void* data, size_t size, unsigned l
                 }
                 continue;
             }
-            throw ReadSocketException(SOCKET_LAST_ERRCODE, SOCKET_LAST_ERROR);
+#endif
+            throw std::runtime_error("Error reading from socket");
         } else if (r == 0) {
-            throw ReadSocketException(0, "Socket closed");
+            throw std::runtime_error("Socket closed");
         }
         data = (char*)data + r;
         s -= r;
@@ -82,7 +121,9 @@ static inline bool tryReadSocket(int socket, void* data, size_t size, unsigned l
 }
 
 static inline void readSocket(int socket, void* data, size_t size) {
-    assert(tryReadSocket(socket, data, size, 0));
+    if (!tryReadSocket(socket, data, size, 0)) {
+        throw std::runtime_error("Error reading from socket");
+    }
 }
 
 ReadSocketException::ReadSocketException(int code, const char* message) {
@@ -169,11 +210,17 @@ void SocketPool::writeMany(unsigned int n, SocketIo* ios) {
             if (io->size > 0) {
                 isWriting = true;
                 int socket = sockets[io->socketIndex];
-                ssize_t s = send(socket, io->data, io->size, 0);
+                ssize_t s = send(socket, (const char*)io->data, io->size, 0);
                 if (s < 0) {
+#ifdef _WIN32
+                    if (WSAGetLastError() == WSAEWOULDBLOCK) {
+                        continue;
+                    }
+#else
                     if (SOCKET_LAST_ERRCODE == EAGAIN) {
                         continue;
                     }
+#endif
                     throw WriteSocketException(SOCKET_LAST_ERRCODE, SOCKET_LAST_ERROR);
                 } else if (s == 0) {
                     throw WriteSocketException(0, "Socket closed");
@@ -295,30 +342,73 @@ SocketServer::SocketServer(int port) {
     const char* host = "0.0.0.0";
     struct sockaddr_in serverAddr;
 
-    socket = ::socket(AF_INET, SOCK_STREAM, 0);
+    #ifdef _WIN32
+    WSADATA wsaData;
+    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
+        throw std::runtime_error("WSAStartup failed: " + std::to_string(WSAGetLastError()));
+    }
+    #endif
+
+    socket = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
     if (socket < 0)
         throw std::runtime_error("Cannot create socket");
 
+    int opt = 1;
+    #ifdef _WIN32
+    int iresult = setsockopt(socket, SOL_SOCKET, SO_REUSEADDR, (char*)&opt, sizeof(opt));
+    if (iresult == SOCKET_ERROR) {
+        closesocket(socket);
+        WSACleanup();
+        throw std::runtime_error("setsockopt failed: " + std::to_string(WSAGetLastError()));
+    }
+    #else
+    if (setsockopt(socket, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
+        close(socket);
+        throw std::runtime_error("setsockopt failed: " + std::string(strerror(errno)));
+    }
+    #endif
+
     memset(&serverAddr, 0, sizeof(serverAddr));
     serverAddr.sin_family = AF_INET;
     serverAddr.sin_port = htons(port);
     serverAddr.sin_addr.s_addr = inet_addr(host);
 
-    int bindResult = bind(socket, (struct sockaddr*)&serverAddr, sizeof(serverAddr));
+    int bindResult;
+    #ifdef _WIN32
+    bindResult = bind(socket, (SOCKADDR*)&serverAddr, sizeof(serverAddr));
+    if (bindResult == SOCKET_ERROR) {
+        int error = WSAGetLastError();
+        closesocket(socket);
+        WSACleanup();
+        throw std::runtime_error("Cannot bind port: " + std::to_string(error));
+    }
+    #else
+    bindResult = bind(socket, (struct sockaddr*)&serverAddr, sizeof(serverAddr));
     if (bindResult < 0) {
-        printf("Cannot bind %s:%d\n", host, port);
-        throw std::runtime_error("Cannot bind port");
+        close(socket);
+        throw std::runtime_error("Cannot bind port: " + std::string(strerror(errno)));
     }
+    #endif
 
-    int listenResult = listen(socket, 1);
+    int listenResult = listen(socket, SOMAXCONN);
     if (listenResult != 0) {
-        printf("Cannot listen %s:%d\n", host, port);
-        throw std::runtime_error("Cannot listen port");
+        #ifdef _WIN32
+        closesocket(socket);
+        WSACleanup();
+        throw std::runtime_error("Cannot listen on port: " + std::to_string(WSAGetLastError()));
+        #else
+        close(socket);
+        throw std::runtime_error("Cannot listen on port: " + std::string(strerror(errno)));
+        #endif
     }
+
     printf("Listening on %s:%d...\n", host, port);
 }
 
 SocketServer::~SocketServer() {
     shutdown(socket, 2);
+    #ifdef _WIN32
+    WSACleanup();
+    #endif
     close(socket);
 }