PERF: Do not make copies if the number of references is only 1

pavanky · pavanky · commit 51b90c78e8b6 · 2015-06-30T11:04:22.000-04:00
diff --git a/src/api/c/assign.cpp b/src/api/c/assign.cpp
@@ -116,9 +116,19 @@ af_err af_assign_seq(af_array *out,
             ARG_ASSERT(2, (index[i].step>=0));
         }
 
-        af_array res;
-        if (*out != lhs) AF_CHECK(af_copy_array(&res, lhs));
-        else             res = lhs;
+        af_array res = 0;
+
+        if (*out != lhs) {
+            int count = 0;
+            AF_CHECK(af_get_data_ref_count(&count, lhs));
+            if (count > 1) {
+                AF_CHECK(af_copy_array(&res, lhs));
+            } else {
+                AF_CHECK(af_retain_array(&res, lhs));
+            }
+        } else {
+            res = lhs;
+        }
 
         try {
 
@@ -190,8 +200,17 @@ af_err af_assign_gen(af_array *out,
         ARG_ASSERT(1, (lhs!=0));
         ARG_ASSERT(4, (rhs!=0));
 
-        if (*out != lhs) AF_CHECK(af_copy_array(&output, lhs));
-        else             output = lhs;
+        if (*out != lhs) {
+            int count = 0;
+            AF_CHECK(af_get_data_ref_count(&count, lhs));
+            if (count > 1) {
+                AF_CHECK(af_copy_array(&output, lhs));
+            } else {
+                AF_CHECK(af_retain_array(&output, lhs));
+            }
+        } else {
+            output = lhs;
+        }
 
         ArrayInfo lInfo = getInfo(lhs);
         ArrayInfo rInfo = getInfo(rhs);
diff --git a/test/memory.cpp b/test/memory.cpp
@@ -263,18 +263,15 @@ TEST(Memory, Assign)
     ASSERT_EQ(lock_bytes, 1 * step_bytes);
 
     {
-        // Should just a copy
         af::array b = af::randu(num / 2);
         a(af::seq(num / 2)) = b;
 
         af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
                           &lock_bytes, &lock_buffers);
 
-        // FIXME: An extra buffer is used because of copy on write
-        // Fix to not perform a copy when the buffer does not have children
-        ASSERT_EQ(alloc_buffers, 3u);
+        ASSERT_EQ(alloc_buffers, 2u);
         ASSERT_EQ(lock_buffers, 2u);
-        ASSERT_EQ(alloc_bytes, 3 * step_bytes);
+        ASSERT_EQ(alloc_bytes, 2 * step_bytes);
         ASSERT_EQ(lock_bytes, 2 * step_bytes);
     }
 
@@ -283,13 +280,140 @@ TEST(Memory, Assign)
     af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
                       &lock_bytes, &lock_buffers);
 
-    ASSERT_EQ(alloc_buffers, 3u);
+    ASSERT_EQ(alloc_buffers, 2u);
     ASSERT_EQ(lock_buffers, 1u);
-    ASSERT_EQ(alloc_bytes, 3 * step_bytes);
+    ASSERT_EQ(alloc_bytes, 2 * step_bytes);
     ASSERT_EQ(lock_bytes, 1 * step_bytes);
 
 }
 
