arrayfire · pavanky · Apr 24, 2016 · Mar 30, 2016 · Mar 30, 2016 · Mar 30, 2016
diff --git a/include/af/algorithm.h b/include/af/algorithm.h
@@ -357,8 +357,6 @@ namespace af
        \return the sorted output
 
        \ingroup sort_func_sort
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI array sort(const array &in, const unsigned dim = 0, const bool isAscending = true);
 
@@ -372,8 +370,6 @@ namespace af
        \param[in] isAscending specifies the sorting order
 
        \ingroup sort_func_sort_index
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI void  sort(array &out, array &indices, const array &in, const unsigned dim = 0,
                      const bool isAscending = true);
@@ -388,8 +384,6 @@ namespace af
        \param[in] isAscending specifies the sorting order
 
        \ingroup sort_func_sort_keys
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI void  sort(array &out_keys, array &out_values, const array &keys, const array &values,
                      const unsigned dim = 0, const bool isAscending = true);
@@ -794,8 +788,6 @@ extern "C" {
        \return \ref AF_SUCCESS if the execution completes properly
 
        \ingroup sort_func_sort
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI af_err af_sort(af_array *out, const af_array in, const unsigned dim, const bool isAscending);
 
@@ -810,8 +802,6 @@ extern "C" {
        \return \ref AF_SUCCESS if the execution completes properly
 
        \ingroup sort_func_sort_index
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI af_err af_sort_index(af_array *out, af_array *indices, const af_array in,
                                const unsigned dim, const bool isAscending);
@@ -827,8 +817,6 @@ extern "C" {
        \return \ref AF_SUCCESS if the execution completes properly
 
        \ingroup sort_func_sort_keys
-
-       \note \p dim is currently restricted to 0.
     */
     AFAPI af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
                                 const af_array keys, const af_array values,

diff --git a/src/api/c/median.cpp b/src/api/c/median.cpp
@@ -68,8 +68,8 @@ static af_array median(const af_array& in, const dim_t dim)
     const Array<T> input = getArray<T>(in);
     Array<T> sortedIn   = sort<T, true>(input, dim);
 
-    int nElems    = input.dims()[0];
-    double mid    = (nElems + 1) / 2;
+    int dimLength = input.dims()[dim];
+    double mid    = (dimLength + 1) / 2;
     af_array left = 0;
 
     af_seq slices[4] = {af_span, af_span, af_span, af_span};
@@ -78,7 +78,7 @@ static af_array median(const af_array& in, const dim_t dim)
     af_array sortedIn_handle = getHandle<T>(sortedIn);
     AF_CHECK(af_index(&left, sortedIn_handle, input.ndims(), slices));
 
-    if (nElems % 2 == 1) {
+    if (dimLength % 2 == 1) {
         // mid-1 is our guy
         if (input.isFloating()) return left;
 
@@ -90,7 +90,7 @@ static af_array median(const af_array& in, const dim_t dim)
         return out;
     } else {
         // ((mid-1)+mid)/2 is our guy
-        dim4  dims = input.dims();
+        dim4 dims = input.dims();
         af_array right = 0;
         slices[dim] = af_make_seq(mid, mid, 1.0);
 
@@ -100,7 +100,8 @@ static af_array median(const af_array& in, const dim_t dim)
         af_array carr   = 0;
         af_array result = 0;
 
-        dim4 cdims = dim4(1, dims[1], dims[2], dims[3]);
+        dim4 cdims = dims;
+        cdims[dim] = 1;
         AF_CHECK(af_constant(&carr, 0.5, cdims.ndims(), cdims.get(), input.isDouble() ? f64 : f32));
 
         if (!input.isFloating()) {
@@ -148,7 +149,7 @@ af_err af_median_all(double *realVal, double *imagVal, const af_array in)
 af_err af_median(af_array* out, const af_array in, const dim_t dim)
 {
     try {
-        ARG_ASSERT(2, (dim>=0 && dim<=0));
+        ARG_ASSERT(2, (dim >= 0 && dim <= 4));
 
         af_array output = 0;
         ArrayInfo info = getInfo(in);

diff --git a/src/api/c/moddims.cpp b/src/api/c/moddims.cpp
@@ -31,7 +31,6 @@ Array<T> modDims(const Array<T>& in, const af::dim4 &newDims)
         Out = copyArray<T>(in);
     }
 
-    Out.modDims(newDims);
     Out.setDataDims(newDims);
 
     return Out;

diff --git a/src/api/c/sort.cpp b/src/api/c/sort.cpp
@@ -42,8 +42,6 @@ af_err af_sort(af_array *out, const af_array in, const unsigned dim, const bool
         af_dtype type = info.getType();
 
         DIM_ASSERT(1, info.elements() > 0);
-        // Only Dim 0 supported
-        ARG_ASSERT(2, dim == 0);
 
         af_array val;
 
@@ -93,8 +91,6 @@ af_err af_sort_index(af_array *out, af_array *indices, const af_array in, const
         af_dtype type = info.getType();
 
         DIM_ASSERT(2, info.elements() > 0);
-        // Only Dim 0 supported
-        ARG_ASSERT(3, dim == 0);
 
         af_array val;
         af_array idx;
@@ -150,6 +146,8 @@ void sort_by_key_tmplt(af_array *okey, af_array *oval, const af_array ikey, cons
     switch(vtype) {
     case f32: sort_by_key<Tk, float  >(okey, oval, ikey, ival, dim, isAscending);  break;
     case f64: sort_by_key<Tk, double >(okey, oval, ikey, ival, dim, isAscending);  break;
+    case c32: sort_by_key<Tk, cfloat >(okey, oval, ikey, ival, dim, isAscending);  break;
+    case c64: sort_by_key<Tk, cdouble>(okey, oval, ikey, ival, dim, isAscending);  break;
     case s32: sort_by_key<Tk, int    >(okey, oval, ikey, ival, dim, isAscending);  break;
     case u32: sort_by_key<Tk, uint   >(okey, oval, ikey, ival, dim, isAscending);  break;
     case s16: sort_by_key<Tk, short  >(okey, oval, ikey, ival, dim, isAscending);  break;
@@ -169,20 +167,20 @@ af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
                       const unsigned dim, const bool isAscending)
 {
     try {
-        ArrayInfo info = getInfo(keys);
-        af_dtype type = info.getType();
+        ArrayInfo kinfo = getInfo(keys);
+        af_dtype ktype = kinfo.getType();
 
         ArrayInfo vinfo = getInfo(values);
 
-        DIM_ASSERT(3, info.elements() > 0);
-        DIM_ASSERT(4, info.dims() == vinfo.dims());
-        // Only Dim 0 supported
-        ARG_ASSERT(5, dim == 0);
+        DIM_ASSERT(3, kinfo.elements() > 0);
+        DIM_ASSERT(4, kinfo.dims() == vinfo.dims());
+
+        TYPE_ASSERT(kinfo.isReal());
 
         af_array oKey;
         af_array oVal;
 
-        switch(type) {
+        switch(ktype) {
             case f32: sort_by_key_tmplt<float  >(&oKey, &oVal, keys, values, dim, isAscending);  break;
             case f64: sort_by_key_tmplt<double >(&oKey, &oVal, keys, values, dim, isAscending);  break;
             case s32: sort_by_key_tmplt<int    >(&oKey, &oVal, keys, values, dim, isAscending);  break;
@@ -193,7 +191,7 @@ af_err af_sort_by_key(af_array *out_keys, af_array *out_values,
             case u64: sort_by_key_tmplt<uintl  >(&oKey, &oVal, keys, values, dim, isAscending);  break;
             case u8:  sort_by_key_tmplt<uchar  >(&oKey, &oVal, keys, values, dim, isAscending);  break;
             case b8:  sort_by_key_tmplt<char   >(&oKey, &oVal, keys, values, dim, isAscending);  break;
-            default:  TYPE_ERROR(1, type);
+            default:  TYPE_ERROR(1, ktype);
         }
         std::swap(*out_keys , oKey);
         std::swap(*out_values , oVal);

diff --git a/src/backend/cpu/Array.hpp b/src/backend/cpu/Array.hpp
@@ -181,6 +181,7 @@ namespace cpu
 
         void setDataDims(const dim4 &new_dims)
         {
+            modDims(new_dims);
             data_dims = new_dims;
         }
 

diff --git a/src/backend/cpu/CMakeLists.txt b/src/backend/cpu/CMakeLists.txt
@@ -93,7 +93,8 @@ FILE(GLOB cpu_headers
     "*.h")
 
 FILE(GLOB cpu_sources
-    "*.cpp")
+    "*.cpp"
+    "kernel/sort_by_key/*.cpp")
 
 LIST(SORT cpu_headers)
 LIST(SORT cpu_sources)

diff --git a/src/backend/cpu/kernel/sort.hpp b/src/backend/cpu/kernel/sort.hpp
@@ -23,7 +23,7 @@ namespace kernel
 
 // Based off of http://stackoverflow.com/a/12399290
 template<typename T, bool isAscending>
-void sort0(Array<T> val)
+void sort0Iterative(Array<T> val)
 {
     // initialize original index locations
     T *val_ptr = val.get();

diff --git a/src/backend/cpu/kernel/sort_by_key.hpp b/src/backend/cpu/kernel/sort_by_key.hpp
@@ -8,79 +8,22 @@
  ********************************************************/
 
 #pragma once
-#include <af/defines.h>
 #include <Array.hpp>
-#include <math.hpp>
-#include <algorithm>
-#include <numeric>
-#include <queue>
 #include <err_cpu.hpp>
-#include <functional>
 
 namespace cpu
 {
 namespace kernel
 {
 
 template<typename Tk, typename Tv, bool isAscending>
-void sort0_by_key(Array<Tk> okey, Array<Tv> oval, Array<uint> oidx,
-                  const Array<Tk> ikey, const Array<Tv> ival)
-{
-    function<bool(Tk, Tk)> op = std::greater<Tk>();
-    if(isAscending) { op = std::less<Tk>(); }
-
-    // Get pointers and initialize original index locations
-        uint *oidx_ptr = oidx.get();
-          Tk *okey_ptr = okey.get();
-          Tv *oval_ptr = oval.get();
-    const Tk *ikey_ptr = ikey.get();
-    const Tv *ival_ptr = ival.get();
-
-    std::vector<uint> seq_vec(oidx.dims()[0]);
-    std::iota(seq_vec.begin(), seq_vec.end(), 0);
-
-    const Tk *comp_ptr = nullptr;
-    auto comparator = [&comp_ptr, &op](size_t i1, size_t i2) {return op(comp_ptr[i1], comp_ptr[i2]);};
-
-    for(dim_t w = 0; w < ikey.dims()[3]; w++) {
-        dim_t okeyW = w * okey.strides()[3];
-        dim_t ovalW = w * oval.strides()[3];
-        dim_t oidxW = w * oidx.strides()[3];
-        dim_t ikeyW = w * ikey.strides()[3];
-        dim_t ivalW = w * ival.strides()[3];
-
-        for(dim_t z = 0; z < ikey.dims()[2]; z++) {
-            dim_t okeyWZ = okeyW + z * okey.strides()[2];
-            dim_t ovalWZ = ovalW + z * oval.strides()[2];
-            dim_t oidxWZ = oidxW + z * oidx.strides()[2];
-            dim_t ikeyWZ = ikeyW + z * ikey.strides()[2];
-            dim_t ivalWZ = ivalW + z * ival.strides()[2];
+void sort0ByKeyIterative(Array<Tk> okey, Array<Tv> oval);
 
-            for(dim_t y = 0; y < ikey.dims()[1]; y++) {
+template<typename Tk, typename Tv, bool isAscending, int dim>
+void sortByKeyBatched(Array<Tk> okey, Array<Tv> oval);
 
-                dim_t okeyOffset = okeyWZ + y * okey.strides()[1];
-                dim_t ovalOffset = ovalWZ + y * oval.strides()[1];
-                dim_t oidxOffset = oidxWZ + y * oidx.strides()[1];
-                dim_t ikeyOffset = ikeyWZ + y * ikey.strides()[1];
-                dim_t ivalOffset = ivalWZ + y * ival.strides()[1];
-
-                uint *ptr = oidx_ptr + oidxOffset;
-                std::copy(seq_vec.begin(), seq_vec.end(), ptr);
-
-                comp_ptr = ikey_ptr + ikeyOffset;
-                std::stable_sort(ptr, ptr + ikey.dims()[0], comparator);
-
-                for (dim_t i = 0; i < oval.dims()[0]; ++i){
-                    uint sortIdx = oidx_ptr[oidxOffset + i];
-                    okey_ptr[okeyOffset + i] = ikey_ptr[ikeyOffset + sortIdx];
-                    oval_ptr[ovalOffset + i] = ival_ptr[ivalOffset + sortIdx];
-                }
-            }
-        }
-    }
-
-    return;
-}
+template<typename Tk, typename Tv, bool isAscending>
+void sort0ByKey(Array<Tk> okey, Array<Tv> oval);
 
 }
 }
diff --git a/src/backend/cpu/kernel/sort_by_key/b8.cpp b/src/backend/cpu/kernel/sort_by_key/b8.cpp
@@ -0,0 +1,19 @@
+/*******************************************************
+ * Copyright (c) 2014, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <kernel/sort_by_key_impl.hpp>
+
+namespace cpu
+{
+namespace kernel
+{
+    INSTANTIATE1(char,true)
+    INSTANTIATE1(char,false)
+}
+}
diff --git a/src/backend/cpu/kernel/sort_by_key/f32.cpp b/src/backend/cpu/kernel/sort_by_key/f32.cpp
@@ -0,0 +1,19 @@
+/*******************************************************
+ * Copyright (c) 2014, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <kernel/sort_by_key_impl.hpp>
+
+namespace cpu
+{
+namespace kernel
+{
+    INSTANTIATE1(float,true)
+    INSTANTIATE1(float,false)
+}
+}
diff --git a/src/backend/cpu/kernel/sort_by_key/f64.cpp b/src/backend/cpu/kernel/sort_by_key/f64.cpp
@@ -0,0 +1,19 @@
+/*******************************************************
+ * Copyright (c) 2014, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <kernel/sort_by_key_impl.hpp>
+
+namespace cpu
+{
+namespace kernel
+{
+    INSTANTIATE1(double,true)
+    INSTANTIATE1(double,false)
+}
+}
diff --git a/src/backend/cpu/kernel/sort_by_key/s16.cpp b/src/backend/cpu/kernel/sort_by_key/s16.cpp
@@ -0,0 +1,19 @@
+/*******************************************************
+ * Copyright (c) 2014, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <kernel/sort_by_key_impl.hpp>
+
+namespace cpu
+{
+namespace kernel
+{
+    INSTANTIATE1(short,true)
+    INSTANTIATE1(short,false)
+}
+}
diff --git a/src/backend/cpu/kernel/sort_by_key/s32.cpp b/src/backend/cpu/kernel/sort_by_key/s32.cpp
@@ -0,0 +1,19 @@
+/*******************************************************
+ * Copyright (c) 2014, ArrayFire
+ * All rights reserved.
+ *
+ * This file is distributed under 3-clause BSD license.
+ * The complete license agreement can be obtained at:
+ * http://arrayfire.com/licenses/BSD-3-Clause
+ ********************************************************/
+
+#include <kernel/sort_by_key_impl.hpp>
+
+namespace cpu
+{
+namespace kernel
+{
+    INSTANTIATE1(int,true)
+    INSTANTIATE1(int,false)
+}
+}
-Original file line number
+Diff line change
@@ Expand Up / @@ -181,6 +181,7 @@ namespace cpu @@
             void setDataDims(const dim4 &new_dims)
             {
+                modDims(new_dims);
                 data_dims = new_dims;
             }
@@ Expand Down @@