
  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 DType
 23 | import Types
 24 | 
 25 | ffi : String -> String
 26 | ffi = libxla "c/xla/shape_util.h"
 27 | 
 28 | %foreign (ffi "MakeShape")
 29 | prim__mkShape : Int -> GCPtr Int64 -> Bits64 -> PrimIO AnyPtr
 30 | 
 31 | export
 32 | mkShape : HasIO io => Types.Shape -> DType -> io Xla.Shape
 33 | mkShape shape dtype = do
 34 |   let dtypeEnum = xlaIdent dtype
 35 |   MkInt64Array shape shapeLen <- mkInt64Array $ cast <$> shape
 36 |   shapePtr <- primIO $ prim__mkShape dtypeEnum shape shapeLen
 37 |   shapePtr <- onCollectAny' shapePtr Shape.delete
 38 |   pure (MkShape shapePtr)
 39 |