FEAT: Add CSR <--> COO Conversion in all backends

src/api/c/sparse.cpp

@@ -183,10 +183,10 @@ af_array createSparseArrayFromDense(
     switch(stype) {
         case AF_STORAGE_CSR:
             return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_CSR>(in));
-        case AF_STORAGE_CSC:
-            return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_CSC>(in));
         case AF_STORAGE_COO:
             return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_COO>(in));
+        case AF_STORAGE_CSC:
+            //return getHandle(sparseConvertDenseToStorage<T, AF_STORAGE_CSC>(in));
         default: AF_ERROR("Storage type is out of range/unsupported", AF_ERR_ARG);
     }
 }
@@ -235,16 +235,11 @@ af_array sparseConvertStorage(const af_array in_, const af_storage destStorage)
 {
     const SparseArray<T> in = getSparseArray<T>(in_);
 
-    // Only destStorage == AF_STORAGE_DENSE is supported
-    // All the other calls are for future when conversions are supported in
-    // the backend
     if(destStorage == AF_STORAGE_DENSE) {
         // Returns a regular af_array, not sparse
         switch(in.getStorage()) {
             case AF_STORAGE_CSR:
                 return getHandle(detail::sparseConvertStorageToDense<T, AF_STORAGE_CSR>(in));
-            case AF_STORAGE_CSC:
-                return getHandle(detail::sparseConvertStorageToDense<T, AF_STORAGE_CSC>(in));
             case AF_STORAGE_COO:
                 return getHandle(detail::sparseConvertStorageToDense<T, AF_STORAGE_COO>(in));
             default:
@@ -255,32 +250,16 @@ af_array sparseConvertStorage(const af_array in_, const af_storage destStorage)
         switch(in.getStorage()) {
             case AF_STORAGE_CSR:
                 return retainSparseHandle<T>(in_);
-            case AF_STORAGE_CSC:
-                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_CSR, AF_STORAGE_CSC>(in));
             case AF_STORAGE_COO:
                 return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_CSR, AF_STORAGE_COO>(in));
             default:
                 AF_ERROR("Invalid storage type of input array", AF_ERR_ARG);
         }
-    } else if(destStorage == AF_STORAGE_CSC) {
-        // Returns a sparse af_array
-        switch(in.getStorage()) {
-            case AF_STORAGE_CSR:
-                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_CSC, AF_STORAGE_CSR>(in));
-            case AF_STORAGE_CSC:
-                return retainSparseHandle<T>(in_);
-            case AF_STORAGE_COO:
-                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_CSC, AF_STORAGE_COO>(in));
-            default:
-                AF_ERROR("Invalid storage type of input array", AF_ERR_ARG);
-        }
     } else if(destStorage == AF_STORAGE_COO) {
         // Returns a sparse af_array
         switch(in.getStorage()) {
             case AF_STORAGE_CSR:
                 return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_COO, AF_STORAGE_CSR>(in));
-            case AF_STORAGE_CSC:
-                return getHandle(detail::sparseConvertStorageToStorage<T, AF_STORAGE_COO, AF_STORAGE_CSC>(in));
             case AF_STORAGE_COO:
                 return retainSparseHandle<T>(in_);
             default:
@@ -295,19 +274,24 @@ af_array sparseConvertStorage(const af_array in_, const af_storage destStorage)
 af_err af_sparse_convert_to(af_array *out, const af_array in,
                             const af_storage destStorage)
 {
-    // Right now dest_storage can only be AF_STORAGE_DENSE
     try {
+        // Handle dense case
+        const ArrayInfo& info = getInfo(in, false, true);
+        if(!info.isSparse()) {    // If input is dense
+            return af_create_sparse_array_from_dense(out, in, destStorage);
+        }
+
         af_array output = 0;
 
         const SparseArrayBase base = getSparseArrayBase(in);
 
-        // Dense not allowed as input -> Should never happen
-        // To convert from dense to type, use the create* functions
-        ARG_ASSERT(1, base.getStorage() != AF_STORAGE_DENSE);
+        // Dense not allowed as input -> Should never happen with SparseArrayBase
+        // CSC is currently not supported
+        ARG_ASSERT(1, base.getStorage() != AF_STORAGE_DENSE
+                   && base.getStorage() != AF_STORAGE_CSC);
 
-        // Right now dest_storage can only be AF_STORAGE_DENSE
-        // TODO: Add support for [CSR, CSC, COO] <-> [CSR, CSC, COO] in backends
-        ARG_ASSERT(1, destStorage == AF_STORAGE_DENSE);
+        // Conversion to and from CSC is not supported
+        ARG_ASSERT(2, destStorage != AF_STORAGE_CSC);
 
         if(base.getStorage() == destStorage) {
             // Return a reference

src/backend/SparseArray.cpp

@@ -124,19 +124,19 @@ SparseArray<T> createDeviceDataSparseArray(
         const af::dim4 &_dims, const dim_t nNZ,
         const T * const _values,
         const int * const _rowIdx, const int * const _colIdx,
-        const af::storage _storage)
+        const af::storage _storage, const bool _copy)
 {
-    return SparseArray<T>(_dims, nNZ, _values, _rowIdx, _colIdx, _storage, false);
+    return SparseArray<T>(_dims, nNZ, _values, _rowIdx, _colIdx, _storage, _copy);
 }
 
 template<typename T>
 SparseArray<T> createArrayDataSparseArray(
         const af::dim4 &_dims,
         const Array<T> &_values,
         const Array<int> &_rowIdx, const Array<int> &_colIdx,
-        const af::storage _storage)
+        const af::storage _storage, const bool _copy)
 {
-    return SparseArray<T>(_dims, _values, _rowIdx, _colIdx, _storage, false);
+    return SparseArray<T>(_dims, _values, _rowIdx, _colIdx, _storage, _copy);
 }
 
 template<typename T>
@@ -223,12 +223,12 @@ SparseArray<T>::~SparseArray()
             const af::dim4 &_dims, const dim_t _nNZ,                                                \
             const T * const _values,                                                                \
             const int * const _rowIdx, const int * const _colIdx,                                   \
-            const af::storage _storage);                                                            \
+            const af::storage _storage, const bool _copy);                                          \
     template SparseArray<T> createArrayDataSparseArray<T>(                                          \
             const af::dim4 &_dims,                                                                  \
             const Array<T> &_values,                                                                \
             const Array<int> &_rowIdx, const Array<int> &_colIdx,                                   \
-            const af::storage _storage);                                                            \
+            const af::storage _storage, const bool _copy);                                          \
     template SparseArray<T> *initSparseArray<T>();                                                  \
     template void destroySparseArray<T>(SparseArray<T> *sparse);                                    \
                                                                                                     \

