From f02060da1e58ca701e2e37e732c2a8b9f685f1c9 Mon Sep 17 00:00:00 2001
From: Hannan Naeem <63553569+HannanNaeem@users.noreply.github.com>
Date: Wed, 13 Mar 2024 16:12:22 -0500
Subject: [PATCH] ufuncs + tests: added np equivalent copyto (#264)

---
 pykokkos/__init__.py   |  1 +
 pykokkos/lib/ufuncs.py | 50 +++++++++++++++++++++++++++++++++
 tests/test_ufuncs.py   | 64 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 115 insertions(+)

diff --git a/pykokkos/__init__.py b/pykokkos/__init__.py
index cdc9bcc8..440a836c 100644
--- a/pykokkos/__init__.py
+++ b/pykokkos/__init__.py
@@ -20,6 +20,7 @@
                                  sqrt,
                                  sign,
                                  add,
+                                 copyto,
                                  subtract,
                                  dot,
                                  multiply,
diff --git a/pykokkos/lib/ufuncs.py b/pykokkos/lib/ufuncs.py
index 20aeb0a3..72c302d8 100644
--- a/pykokkos/lib/ufuncs.py
+++ b/pykokkos/lib/ufuncs.py
@@ -1320,6 +1320,56 @@ def subtract(viewA, valB):
 
     return out
 
+@pk.workunit
+def copyto_impl_2d(tid, viewA, viewB):
+    r_idx : int = tid / viewA.extent(1)
+    c_idx : int = tid - r_idx * viewA.extent(1)
+
+    viewA[r_idx][c_idx] = viewB[r_idx][c_idx]
+
+@pk.workunit
+def copyto_impl_1d(tid, viewA, viewB):
+    viewA[tid] = viewB[tid]
+
+def copyto(viewA, viewB):
+    '''
+    copies values of viewB into valueA for corresponding indicies
+
+    Parameters
+    ----------
+    viewA : pykokkos view
+            Input view.
+    valB : pykokkos view or scalar
+            Input view 
+
+    Returns
+    -------
+        Void
+    '''
+
+    if not isinstance(viewA, ViewType):
+        raise ValueError("copyto: Cannot copy to a non-view type")
+    if not isinstance(viewB, ViewType):
+        raise ValueError("copyto: Cannot copy from a non-view type")
+    if viewA.shape != viewB.shape:
+        if not check_broadcastable_impl(viewA, viewB): # if shape is not same check compatibility
+            raise ValueError("copyto: Views must be broadcastable or of the same size. {} against {}".format(viewA.shape, viewB.shape))
+        # check if size is same otherwise broadcast and fix
+        viewA = broadcast_view(viewB, viewA)
+
+    # implementation constraint, for now
+    if viewA.rank() > 2:
+        raise NotImplementedError("copyto: This version of Pykokkos only supports copyto upto 2D views")
+
+    if viewA.rank() == 1:
+        pk.parallel_for(viewA.shape[0], copyto_impl_1d, viewA=viewA, viewB=viewB)
+
+    else:   
+        outRows = viewA.shape[0]
+        outCols = viewA.shape[1] 
+        totalThreads = outRows * outCols
+        pk.parallel_for(totalThreads, copyto_impl_2d, viewA=viewA, viewB=viewB)
+
 @pk.workunit
 def np_matmul_impl_2d_2d(tid, cols, vec_length, viewA, viewB, viewOut):
     r_idx : int = tid / cols
diff --git a/tests/test_ufuncs.py b/tests/test_ufuncs.py
index d8d395cd..4c5c0734 100644
--- a/tests/test_ufuncs.py
+++ b/tests/test_ufuncs.py
@@ -678,6 +678,70 @@ def test_sign_1d_special_cases(in_arr, pk_dtype, numpy_dtype):
     assert_allclose(actual, expected)
 
 
+@pytest.mark.parametrize("pk_ufunc, numpy_ufunc", [
+        (pk.copyto, np.copyto),
+])
+@pytest.mark.parametrize("numpy_dtype", [
+        (np.float64),
+        (np.float32),
+])
+def test_copyto_1d(pk_ufunc, numpy_ufunc, numpy_dtype):
+    N = 4
+    M = 7
+    rng = default_rng(123)
+    np1 = rng.random((N, M)).astype(numpy_dtype)
+    np2 = rng.random((N, M)).astype(numpy_dtype)
+    numpy_ufunc(np1, np2)
+
+    view1 = pk.array(np1)
+    view2 = pk.array(np2)
+    pk_ufunc(view1, view2)
+
+    assert_allclose(np1, view1)
+
+
+@pytest.mark.parametrize("pk_ufunc, numpy_ufunc", [
+        (pk.subtract, np.subtract),
+])
+@pytest.mark.parametrize("numpy_dtype", [
+        (np.float64),
+        (np.float32),
+])
+@pytest.mark.parametrize("test_dim", [
+    [4,3,4,3], [4,3,1,1], [4,3,1,3], [4,3,4,1], [4,3,1], [4,3,3], [4,3], [4]
+])
+def test_copyto_broadcast_2d(pk_ufunc, numpy_ufunc, numpy_dtype, test_dim):
+    np1 = None
+    np2 = None
+    rng = default_rng(123)
+    scalar = 3.0
+
+    if len(test_dim) == 4:
+        np1 = rng.random((test_dim[0], test_dim[1])).astype(numpy_dtype)
+        np2 = rng.random((test_dim[2], test_dim[3])).astype(numpy_dtype)
+    elif len(test_dim) == 3:
+        np1 = rng.random((test_dim[0], test_dim[1])).astype(numpy_dtype)
+        np2 = rng.random((test_dim[2])).astype(numpy_dtype)
+    elif len(test_dim) == 2:
+        np1 = rng.random((test_dim[0], test_dim[1])).astype(numpy_dtype)
+        np2 = scalar # 2d with scalar
+    elif len(test_dim) == 1:
+        np1 = rng.random((test_dim[0])).astype(numpy_dtype)
+        np2 = scalar # 1d with scalar
+    else:
+        raise NotImplementedError("Invalid test conditions: Broadcasting operations are only supported uptil 2D")
+
+    assert np1 is not None and np2 is not None, "Invalid test conditions: Are parameters uptil 2D?"
+    
+    numpy_ufunc(np1, np2)
+
+    view1 = pk.array(np1)
+    view2 = pk.array(np2) if isinstance(np2, np.ndarray) else np2
+    pk_ufunc(view1, view2)
+
+    assert_allclose(np1, view1)
+
+
 @pytest.mark.parametrize("input_dtype", [
         pk.double, pk.float,
 ])