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 Compiler.Xla.ShapeUtil
18 |
19 | import Compiler.FFI
20 | import Compiler.Xla.Shape
21 | import Compiler.Xla.XlaData
22 | import Types
23 |
24 | ffi : String -> String
25 | ffi = libxla "c/xla/shape_util.h"
26 |
27 | %foreign (ffi "MakeShape")
28 | prim__mkShape : Int -> GCPtr Int -> Int -> PrimIO AnyPtr
29 |
30 | export
31 | mkShape : (HasIO io, Primitive dtype) => Types.Shape -> io Xla.Shape
32 | mkShape shape = do
33 |   let dtypeEnum = xlaIdentifier {dtype}
34 |   MkIntArray shapeArrayPtr <- mkIntArray shape
35 |   shapePtr <- primIO $ prim__mkShape dtypeEnum shapeArrayPtr (cast $ length shape)
36 |   shapePtr <- onCollectAny' shapePtr Shape.delete
37 |   pure (MkShape shapePtr)
38 |