From 5577cada211997879d695df3edc17d092e07cb04 Mon Sep 17 00:00:00 2001
From: matteo serva <matteo.serva@gmail.com>
Date: Thu, 20 Jun 2024 13:51:07 +0200
Subject: [PATCH] Adding support for unified memory

---
 ggml/src/ggml-cuda.cu | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/ggml/src/ggml-cuda.cu b/ggml/src/ggml-cuda.cu
index c73ae40d4..b4b3e94c3 100644
--- a/ggml/src/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda.cu
@@ -558,7 +558,14 @@ GGML_CALL static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffe
     size = std::max(size, (size_t)1); // cudaMalloc returns null for size 0
 
     void * dev_ptr;
-    cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device);
+    cudaError_t err;
+    if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr)
+    {
+        err = cudaMallocManaged(&dev_ptr, size);
+    }
+    else {
+        err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device);
+    }
     if (err != cudaSuccess) {
         // clear the error
         cudaGetLastError();