src/backend/SparseArray.hpp

@@ -218,13 +218,13 @@ class SparseArray
             const af::dim4 &_dims, const dim_t nNZ,
             const T * const _values,
             const int * const _rowIdx, const int * const _colIdx,
-            const af::storage _storage);
+            const af::storage _storage, const bool _copy);
 
     friend SparseArray<T> createArrayDataSparseArray<T>(
             const af::dim4 &_dims,
             const Array<T> &_values,
             const Array<int> &_rowIdx, const Array<int> &_colIdx,
-            const af::storage _storage);
+            const af::storage _storage, const bool _copy);
 
     friend SparseArray<T> *initSparseArray<T>();
 

src/backend/cpu/kernel/sparse.hpp

@@ -11,6 +11,9 @@
 #include <Array.hpp>
 #include <utility.hpp>
 #include <math.hpp>
+#include <kernel/sort_helper.hpp>
+#include <algorithm>
+#include <tuple>
 
 namespace cpu
 {
@@ -93,5 +96,99 @@ struct csr_dense
     }
 };
 
+// Modified code from sort helper
+template <typename T>
+using SpKeyIndexPair = std::tuple<int, T, int>; // sorting index, value, other index
+
+template <typename T>
+struct SpKIPCompareK
+{
+    bool operator()(const SpKeyIndexPair<T> &lhs, const SpKeyIndexPair<T> &rhs)
+    {
+        int lhsVal = std::get<0>(lhs);
+        int rhsVal = std::get<0>(rhs);
+        // Always returns ascending
+        return (lhsVal < rhsVal);
+    }
+};
+
+template<typename T>
+void csr_coo(Array<T> ovalues, Array<int> orowIdx, Array<int> ocolIdx,
+             Array<T> const ivalues, Array<int> const irowIdx, Array<int> const icolIdx)
+{
+    // First calculate the linear index
+    T         * const ovPtr = ovalues.get();
+    int       * const orPtr = orowIdx.get();
+    int       * const ocPtr = ocolIdx.get();
+
+    T   const * const ivPtr = ivalues.get();
+    int const * const irPtr = irowIdx.get();
+    int const * const icPtr = icolIdx.get();
+
+    // Create cordinate form of the row array
+    for(int i = 0; i < (int)irowIdx.elements() - 1; i++) {
+        std::fill_n(orPtr + irPtr[i], irPtr[i + 1] - irPtr[i], i);
+    }
+
+    // Sort the coordinate form using column index
+    // Uses code from sort_by_key kernels
+    typedef SpKeyIndexPair<T> CurrentPair;
+    int size = ovalues.dims()[0];
+    size_t bytes = size * sizeof(CurrentPair);
+    CurrentPair *pairKeyVal = (CurrentPair *)memAlloc<char>(bytes);
+
+    for(int x = 0; x < size; x++) {
+       pairKeyVal[x] = std::make_tuple(icPtr[x], ivPtr[x], orPtr[x]);
+    }
+
+    std::stable_sort(pairKeyVal, pairKeyVal + size, SpKIPCompareK<T>());
+
+    for(int x = 0; x < (int)ovalues.elements(); x++) {
+        std::tie(ocPtr[x], ovPtr[x], orPtr[x]) = pairKeyVal[x];
+    }
+
+    memFree((char *)pairKeyVal);
+}
+
+template<typename T>
+void coo_csr(Array<T> ovalues, Array<int> orowIdx, Array<int> ocolIdx,
+             Array<T> const ivalues, Array<int> const irowIdx, Array<int> const icolIdx)
+{
+    T         * const ovPtr = ovalues.get();
+    int       * const orPtr = orowIdx.get();
+    int       * const ocPtr = ocolIdx.get();
+
+    T   const * const ivPtr = ivalues.get();
+    int const * const irPtr = irowIdx.get();
+    int const * const icPtr = icolIdx.get();
+
+    // Sort the colidx and values based on rowIdx
+    // Uses code from sort_by_key kernels
+    typedef SpKeyIndexPair<T> CurrentPair;
+    int size = ovalues.dims()[0];
+    size_t bytes = size * sizeof(CurrentPair);
+    CurrentPair *pairKeyVal = (CurrentPair *)memAlloc<char>(bytes);
+
+    for(int x = 0; x < size; x++) {
+       pairKeyVal[x] = std::make_tuple(irPtr[x], ivPtr[x], icPtr[x]);
+    }
+
+    std::stable_sort(pairKeyVal, pairKeyVal + size, SpKIPCompareK<T>());
+
+    ovPtr[0] = 0;
+    for(int x = 0; x < (int)ovalues.elements(); x++) {
+        int row = -2; // Some value that will make orPtr[row + 1] error out
+        std::tie(row, ovPtr[x], ocPtr[x]) = pairKeyVal[x];
+        orPtr[row + 1]++;
+    }
+
+    // Compress row storage
+    for(int x = 1; x < (int)orowIdx.elements(); x++) {
+        orPtr[x] += orPtr[x - 1];
+    }
+
+    memFree((char *)pairKeyVal);
+}
+
 }
 }

