0 | {--
 1 | Copyright (C) 2022  Joel Berkeley
 2 |
 3 | This program is free software: you can redistribute it and/or modify
 4 | it under the terms of the GNU Affero General Public License as published
 5 | by the Free Software Foundation, either version 3 of the License, or
 6 | (at your option) any later version.
 7 |
 8 | This program is distributed in the hope that it will be useful,
 9 | but WITHOUT ANY WARRANTY; without even the implied warranty of
10 | MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 | GNU Affero General Public License for more details.
12 |
13 | You should have received a copy of the GNU Affero General Public License
14 | along with this program.  If not, see <https://www.gnu.org/licenses/>.
15 | --}
16 | ||| For internal spidr use only.
17 | module Spidr.Compiler.LiteralRW
18 |
19 | import Control.Monad.State
20 |
21 | import Spidr.Compiler.MLIR.IR.MLIRContext
22 | import Spidr.Compiler.MLIR.IR.Types
23 | import Spidr.Compiler.MLIR.IR.BuiltinTypes
24 | import Spidr.Compiler.MLIR.IR.BuiltinAttributes
25 | import Spidr.Compiler.Array
26 | import Spidr.Compiler.DType
27 | import Spidr.Compiler.FFI
28 | import Spidr.Data.Literal
29 | import Spidr.DType
30 | import Spidr.Data.List
31 |
32 | export
33 | dtypeIsArrayType : (dtype : DType) -> ArrayType (idrisType dtype)
34 | dtypeIsArrayType PRED = %search
35 | dtypeIsArrayType S8   = %search
36 | dtypeIsArrayType S16  = %search
37 | dtypeIsArrayType S32  = %search
38 | dtypeIsArrayType S64  = %search
39 | dtypeIsArrayType U8   = %search
40 | dtypeIsArrayType U16  = %search
41 | dtypeIsArrayType U32  = %search
42 | dtypeIsArrayType U64  = %search
43 | dtypeIsArrayType F64  = %search
44 |
45 | export
46 | mkDenseElementsAttr :
47 |   HasIO io =>
48 |   MLIRContext ->
49 |   {shape : _} -> {dtype : _} ->
50 |   Literal shape (idrisType dtype) ->
51 |   io DenseElementsAttr
52 | mkDenseElementsAttr ctx lit = do
53 |   type <- RankedTensorType.get shape !(mlirType ctx dtype)
54 |   arr <- let at = dtypeIsArrayType dtype in fromList $ toList lit
55 |   case dtype of
56 |     PRED => DenseElementsAttr.ArrayRef.Bool.get type arr
57 |     S8   => DenseElementsAttr.ArrayRef.Int8.get type arr
58 |     S16  => DenseElementsAttr.ArrayRef.Int16.get type arr
59 |     S32  => DenseElementsAttr.ArrayRef.Int32.get type arr
60 |     S64  => DenseElementsAttr.ArrayRef.Int64.get type arr
61 |     U8   => DenseElementsAttr.ArrayRef.UInt8.get type arr
62 |     U16  => DenseElementsAttr.ArrayRef.UInt16.get type arr
63 |     U32  => DenseElementsAttr.ArrayRef.UInt32.get type arr
64 |     U64  => DenseElementsAttr.ArrayRef.UInt64.get type arr
65 |     F64  => DenseElementsAttr.ArrayRef.Double.get type arr
66 |
67 | export
68 | read : {shape : _} -> (dtype : DType) -> Array (idrisType dtype) -> Literal shape (idrisType dtype)
69 | read dtype arr =
70 |   let arrIndices = evalState Z $ traverse (\() => State.get <* modify S) (pure ())
71 |       at = dtypeIsArrayType dtype
72 |    in Array.get arr <$> arrIndices
73 |