From 3fb866d5b54c901f20f401b5bc31049e08fbd28c Mon Sep 17 00:00:00 2001
From: Umar Arshad <umar@arrayfire.com>
Date: Wed, 4 Oct 2023 12:47:29 -0400
Subject: [PATCH] Fix OpenCL memory migration on devices with different
 contexts

---
 src/backend/opencl/Array.cpp |  6 ++++--
 test/array.cpp               | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/src/backend/opencl/Array.cpp b/src/backend/opencl/Array.cpp
index 21dec5166c..b4b6bcd5a9 100644
--- a/src/backend/opencl/Array.cpp
+++ b/src/backend/opencl/Array.cpp
@@ -201,10 +201,12 @@ void checkAndMigrate(Array<T> &arr) {
         AF_TRACE("Migrating array from {} to {}.", arr_id, cur_id);
         auto migrated_data           = memAlloc<T>(arr.elements());
         void *mapped_migrated_buffer = getQueue().enqueueMapBuffer(
-            *migrated_data, CL_TRUE, CL_MAP_READ, 0, arr.elements());
+            *migrated_data, CL_TRUE, CL_MAP_WRITE_INVALIDATE_REGION, 0,
+            sizeof(T) * arr.elements());
         setDevice(arr_id);
         Buffer &buf = *arr.get();
-        getQueue().enqueueReadBuffer(buf, CL_TRUE, 0, arr.elements(),
+        getQueue().enqueueReadBuffer(buf, CL_TRUE, 0,
+                                     sizeof(T) * arr.elements(),
                                      mapped_migrated_buffer);
         setDevice(cur_id);
         getQueue().enqueueUnmapMemObject(*migrated_data,
diff --git a/test/array.cpp b/test/array.cpp
index bcf6fa997e..b68f06820a 100644
--- a/test/array.cpp
+++ b/test/array.cpp
@@ -501,6 +501,29 @@ TEST(DeviceId, Different) {
     deviceGC();
 }
 
+TEST(Device, MigrateAllDevicesToAllDevices) {
+    int ndevices = getDeviceCount();
+    if (ndevices < 2) GTEST_SKIP() << "Skipping mult-GPU test";
+
+    for (int i = 0; i < ndevices; i++) {
+        for (int j = 0; j < ndevices; j++) {
+            setDevice(i);
+            array a = constant(i * 255, 10, 10);
+            a.eval();
+
+            setDevice(j);
+            array b = constant(j * 256, 10, 10);
+            b.eval();
+
+            array c = a + b;
+
+            std::vector<float> gold(10 * 10, i * 255 + j * 256);
+
+            ASSERT_VEC_ARRAY_EQ(gold, dim4(10, 10), c);
+        }
+    }
+}
+
 TEST(Device, empty) {
     array a = array();
     ASSERT_EQ(a.device<float>(), nullptr);