src/backend/cpu/sparse.cpp

@@ -119,7 +119,7 @@ SPARSE_FUNC(csrcsc, cdouble,z)
 #endif // USE_MKL
 
 ////////////////////////////////////////////////////////////////////////////////
-// Common to MKL and Not MKL
+// Common Funcs for MKL and Non-MKL Code Paths
 ////////////////////////////////////////////////////////////////////////////////
 
 // Partial template specialization of sparseConvertDenseToStorage for COO
@@ -134,11 +134,11 @@ SparseArray<T> sparseConvertDenseToCOO(const Array<T> &in)
 
     dim_t nNZ = nonZeroIdx.elements();
 
-    Array<int> constNNZ = createValueArray<int>(dim4(nNZ), nNZ);
-    constNNZ.eval();
+    Array<int> constDim = createValueArray<int>(dim4(nNZ), in.dims()[0]);
+    constDim.eval();
 
-    Array<int> rowIdx = arithOp<int, af_mod_t>(nonZeroIdx, constNNZ, nonZeroIdx.dims());
-    Array<int> colIdx = arithOp<int, af_div_t>(nonZeroIdx, constNNZ, nonZeroIdx.dims());
+    Array<int> rowIdx = arithOp<int, af_mod_t>(nonZeroIdx, constDim, nonZeroIdx.dims());
+    Array<int> colIdx = arithOp<int, af_div_t>(nonZeroIdx, constDim, nonZeroIdx.dims());
 
     Array<T> values = copyArray<T>(in);
     values.modDims(dim4(values.elements()));
@@ -337,20 +337,34 @@ Array<T> sparseConvertStorageToDense(const SparseArray<T> &in_)
 ////////////////////////////////////////////////////////////////////////////////
 
 ////////////////////////////////////////////////////////////////////////////////
-// Common to MKL and Not MKL
+// Common Funcs for MKL and Non-MKL Code Paths
 ////////////////////////////////////////////////////////////////////////////////
-template<typename T, af_storage src, af_storage dest>
+template<typename T, af_storage dest, af_storage src>
 SparseArray<T> sparseConvertStorageToStorage(const SparseArray<T> &in)
 {
-    // Dummy function
-    // TODO finish this function when support is required
-    AF_ERROR("CPU Backend only supports Dense to CSR or COO", AF_ERR_NOT_SUPPORTED);
-
     in.eval();
 
-    SparseArray<T> dense = createEmptySparseArray<T>(in.dims(), (int)in.getNNZ(), dest);
+    SparseArray<T> converted = createEmptySparseArray<T>(in.dims(), (int)in.getNNZ(), dest);
+    converted.eval();
 
-    return dense;
+    function<void (Array<T>, Array<int>, Array<int>,
+                   Array<T> const, Array<int> const, Array<int> const)
+            > converter;
+
+    if(src == AF_STORAGE_CSR && dest == AF_STORAGE_COO) {
+        converter = kernel::csr_coo<T>;
+    } else if (src == AF_STORAGE_COO && dest == AF_STORAGE_CSR) {
+        converter = kernel::coo_csr<T>;
+    } else {
+        // Should never come here
+        AF_ERROR("CPU Backend invalid conversion combination", AF_ERR_NOT_SUPPORTED);
+    }
+
+    getQueue().enqueue(converter,
+                       converted.getValues(), converted.getRowIdx(), converted.getColIdx(),
+                       in.getValues(), in.getRowIdx(), in.getColIdx());
+
+    return converted;
 }
 
 #define INSTANTIATE_TO_STORAGE(T, S)                                                                        \