add: Expand stage

Author: IgorErin

src/GraphBLAS-sharp.Backend/Matrix/CSRMatrix/SpGEMM/Expand.fs

@@ -73,99 +73,43 @@ module Expand =
 
             requiredRawsLengths
 
-    let expandRightMatrixValuesIndices (clContext: ClContext) workGroupSize =
-
-        let kernel =
-            <@ fun (ndRange: Range1D) length (globalRightMatrixValuesPositions: Indices) (requiredRightMatrixValuesPointers: Indices) (globalPositions: Indices) (result: Indices) ->
+    let extractLeftMatrixRequiredValuesAndColumns (clContext: ClContext) workGroupSize =
 
-                let gid = ndRange.GlobalID0
+        let getUniqueBitmap =
+            ClArray.getUniqueBitmap clContext workGroupSize
 
-                if gid < length then
-                    // index corresponding to the position of pointers
-                    let positionIndex = globalPositions.[gid] - 1
+        let prefixSumExclude =
+            PrefixSum.standardExcludeInplace clContext workGroupSize
 
-                    // the position of the beginning of a new line of pointers
-                    let sourcePosition = globalRightMatrixValuesPositions.[positionIndex]
+        let indicesScatter =
+            Scatter.runInplace clContext workGroupSize
 
-                    // offset from the source pointer
-                    let offsetFromSourcePosition = gid - sourcePosition
+        let dataScatter =
+            Scatter.runInplace clContext workGroupSize
 
-                    // pointer to the first element in the row of the right matrix from which
-                    // the offset will be counted to get pointers to subsequent elements in this row
-                    let sourcePointer = requiredRightMatrixValuesPointers.[positionIndex]
+        fun (processor: MailboxProcessor<_>) (leftMatrix: ClMatrix.CSR<'a>) (globalRightMatrixRawsStartPositions: Indices) ->
 
-                    // adding up the mix with the source pointer,
-                    // we get a pointer to a specific element in the raw
-                    result.[gid] <- sourcePointer + offsetFromSourcePosition @>
+            let leftMatrixRequiredPositions, resultLength  =
+                let bitmap =
+                    getUniqueBitmap processor DeviceOnly globalRightMatrixRawsStartPositions
 
-        let kernel = clContext.Compile kernel
+                let length = (prefixSumExclude processor bitmap).ToHostAndFree processor
 
-        fun (processor: MailboxProcessor<_>) (resultLength: int) (globalRightMatrixRawsStartPositions: Indices) (requiredRightMatrixValuesPointers: Indices) (globalPositions: Indices) ->
+                bitmap, length
 
-            let globalRightMatrixValuesPointers =
+            let requiredLeftMatrixValues =
                 clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
 
-            let kernel = kernel.GetKernel()
-
-            let ndRange =
-                Range1D.CreateValid(resultLength, workGroupSize)
-
-            processor.Post(
-                Msg.MsgSetArguments
-                    (fun () ->
-                        kernel.KernelFunc
-                            ndRange
-                            resultLength
-                            globalRightMatrixRawsStartPositions
-                            requiredRightMatrixValuesPointers
-                            globalPositions
-                            globalRightMatrixValuesPointers)
-            )
-
-            processor.Post <| Msg.CreateRunMsg<_, _> kernel
-
-            globalRightMatrixValuesPointers
-
-    let getResultRowPointers (clContext: ClContext) workGroupSize =
-
-        let kernel =
-            <@ fun (ndRange: Range1D) length (leftMatrixRowPointers: Indices) (globalArrayRightMatrixRawPointers: Indices) (result: Indices) ->
-
-                let gid = ndRange.GlobalID0
-
-                if gid < length then
-                    let rowPointer = leftMatrixRowPointers.[gid]
-                    let globalPointer = globalArrayRightMatrixRawPointers.[rowPointer]
-
-                    result.[gid] <- globalPointer
-                    @>
-
-        let kernel = clContext.Compile kernel
-
-        fun (processor: MailboxProcessor<_>) (leftMatrixRowPointers: Indices) (globalArrayRightMatrixRawPointers: Indices) ->
-
-            let result =
-                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, leftMatrixRowPointers.Length)
-
-            let kernel = kernel.GetKernel()
+            indicesScatter processor leftMatrixRequiredPositions leftMatrix.Values requiredLeftMatrixValues
 