+TEST(Memory, AssignLoop)
+{
+    size_t alloc_bytes, alloc_buffers;
+    size_t lock_bytes, lock_buffers;
+
+    cleanSlate(); // Clean up everything done so far
+
+    const int num = step_bytes / sizeof(float);
+    const int cols = 100;
+
+    af::array a = af::randu(num, cols);
+
+    af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                      &lock_bytes, &lock_buffers);
+
+    ASSERT_EQ(alloc_buffers, 1u);
+    ASSERT_EQ(lock_buffers, 1u);
+    ASSERT_EQ(alloc_bytes, cols * step_bytes);
+    ASSERT_EQ(lock_bytes, cols * step_bytes);
+
+    for (int i = 0; i < cols; i++) {
+
+        af::array b = af::randu(num);
+        a(af::span, i) = b;
+
+        af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                          &lock_bytes, &lock_buffers);
+
+        ASSERT_EQ(alloc_buffers, 2u); // 3 because you need another scratch space for b
+        ASSERT_EQ(lock_buffers, 2u);
+        ASSERT_EQ(alloc_bytes, (cols + 1) * step_bytes);
+        ASSERT_EQ(lock_bytes, (cols + 1) * step_bytes);
+    }
+}
+
+TEST(Memory, AssignRef)
+{
+    size_t alloc_bytes, alloc_buffers;
+    size_t lock_bytes, lock_buffers;
+
+    cleanSlate(); // Clean up everything done so far
+
+    const int num = step_bytes / sizeof(float);
+
+    af::array a = af::randu(num);
+    af::array a_ref = a;
+
+    af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                      &lock_bytes, &lock_buffers);
+
+    ASSERT_EQ(alloc_buffers, 1u);
+    ASSERT_EQ(lock_buffers, 1u);
+    ASSERT_EQ(alloc_bytes, 1 * step_bytes);
+    ASSERT_EQ(lock_bytes, 1 * step_bytes);
+
+    {
+        af::array b = af::randu(num / 2);
+        // This should do a full copy of a
+        a(af::seq(num / 2)) = b;
+
+        af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                          &lock_bytes, &lock_buffers);
+
+        ASSERT_EQ(alloc_buffers, 3u);
+        ASSERT_EQ(lock_buffers, 3u);
+        ASSERT_EQ(alloc_bytes, 3 * step_bytes);
+        ASSERT_EQ(lock_bytes, 3 * step_bytes);
+    }
+
+
+    // b should not have deleted a
+    af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                      &lock_bytes, &lock_buffers);
+
+    ASSERT_EQ(alloc_buffers, 3u);
+    ASSERT_EQ(lock_buffers, 2u); // a_ref
+    ASSERT_EQ(alloc_bytes, 3 * step_bytes);
+    ASSERT_EQ(lock_bytes, 2 * step_bytes); // a_ref
+
+}
+
+TEST(Memory, AssignRefLoop)
+{
+    size_t alloc_bytes, alloc_buffers;
+    size_t lock_bytes, lock_buffers;
+
+    cleanSlate(); // Clean up everything done so far
+
+    const int num = step_bytes / sizeof(float);
+    const int cols = 100;
+
+    af::array a = af::randu(num, cols);
+    af::array a_ref = a;
+
+    af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                      &lock_bytes, &lock_buffers);
+
+    ASSERT_EQ(alloc_buffers, 1u);
+    ASSERT_EQ(lock_buffers, 1u);
+    ASSERT_EQ(alloc_bytes, cols * step_bytes);
+    ASSERT_EQ(lock_bytes, cols * step_bytes);
+
+    for (int i = 0; i < cols; i++) {
+
+        af::array b = af::randu(num);
+        a(af::span, i) = b;
+
+        af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                          &lock_bytes, &lock_buffers);
+
+        ASSERT_EQ(alloc_buffers, 3u);
+        ASSERT_EQ(lock_buffers, 3u);
+        ASSERT_EQ(alloc_bytes, (2 * cols + 1) * step_bytes);
+        ASSERT_EQ(lock_bytes, (2 * cols + 1) * step_bytes);
+    }
+
+
+    // b should not have deleted a
+    af::deviceMemInfo(&alloc_bytes, &alloc_buffers,
+                      &lock_bytes, &lock_buffers);
+
+    ASSERT_EQ(alloc_buffers, 3u);
+    ASSERT_EQ(lock_buffers, 2u); // a_ref
+    ASSERT_EQ(alloc_bytes, (2 * cols + 1) * step_bytes);
+    ASSERT_EQ(lock_bytes, 2 * cols * step_bytes); // a_ref
+
+}
 
 TEST(Memory, device)
 {
