Radix returns sorted keys

Author: kirillgarbar

src/GraphBLAS-sharp.Backend/Common/Sort/Radix.fs

@@ -268,7 +268,7 @@ module Radix =
                 failwith "Mismatch of key lengths and value. Lengths must be the same"
 
             if values.Length <= 1 then
-                dataCopy processor allocationMode values
+                copy processor DeviceOnly keys, dataCopy processor DeviceOnly values
             else
                 let firstKeys = copy processor DeviceOnly keys
 
@@ -321,11 +321,10 @@ module Radix =
                     localOffset.Free processor
                     shift.Free processor
 
-                (fst keysPair).Free processor
                 (snd keysPair).Free processor
                 (snd valuesPair).Free processor
 
-                (fst valuesPair)
+                (fst keysPair, fst valuesPair)
 
     let runByKeysStandard clContext workGroupSize =
         runByKeys clContext workGroupSize defaultBitCount

src/GraphBLAS-sharp.Backend/Matrix/SpGeMM/Expand.fs

@@ -181,21 +181,31 @@ module Expand =
 
         fun (processor: MailboxProcessor<_>) (values: ClArray<'a>) (columns: ClArray<int>) (rows: ClArray<int>) ->
             // sort by columns
-            let valuesSortedByColumns =
+            let keys, valuesSortedByColumns =
                 sortByKeyValues processor DeviceOnly columns values
 
-            let rowsSortedByColumns =
+            keys.Free processor
+
+            let keys, rowsSortedByColumns =
                 sortByKeyIndices processor DeviceOnly columns rows
 
+            keys.Free processor
+
+            //TODO: already sorted above?
             let sortedColumns = sortKeys processor columns
 
             // sort by rows
-            let valuesSortedByRows =
+            let keys, valuesSortedByRows =
                 sortByKeyValues processor DeviceOnly rowsSortedByColumns valuesSortedByColumns
 
-            let columnsSortedByRows =
+            keys.Free processor
+
+            let keys, columnsSortedByRows =
                 sortByKeyIndices processor DeviceOnly rowsSortedByColumns sortedColumns
 
+            keys.Free processor
+
+            //TODO: already sorted above?
             let sortedRows = sortKeys processor rowsSortedByColumns
 
             valuesSortedByColumns.Free processor

tests/GraphBLAS-sharp.Tests/Backend/Common/Sort/Radix.fs

@@ -17,8 +17,13 @@ let processor = Context.defaultContext.Queue
 
 let context = Context.defaultContext.ClContext
 
-let checkResultByKeys (inputArray: (int * 'a) []) (actualValues: 'a []) =
-    let expectedValues = Seq.sortBy fst inputArray |> Seq.map snd
+let checkResultByKeys (inputArray: (int * 'a) []) (actualKeys: int []) (actualValues: 'a []) =
+    let expected = Seq.sortBy fst inputArray
+    let expectedKeys = expected |> Seq.map fst
+    let expectedValues = expected |> Seq.map snd
+
+    "Keys must be the same"
+    |> Expect.sequenceEqual expectedKeys actualKeys
 
     "Values must be the same"
     |> Expect.sequenceEqual expectedValues actualValues
@@ -32,12 +37,13 @@ let makeTestByKeys<'a when 'a: equality> sortFun (array: (int * 'a) []) =
         let clKeys = keys.ToDevice context
         let clValues = values.ToDevice context
 
-        let clActualValues: ClArray<'a> =
+        let clActualKeys, clActualValues: ClArray<int> * ClArray<'a> =
             sortFun processor HostInterop clKeys clValues
 
+        let actualKeys = clActualKeys.ToHostAndFree processor
         let actualValues = clActualValues.ToHostAndFree processor
 
-        checkResultByKeys array actualValues
+        checkResultByKeys array actualKeys actualValues
 
 let createTestByKeys<'a when 'a: equality and 'a: struct> =
     let sort =