-            let ndRange =
-                Range1D.CreateValid( leftMatrixRowPointers.Length, workGroupSize)
+            let requiredLeftMatrixColumns =
+                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
 
-            processor.Post(
-                Msg.MsgSetArguments
-                    (fun () ->
-                        kernel.KernelFunc
-                            ndRange
-                            leftMatrixRowPointers.Length
-                            leftMatrixRowPointers
-                            globalArrayRightMatrixRawPointers
-                            result)
-            )
+            dataScatter processor leftMatrixRequiredPositions leftMatrix.Columns requiredLeftMatrixColumns
 
-            processor.Post <| Msg.CreateRunMsg<_, _> kernel
+            leftMatrixRequiredPositions.Free processor
 
-            result
+            requiredLeftMatrixColumns, requiredLeftMatrixValues
 
     let getGlobalMap (clContext: ClContext) workGroupSize =
 
@@ -183,49 +127,56 @@ module Expand =
             // Insert units at the beginning of new lines (source positions)
             assignUnits processor globalRightMatrixValuesPositions globalPositions
 
-            // Apply the prefix sum,
+            // Apply the prefix sum, SIDE EFFECT!!!
             // get an array where different sub-arrays of pointers to elements of the same row differ in values
             (prefixSum processor globalPositions).Free processor
 
             globalPositions
 
-    let extractLeftMatrixRequiredValuesAndColumns (clContext: ClContext) workGroupSize =
-
-        let getUniqueBitmap =
-            ClArray.getUniqueBitmap clContext workGroupSize
+    let getResultRowPointers (clContext: ClContext) workGroupSize =
 
-        let prefixSumExclude =
-            PrefixSum.standardExcludeInplace clContext workGroupSize
+        let kernel =
+            <@ fun (ndRange: Range1D) length (leftMatrixRowPointers: Indices) (globalArrayRightMatrixRawPointers: Indices) (result: Indices) ->
 
-        let indicesScatter =
-            Scatter.runInplace clContext workGroupSize
+                let gid = ndRange.GlobalID0
 
-        let dataScatter =
-            Scatter.runInplace clContext workGroupSize
+                // do not touch the last element
+                if gid < length - 1 then
+                    let rowPointer = leftMatrixRowPointers.[gid]
+                    let globalPointer = globalArrayRightMatrixRawPointers.[rowPointer]
 
-        fun (processor: MailboxProcessor<_>) (leftMatrix: ClMatrix.CSR<'a>) (globalRightMatrixRawsStartPositions: Indices) ->
+                    result.[gid] <- globalPointer @>
 
-            let leftMatrixRequiredPositions, resultLength  =
-                let bitmap =
-                    getUniqueBitmap processor DeviceOnly globalRightMatrixRawsStartPositions
+        let kernel = clContext.Compile kernel
 
-                let length = (prefixSumExclude processor bitmap).ToHostAndFree processor
+        let createResultPointersBuffer = ClArray.create clContext workGroupSize
 
-                bitmap, length
+        fun (processor: MailboxProcessor<_>) (globalLength: int) (leftMatrixRowPointers: Indices) (globalRightMatrixRowPointers: Indices) ->
 
-            let requiredLeftMatrixValues =
-                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+            // The last element must be equal to the length of the global array.
+            let result =
+                createResultPointersBuffer processor DeviceOnly leftMatrixRowPointers.Length globalLength
 
-            indicesScatter processor leftMatrixRequiredPositions leftMatrix.Values requiredLeftMatrixValues
+            let kernel = kernel.GetKernel()
 
-            let requiredLeftMatrixColumns =
-                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, resultLength)
+            // do not touch the last element
+            let ndRange =
+                Range1D.CreateValid(leftMatrixRowPointers.Length - 1, workGroupSize)
 
-            dataScatter processor leftMatrixRequiredPositions leftMatrix.Columns requiredLeftMatrixColumns
+            processor.Post(
+                Msg.MsgSetArguments
+                    (fun () ->
+                        kernel.KernelFunc
+                            ndRange
+                            leftMatrixRowPointers.Length
+                            leftMatrixRowPointers
+                            globalRightMatrixRowPointers
+                            result)
+            )
 
