Skip to content

Commit 6da457b

Browse files
IgorErinIgorErin
committed
add: reduceByKey2D
1 parent 890bdd1 commit 6da457b

5 files changed

Lines changed: 366 additions & 12 deletions

File tree

‎src/GraphBLAS-sharp.Backend/Common/Sum.fs‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,3 +470,149 @@ module Reduce =
470470
processor.Post(Msg.CreateRunMsg<_, _>(kernel))
471471

472472
reducedKeys, reducedValues
473+
474+
module ByKey2D =
475+
/// <summary>
476+
/// Reduce an array of values by 2D keys using a single work item.
477+
/// </summary>
478+
/// <param name="clContext">ClContext.</param>
479+
/// <param name="workGroupSize">Work group size.</param>
480+
/// <param name="reduceOp">Operation for reducing values.</param>
481+
/// <remarks>
482+
/// The length of the result must be calculated in advance.
483+
/// </remarks>
484+
let sequential (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a>) =
485+
486+
let kernel =
487+
<@ fun (ndRange: Range1D) length (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (firstReducedKeys: ClArray<int>) (secondReducedKeys: ClArray<int>) ->
488+
489+
let gid = ndRange.GlobalID0
490+
491+
if gid = 0 then
492+
let mutable firstCurrentKey = firstKeys.[0]
493+
let mutable secondCurrentKey = secondKeys.[0]
494+
495+
let mutable segmentResult = values.[0]
496+
let mutable segmentCount = 0
497+
498+
for i in 1 .. length - 1 do
499+
if firstCurrentKey = firstKeys.[i]
500+
&& secondCurrentKey = secondKeys.[i] then
501+
segmentResult <- (%reduceOp) segmentResult values.[i]
502+
else
503+
reducedValues.[segmentCount] <- segmentResult
504+
505+
firstReducedKeys.[segmentCount] <- firstCurrentKey
506+
secondReducedKeys.[segmentCount] <- secondCurrentKey
507+
508+
segmentCount <- segmentCount + 1
509+
firstCurrentKey <- firstKeys.[i]
510+
secondCurrentKey <- secondKeys.[i]
511+
segmentResult <- values.[i]
512+
513+
firstReducedKeys.[segmentCount] <- firstCurrentKey
514+
secondReducedKeys.[segmentCount] <- secondCurrentKey
515+
516+
reducedValues.[segmentCount] <- segmentResult @>
517+
518+
let kernel = clContext.Compile kernel
519+
520+
fun (processor: MailboxProcessor<_>) allocationMode (resultLength: int) (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) ->
521+
522+
let reducedValues =
523+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
524+
525+
let firstReducedKeys =
526+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
527+
528+
let secondReducedKeys =
529+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
530+
531+
let ndRange =
532+
Range1D.CreateValid(resultLength, workGroupSize)
533+
534+
let kernel = kernel.GetKernel()
535+
536+
processor.Post(
537+
Msg.MsgSetArguments
538+
(fun () -> kernel.KernelFunc ndRange firstKeys.Length firstKeys secondKeys values reducedValues firstReducedKeys secondReducedKeys)
539+
)
540+
541+
processor.Post(Msg.CreateRunMsg<_, _>(kernel))
542+
543+
firstReducedKeys, secondReducedKeys, reducedValues
544+
545+
/// <summary>
546+
/// Reduces values by key. Each segment is reduced by one work item.
547+
/// </summary>
548+
/// <param name="clContext">ClContext.</param>
549+
/// <param name="workGroupSize">Work group size.</param>
550+
/// <param name="reduceOp">Operation for reducing values.</param>
551+
/// <remarks>
552+
/// The length of the result must be calculated in advance.
553+
/// </remarks>
554+
let segmentSequential<'a> (clContext: ClContext) workGroupSize (reduceOp: Expr<'a -> 'a -> 'a>) =
555+
556+
let kernel =
557+
<@ fun (ndRange: Range1D) uniqueKeyCount keysLength (offsets: ClArray<int>) (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) (reducedValues: ClArray<'a>) (firstReducedKeys: ClArray<int>) (secondReducedKeys: ClArray<int>) ->
558+
559+
let gid = ndRange.GlobalID0
560+
561+
if gid < uniqueKeyCount then
562+
let startPosition = offsets.[gid]
563+
564+
let firstSourceKey = firstKeys.[startPosition]
565+
let secondSourceKey = secondKeys.[startPosition]
566+
567+
let mutable sum = values.[startPosition]
568+
569+
let mutable currentPosition = startPosition + 1
570+
571+
while currentPosition < keysLength
572+
&& firstSourceKey = firstKeys.[currentPosition]
573+
&& secondSourceKey = secondKeys.[currentPosition] do
574+
575+
sum <- (%reduceOp) sum values.[currentPosition]
576+
currentPosition <- currentPosition + 1
577+
578+
reducedValues.[gid] <- sum
579+
firstReducedKeys.[gid] <- firstSourceKey
580+
secondReducedKeys.[gid] <- secondSourceKey @>
581+
582+
let kernel = clContext.Compile kernel
583+
584+
fun (processor: MailboxProcessor<_>) allocationMode (resultLength: int) (offsets: ClArray<int>) (firstKeys: ClArray<int>) (secondKeys: ClArray<int>) (values: ClArray<'a>) ->
585+
586+
let reducedValues =
587+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
588+
589+
let firstReducedKeys =
590+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
591+
592+
let secondReducedKeys =
593+
clContext.CreateClArrayWithSpecificAllocationMode(allocationMode, resultLength)
594+
595+
let ndRange =
596+
Range1D.CreateValid(resultLength, workGroupSize)
597+
598+
let kernel = kernel.GetKernel()
599+
600+
processor.Post(
601+
Msg.MsgSetArguments
602+
(fun () ->
603+
kernel.KernelFunc
604+
ndRange
605+
resultLength
606+
firstKeys.Length
607+
offsets
608+
firstKeys
609+
secondKeys
610+
values
611+
reducedValues
612+
firstReducedKeys
613+
secondReducedKeys)
614+
)
615+
616+
processor.Post(Msg.CreateRunMsg<_, _>(kernel))
617+
618+
firstReducedKeys, secondReducedKeys, reducedValues

‎src/GraphBLAS-sharp.Backend/Matrix/CSRMatrix/Matrix.fs‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ open Microsoft.FSharp.Quotations
88
open GraphBLAS.FSharp.Backend.Objects
99
open GraphBLAS.FSharp.Backend.Objects.ClMatrix
1010
open GraphBLAS.FSharp.Backend.Objects.ClContext
11+
open GraphBLAS.FSharp.Backend.Objects.ArraysExtensions
1112

1213
module Matrix =
1314
let private expandRowPointers (clContext: ClContext) workGroupSize =
@@ -153,3 +154,51 @@ module Matrix =
153154
fun (queue: MailboxProcessor<_>) (matrixLeft: ClMatrix.CSR<'a>) (matrixRight: ClMatrix.CSC<'b>) (mask: ClMatrix.COO<_>) ->
154155

155156
run queue matrixLeft matrixRight mask
157+
158+
let spgemm
159+
(clContext: ClContext)
160+
workGroupSize
161+
(opAdd: Expr<'c -> 'c -> 'c>)
162+
(opMul: Expr<'a -> 'b -> 'c>)
163+
=
164+
165+
let expand = SpGEMM.Expand.run clContext workGroupSize opMul
166+
167+
let expandRowPointers = expandRowPointers clContext workGroupSize
168+
169+
let sortData = Sort.Radix.runByKeysStandard clContext workGroupSize
170+
171+
let sortKeys = Sort.Radix.runByKeysStandard clContext workGroupSize
172+
173+
let reduceByKey = Reduce.ByKey.segmentSequential clContext workGroupSize opAdd
174+
175+
fun (processor: MailboxProcessor<_>) (leftMatrix: ClMatrix.CSR<'a>) (rightMatrix: ClMatrix.CSR<'b>) ->
176+
let multiplicationResult, columns, rowPointers =
177+
expand processor leftMatrix rightMatrix
178+
179+
let rows =
180+
expandRowPointers processor DeviceOnly rowPointers columns.Length leftMatrix.RowCount
181+
182+
rowPointers.Free processor
183+
184+
// sorting
185+
let sortData = sortData processor
186+
let sortKeys = sortKeys processor
187+
188+
// by columns
189+
let valuesSortedByColumns = sortData columns multiplicationResult
190+
let byKeSortedRows = sortKeys columns rows
191+
192+
multiplicationResult.Free processor
193+
rows.Free processor
194+
195+
// by rows
196+
let values = sortData byKeSortedRows valuesSortedByColumns
197+
let columns = sortKeys byKeSortedRows columns
198+
199+
// reduce
200+
201+
202+
203+
()
204+

‎tests/GraphBLAS-sharp.Tests/Common/Reduce/ReduceByKey.fs‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ module GraphBLAS.FSharp.Tests.Backend.Common.Reduce.ByKey
22

33
open Expecto
44
open GraphBLAS.FSharp.Backend.Common
5+
open GraphBLAS.FSharp.Test
56
open GraphBLAS.FSharp.Tests
67
open GraphBLAS.FSharp.Backend.Objects.ClContext
78
open Brahma.FSharp
@@ -185,3 +186,156 @@ let sequentialSegmentTests =
185186
createTestSequentialSegments<bool> (=) (&&) <@ (&&) @> ]
186187

187188
testList "Sequential segments" [ addTests; mulTests ]
189+
190+
let checkResult2D isEqual firstActualKeys secondActualKeys actualValues firstKeys secondKeys values reduceOp =
191+
192+
let expectedFirstKeys, expectedSecondKeys, expectedValues =
193+
HostPrimitives.reduceByKey2D firstKeys secondKeys values reduceOp
194+
195+
"First keys must be the same"
196+
|> Utils.compareArrays (=) firstActualKeys expectedFirstKeys
197+
198+
"Second keys must be the same"
199+
|> Utils.compareArrays (=) secondActualKeys expectedSecondKeys
200+
201+
"Values must the same"
202+
|> Utils.compareArrays isEqual actualValues expectedValues
203+
204+
let makeTest2D isEqual reduce reduceOp (array: (int * int * 'a) []) =
205+
let firstKeys, secondKeys, values =
206+
array
207+
|> Array.sortBy (fun (fst, snd, _) -> fst, snd)
208+
|> Array.unzip3
209+
210+
if firstKeys.Length > 0 then
211+
let clFirstKeys =
212+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, firstKeys)
213+
214+
let clSecondKeys =
215+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, secondKeys)
216+
217+
let clValues =
218+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, values)
219+
220+
let resultLength = Array.length <| Array.distinctBy (fun (fst, snd, _) -> (fst, snd)) array
221+
222+
let clFirstActualKeys, clSecondActualKeys, clActualValues: ClArray<int> * ClArray<int> * ClArray<'a> =
223+
reduce processor HostInterop resultLength clFirstKeys clSecondKeys clValues
224+
225+
clValues.Free processor
226+
clFirstKeys.Free processor
227+
clSecondKeys.Free processor
228+
229+
let actualValues = clActualValues.ToHostAndFree processor
230+
let firstActualKeys = clFirstActualKeys.ToHostAndFree processor
231+
let secondActualKeys = clSecondActualKeys.ToHostAndFree processor
232+
233+
checkResult2D isEqual firstActualKeys secondActualKeys actualValues firstKeys secondKeys values reduceOp
234+
235+
let createTestSequential2D<'a> (isEqual: 'a -> 'a -> bool) reduceOp reduceOpQ =
236+
237+
let reduce =
238+
Reduce.ByKey2D.sequential context Utils.defaultWorkGroupSize reduceOpQ
239+
240+
makeTest2D isEqual reduce reduceOp
241+
|> testPropertyWithConfig { config with arbitrary = [ typeof<Generators.ArrayOfDistinctKeys> ]; endSize = 10 } $"test on {typeof<'a>}"
242+
243+
let sequential2DTest =
244+
let addTests =
245+
testList
246+
"add tests"
247+
[ createTestSequential2D<int> (=) (+) <@ (+) @>
248+
createTestSequential2D<byte> (=) (+) <@ (+) @>
249+
250+
if Utils.isFloat64Available context.ClDevice then
251+
createTestSequential2D<float> Utils.floatIsEqual (+) <@ (+) @>
252+
253+
createTestSequential2D<float32> Utils.float32IsEqual (+) <@ (+) @>
254+
createTestSequential2D<bool> (=) (||) <@ (||) @> ]
255+
256+
let mulTests =
257+
testList
258+
"mul tests"
259+
[ createTestSequential2D<int> (=) (*) <@ (*) @>
260+
createTestSequential2D<byte> (=) (*) <@ (*) @>
261+
262+
if Utils.isFloat64Available context.ClDevice then
263+
createTestSequential2D<float> Utils.floatIsEqual (*) <@ (*) @>
264+
265+
createTestSequential2D<float32> Utils.float32IsEqual (*) <@ (*) @>
266+
createTestSequential2D<bool> (=) (&&) <@ (&&) @> ]
267+
268+
testList "Sequential 2D" [ addTests; mulTests ]
269+
270+
let makeTestSequentialSegments2D isEqual reduce reduceOp (array: (int * int * 'a) []) =
271+
272+
let firstKeys, secondKeys, values =
273+
array
274+
|> Array.sortBy (fun (fst, snd, _) -> fst, snd)
275+
|> Array.unzip3
276+
277+
if firstKeys.Length > 0 then
278+
let offsets =
279+
array
280+
|> Array.map (fun (fst, snd, _) -> fst, snd)
281+
|> HostPrimitives.getUniqueBitmapFirstOccurrence
282+
|> HostPrimitives.getBitPositions
283+
284+
let resultLength = offsets.Length
285+
286+
let firstKeys, secondKeys, values = Array.unzip3 array
287+
288+
let clOffsets =
289+
context.CreateClArrayWithSpecificAllocationMode(HostInterop, offsets)
290+
291+
let clFirstKeys =
292+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, firstKeys)
293+
294+
let clSecondKeys =
295+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, secondKeys)
296+
297+
let clValues =
298+
context.CreateClArrayWithSpecificAllocationMode(DeviceOnly, values)
299+
300+
let clFirstActualKeys, clSecondActualKeys, clReducedValues: ClArray<int> * ClArray<int> * ClArray<'a> =
301+
reduce processor DeviceOnly resultLength clOffsets clFirstKeys clSecondKeys clValues
302+
303+
let reducedFirsKeys = clFirstActualKeys.ToHostAndFree processor
304+
let reducesSecondKeys = clSecondActualKeys.ToHostAndFree processor
305+
let reducedValues = clReducedValues.ToHostAndFree processor
306+
307+
checkResult2D isEqual reducedFirsKeys reducesSecondKeys reducedValues firstKeys secondKeys values reduceOp
308+
309+
let createTestSequentialSegments2D<'a> (isEqual: 'a -> 'a -> bool) reduceOp reduceOpQ =
310+
let reduce =
311+
Reduce.ByKey2D.segmentSequential context Utils.defaultWorkGroupSize reduceOpQ
312+
313+
makeTestSequentialSegments2D isEqual reduce reduceOp
314+
|> testPropertyWithConfig { config with arbitrary = [ typeof<Generators.ArrayOfDistinctKeys> ] } $"test on {typeof<'a>}"
315+
316+
let sequentialSegmentTests2D =
317+
let addTests =
318+
testList
319+
"add tests"
320+
[ createTestSequentialSegments2D<int> (=) (+) <@ (+) @>
321+
createTestSequentialSegments2D<byte> (=) (+) <@ (+) @>
322+
323+
if Utils.isFloat64Available context.ClDevice then
324+
createTestSequentialSegments2D<float> Utils.floatIsEqual (+) <@ (+) @>
325+
326+
createTestSequentialSegments2D<float32> Utils.float32IsEqual (+) <@ (+) @>
327+
createTestSequentialSegments2D<bool> (=) (||) <@ (||) @> ]
328+
329+
let mulTests =
330+
testList
331+
"mul tests"
332+
[ createTestSequentialSegments2D<int> (=) (*) <@ (*) @>
333+
createTestSequentialSegments2D<byte> (=) (*) <@ (*) @>
334+
335+
if Utils.isFloat64Available context.ClDevice then
336+
createTestSequentialSegments2D<float> Utils.floatIsEqual (*) <@ (*) @>
337+
338+
createTestSequentialSegments2D<float32> Utils.float32IsEqual (*) <@ (*) @>
339+
createTestSequentialSegments2D<bool> (=) (&&) <@ (&&) @> ]
340+
341+
testList "Sequential segments 2D" [ addTests; mulTests ]

‎tests/GraphBLAS-sharp.Tests/Helpers.fs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -233,8 +233,8 @@ module HostPrimitives =
233233
|> Array.mapi (fun index bit -> if bit = 1 then Some index else None)
234234
|> Array.choose id
235235

236-
let reduceByKey keys value reduceOp =
237-
let zipped = Array.zip keys value
236+
let reduceByKey keys values reduceOp =
237+
let zipped = Array.zip keys values
238238

239239
Array.distinct keys
240240
|> Array.map
@@ -247,6 +247,12 @@ module HostPrimitives =
247247
|> Array.map (fun (key, values) -> key, Array.reduce reduceOp values)
248248
|> Array.unzip
249249

250+
let reduceByKey2D firstKeys secondKeys values reduceOp =
251+
Array.zip firstKeys secondKeys
252+
|> fun compactedKeys -> reduceByKey compactedKeys values reduceOp
253+
||> Array.map2 (fun (fst, snd) value -> fst, snd, value)
254+
|> Array.unzip3
255+
250256
module Context =
251257
type TestContext =
252258
{ ClContext: ClContext

0 commit comments

Comments
 (0)