-            leftMatrixRequiredPositions.Free processor
+            processor.Post <| Msg.CreateRunMsg<_, _> kernel
 
-            requiredLeftMatrixColumns, requiredLeftMatrixValues
+            result
 
     let processPositions (clContext: ClContext) workGroupSize =
 
@@ -242,6 +193,8 @@ module Expand =
 
         let getGlobalPositions = getGlobalMap clContext workGroupSize
 
+        let getRowPointers = getResultRowPointers clContext workGroupSize
+
         let getRequiredRightMatrixValuesPointers =
             processLeftMatrixColumnsAndRightMatrixRawPointers clContext workGroupSize requiredRawPointers
 
@@ -250,15 +203,14 @@ module Expand =
             let requiredRawsLengths =
                 getRequiredRawsLengths processor leftMatrix.Columns rightMatrix.RowPointers
 
-            // global expanded array length (sum of previous length)
+            // global expanded array length (sum of previous length) SIDE EFFECT!!!
             let globalLength =
                 (prefixSumExclude processor requiredRawsLengths).ToHostAndFree processor
 
             // rename array after side effect of prefix sum include
             // positions in global array for right matrix raws with duplicates
             let globalRightMatrixRowsStartPositions = requiredRawsLengths
 
-
             /// Extract required left matrix columns and values by global right matrix pointers.
             /// Then get required right matrix rows (pointers to rows) by required left matrix columns.
 
@@ -277,13 +229,70 @@ module Expand =
             let globalRightMatrixRawsPointersWithoutDuplicates =
                 removeDuplications processor globalRightMatrixRowsStartPositions
 
+            // RESULT row pointers into result expanded (obtained by multiplication) array
+            let resultRowPointers =
+                getRowPointers processor globalLength leftMatrix.RowPointers globalRightMatrixRowsStartPositions
+
             globalRightMatrixRowsStartPositions.Free processor
 
             // int map to distinguish different raws in a general array. 1 for first, 2 for second and so forth...
             let globalMap =
                 getGlobalPositions processor globalLength globalRightMatrixRawsPointersWithoutDuplicates
 
-            globalMap, globalRightMatrixRawsPointersWithoutDuplicates, requiredLeftMatrixValues, requiredRightMatrixRawPointers
+            globalMap, globalRightMatrixRawsPointersWithoutDuplicates, requiredLeftMatrixValues, requiredRightMatrixRawPointers, resultRowPointers
+
+    let expandRightMatrixValuesIndices (clContext: ClContext) workGroupSize =
+
+        let kernel =
+            <@ fun (ndRange: Range1D) length (globalRightMatrixValuesPositions: Indices) (requiredRightMatrixValuesPointers: Indices) (globalPositions: Indices) (result: Indices) ->
+
+                let gid = ndRange.GlobalID0
+
+                if gid < length then
+                    // index corresponding to the position of pointers
+                    let positionIndex = globalPositions.[gid] - 1
+
+                    // the position of the beginning of a new line of pointers
+                    let sourcePosition = globalRightMatrixValuesPositions.[positionIndex]
+
+                    // offset from the source pointer
+                    let offsetFromSourcePosition = gid - sourcePosition
+
+                    // pointer to the first element in the row of the right matrix from which
+                    // the offset will be counted to get pointers to subsequent elements in this row
+                    let sourcePointer = requiredRightMatrixValuesPointers.[positionIndex]
+
+                    // adding up the mix with the source pointer,
+                    // we get a pointer to a specific element in the raw
+                    result.[gid] <- sourcePointer + offsetFromSourcePosition @>
+
+        let kernel = clContext.Compile kernel
+
+        fun (processor: MailboxProcessor<_>) (globalRightMatrixRawsStartPositions: Indices) (requiredRightMatrixValuesPointers: Indices) (globalMap: Indices) ->
+
+            let globalRightMatrixValuesPointers =
+                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, globalMap.Length)
+
+            let kernel = kernel.GetKernel()
+
+            let ndRange =
+                Range1D.CreateValid(globalMap.Length, workGroupSize)
+
+            processor.Post(
+                Msg.MsgSetArguments
+                    (fun () ->
+                        kernel.KernelFunc
+                            ndRange
+                            globalMap.Length
+                            globalRightMatrixRawsStartPositions
+                            requiredRightMatrixValuesPointers
+                            globalMap
+                            globalRightMatrixValuesPointers)
+            )
+
+            processor.Post <| Msg.CreateRunMsg<_, _> kernel
+
+            globalRightMatrixValuesPointers
 
     let expandLeftMatrixValues (clContext: ClContext) workGroupSize =
 
@@ -330,56 +339,51 @@ module Expand =
 
         let gatherIndices = Gather.run clContext workGroupSize
 
-        fun (processor: MailboxProcessor<_>) (globalLength: int) (globalPositions: Indices) (rightMatrixValues: Values<'a>) (rightMatrixColumns: Indices) ->
+        fun (processor: MailboxProcessor<_>) (globalPositions: Indices) (rightMatrix: ClMatrix.CSR<'a>) ->
             // gather all required right matrix values
             let extendedRightMatrixValues =
-                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, globalLength)
+                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, globalPositions.Length)
 
-            gatherRightMatrixData processor globalPositions rightMatrixValues extendedRightMatrixValues
+            gatherRightMatrixData processor globalPositions rightMatrix.Values extendedRightMatrixValues
 
             // gather all required right matrix column indices
             let extendedRightMatrixColumns =
-                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, globalLength)
+                clContext.CreateClArrayWithSpecificAllocationMode(DeviceOnly, globalPositions.Length)
 
-            gatherIndices processor globalPositions rightMatrixColumns extendedRightMatrixColumns
+            gatherIndices processor globalPositions rightMatrix.Columns extendedRightMatrixColumns
 
             extendedRightMatrixValues, extendedRightMatrixColumns
 
     let run (clContext: ClContext) workGroupSize (multiplication: Expr<'a -> 'b -> 'c>) =
 
         let processPositions = processPositions clContext workGroupSize
 
-        let getRightMatrixValuesPointers =
+        let expandLeftMatrixValues =
+            expandLeftMatrixValues clContext workGroupSize
+
+        let expandRightMatrixValuesPointers =
             expandRightMatrixValuesIndices clContext workGroupSize
 
         let getRightMatrixColumnsAndValues =
             getRightMatrixColumnsAndValues clContext workGroupSize
 
-        let expandLeftMatrixValues =
-            expandLeftMatrixValues clContext workGroupSize
-
         let map2 = ClArray.map2 clContext workGroupSize multiplication
 
-        let getRawPointers = getResultRowPointers clContext workGroupSize
-
         fun (processor: MailboxProcessor<_>) (leftMatrix: ClMatrix.CSR<'a>) (rightMatrix: ClMatrix.CSR<'b>) ->
 
-            let globalMap, globalRightMatrixRowsPointers, requiredLeftMatrixValues, requiredRightMatrixRowPointers
+            let globalMap, globalRightMatrixRowsPointers, requiredLeftMatrixValues, requiredRightMatrixRowPointers, resultRowPointers
                 = processPositions processor leftMatrix rightMatrix
 
-            // left matrix values correspondingly to right matrix values // TODO()
+            // left matrix values correspondingly to right matrix values
             let extendedLeftMatrixValues =
-                expandLeftMatrixValues processor globalMap leftMatrix.Values
-
-            let resultRowPointers =
-                getRawPointers processor leftMatrix.RowPointers globalRightMatrixRowsPointers
+                expandLeftMatrixValues processor globalMap requiredLeftMatrixValues
 
             // extended pointers to all required right matrix numbers
             let globalRightMatrixValuesPointers =
-                getRightMatrixValuesPointers processor globalMap.Length globalRightMatrixRowsPointers requiredRightMatrixRowPointers globalMap
+                expandRightMatrixValuesPointers processor globalRightMatrixRowsPointers requiredRightMatrixRowPointers globalMap
 
             let extendedRightMatrixValues, extendedRightMatrixColumns =
-                getRightMatrixColumnsAndValues processor globalMap.Length globalRightMatrixValuesPointers rightMatrix.Values rightMatrix.Columns
+                getRightMatrixColumnsAndValues processor globalRightMatrixValuesPointers rightMatrix
 
             /// Multiplication
             let multiplicationResult  =