17 | module Compiler.Eval

 19 | import Data.IOArray

 21 | import Compiler.Enzyme.MLIR.Dialect.Ops

 22 | import Compiler.LLVM.ADT.APFloat

 23 | import Compiler.LLVM.ADT.APInt

 24 | import Compiler.LLVM.Support.RawOStream

 25 | import Compiler.MLIR.Dialect.Func.IR.FuncOps

 26 | import Compiler.MLIR.IR

 27 | import Compiler.MLIR.Pass.PassManager

 28 | import Compiler.Stablehlo.Dialect.ChloOps

 29 | import Compiler.Stablehlo.Dialect.Serialization

 30 | import Compiler.Stablehlo.Dialect.StablehloAttrs

 31 | import Compiler.Stablehlo.Dialect.StablehloEnums

 32 | import Compiler.Stablehlo.Dialect.StablehloOps

 33 | import Compiler.Stablehlo.Dialect.Version

 34 | import Compiler.Xla.Client.ExecutableBuildOptions

 35 | import Compiler.Xla.PJRT.C.PjrtCApi

 36 | import Compiler.Xla.PJRT.PjrtExecutable

 37 | import Compiler.DType

 38 | import Compiler.IR

 39 | import Compiler.Array

 40 | import Compiler.FFI

 41 | import Compiler.String

 42 | import Compiler.LiteralRW

 43 | import DType

 44 | import Literal

 45 | import Util

 46 | import Device

 49 | data Err

 50 |   = OutOfBounds Nat Nat

 51 |   | ValueNotFound Nat

 52 |   | PjrtErr PjrtError

 53 |   | MlirPassError String

 54 |   | InvalidHloError String

 56 | data BoundSet : Type where

 57 |   Parameters : Block -> BoundSet

 58 |   OpLike : {auto iface : Op a} -> a -> BoundSet

 61 | Show Err where

 62 |   show (OutOfBounds idx size) = "Index \{show idx} is out of bounds for array of size \{show size}"

 63 |   show (ValueNotFound idx) = "Value not found at index \{show idx}"

 64 |   show (PjrtErr err) = show err

 65 |   show (MlirPassError err) = "MlirPassError: \{err}"

 66 |   show (InvalidHloError err) = "InvalidHloError: \{err}"

 68 | public export 0

 69 | ErrIO : Type -> Type

 70 | ErrIO = EitherT Err IO

 72 | set : IOArray a -> Nat -> a -> ErrIO ()

 73 | set cache idx x = do

 74 |   False <- writeArray cache (cast idx) x | True => right ()

 75 |   left $ OutOfBounds idx (cast $ max cache)

 77 | get : IOArray a -> Nat -> ErrIO a

 78 | get cache idx = do

 79 |   Nothing <- readArray cache (cast idx) | Just x => right x

 80 |   let max = cast (max cache)

 81 |   left $ if idx >= max then OutOfBounds idx max else ValueNotFound idx

 83 | iboundset : Nat -> BoundSet -> ErrIO Value.Value

 84 | iboundset pos (Parameters block) = cast <$> getArgument block pos

 85 | iboundset pos (OpLike x {iface}) = cast <$> (flip getOpResult pos =<< getOperation x)

 87 | itype : IOArray BoundSet -> MLIRContext -> ValueType -> ErrIO Type_

 88 | itype cache ctx (TensorType shape dtype) = cast <$> RankedTensorType.get shape !(mlirType ctx dtype)

 89 | itype cache ctx (TypeRef pos tag) = getType =<< (iboundset pos =<< get cache tag) {m = ErrIO}

 91 | itypes : IOArray BoundSet -> MLIRContext -> Vect n ValueType -> ErrIO TypeRange

 92 | itypes cache ctx types = mkTypeRange =<< traverse (itype cache ctx) (toList types)

 94 | 0 Finalizer : Type -> Type

 95 | Finalizer a = OpBuilder -> Location -> ValueRange -> ErrIO a

 98 | interpretBody :

 99 |   {auto moduleOp : ModuleOp} ->

100 |   IOArray BoundSet ->

101 |   MLIRContext ->

102 |   Location ->

103 |   Block ->

104 |   Fn arity ->

105 |   Finalizer a ->

106 |   ErrIO ()

109 | interpretFunc :

110 |   {auto moduleOp : ModuleOp} ->

111 |   MLIRContext ->

112 |   Location ->

113 |   IOArray BoundSet ->

114 |   String ->

115 |   Fn n ->

116 |   ErrIO FuncOp

117 | interpretFunc ctx uloc cache name f = do

118 |   fnType <- FunctionType.get ctx !(itypes cache ctx f.paramTypes) !(itypes cache ctx f.resultTypes)

119 |   fn <- FuncOp.create uloc name fnType

120 |   interpretBody cache ctx uloc !(addEntryBlock fn) f FuncOps.ReturnOp.create

121 |   pushBack moduleOp !(getOperation fn)

122 |   pure fn

124 | interpretBody cache ctx uloc block f finalizer = do

125 |   builder <- atBlockEnd block

126 |   set cache f.tag (Parameters block)

127 |   for_ (reverse f.env.ops) $ \(i, expr) => set cache i !(iop expr)

128 |   results <- traverse ivalue f.results

129 |   ignore $ finalizer builder uloc !(mkValueRange $ toList results)

131 |   where

133 |   iop : {auto builder : OpBuilder} -> Op -> ErrIO BoundSet

135 |   iopref : {auto builder : OpBuilder} -> OpRef -> ErrIO BoundSet

136 |   iopref (BoundSet k) = get cache k

137 |   iopref (Concrete x) = iop x

139 |   ivalue : {auto builder : OpBuilder} -> IR.Value -> ErrIO Value.Value

140 |   ivalue (V pos op) = iopref op >>= iboundset pos

142 |   addArguments : Fn n -> Block -> ErrIO ()

143 |   addArguments f body = for_ f.paramTypes $ \t => addArgument body !(itype cache ctx t) uloc

145 |   iop (NamedFunc f) = OpLike <$> interpretFunc ctx uloc cache "func\{show f.tag}" f

146 |   iop (CallByName fTag resTys xs) = do

147 |     resTys <- mkTypeRange !(traverse (itype cache ctx) resTys)

148 |     let name = "func\{show fTag}"

149 |     args <- mkValueRange !(traverse ivalue $ toList xs)

150 |     OpLike <$> CallOp.create builder uloc name resTys args

151 |   iop (Grad shape f x) = do

152 |     revInit <- ivalue $ V 0 $ Concrete $ Lit [] F64 1.0

153 |     args <- mkValueRange [!(ivalue x), cast revInit]

154 |     retTys <- mkTypeRange [!(itype cache ctx $ TensorType shape F64)]

155 |     op <- AutoDiffRegionOp.create builder ctx uloc retTys args EnzymeActive EnzymeActivenoneed

156 |     body <- emplaceBlock $ getBody op

157 |     addArguments f body

158 |     interpretBody cache ctx uloc body f YieldOp.create

159 |     pure $ OpLike op

160 |   iop (Vectorize outTys batch fTag xs) = do

161 |     let fName = "func\{show fTag}"

162 |     outTys <- mkTypeRange !(traverse (itype cache ctx) outTys)

163 |     inputs <- mkValueRange !(traverse ivalue xs)

164 |     OpLike <$> BatchOp.create builder ctx uloc outTys fName inputs batch

165 |   iop (MinValue dtype) = do

166 |     apInt <-

167 |       if isSigned dtype

168 |       then getSignedMinValue (numBits dtype)

169 |       else getMinValue (numBits dtype)

170 |     type <- cast <$> RankedTensorType.get [] !(mlirType ctx dtype)

171 |     attr <- APInt.get type apInt

172 |     OpLike <$> ConstantOp.create builder uloc attr

173 |   iop (MaxValue dtype) = do

174 |     apInt <-

175 |       if isSigned dtype

176 |       then getSignedMaxValue (numBits dtype)

177 |       else getMaxValue (numBits dtype)

178 |     type <- cast <$> RankedTensorType.get [] !(mlirType ctx dtype)

179 |     attr <- APInt.get type apInt

180 |     OpLike <$> ConstantOp.create builder uloc attr

181 |   iop MinFiniteFloat = do

182 |     type <- cast <$> RankedTensorType.get [] !(mlirType ctx F64)

183 |     attr <- APFloat.get type !(getLargest True)

184 |     OpLike <$> ConstantOp.create builder uloc attr

185 |   iop MaxFiniteFloat = do

186 |     type <- cast <$> RankedTensorType.get [] !(mlirType ctx F64)

187 |     attr <- APFloat.get type !(getLargest False)

188 |     OpLike <$> ConstantOp.create builder uloc attr

189 |   iop (Lit shape dtype lit) =

190 |     OpLike <$> ConstantOp.create builder uloc !(mkDenseElementsAttr ctx lit)

191 |   iop (Broadcast bs x) = do

192 |     x <- ivalue x

193 |     xType <- getType x

194 |     let xType = unsafeCastToRankedTensorType xType

195 |         from = cast {to = Nat} <$> Array.toList !(getShape xType)

196 |         to : Shape = case bs of

197 |            Explicit to => to

198 |            AddLeading lead => lead ++ from

199 |     resTy <- clone xType !(fromList $ cast {to = Int64} <$> to)

200 |     if elem 0 to

201 |       then do

202 |         attr <- DenseElementsAttr.ArrayRef.Attribute.get resTy []

203 |         OpLike <$> ConstantOp.create builder uloc attr

204 |       else do

205 |         let broadcastDims = Prelude.map (+ length to `minus` length from) $ List.range $ length from

206 |         OpLike <$> BroadcastInDimOp.create builder uloc (cast resTy) x broadcastDims

207 |   iop (UnaryElementwise f x) = do

208 |     let W mkop @{iface} = f.create

209 |     OpLike <$> mkop builder uloc !(ivalue x)

211 |     where

213 |     data Wrap : Type where

214 |       W : forall a . (OpBuilder -> Location -> Value.Value -> ErrIO a) -> Op a => Wrap

216 |     (.create) : UnaryOp -> Wrap

217 |     (.create) = \case

218 |       Abs        => W AbsOp.create

219 |       Ceil       => W CeilOp.create

220 |       Cos        => W CosineOp.create

221 |       Exp        => W ExpOp.create

222 |       Floor      => W FloorOp.create

223 |       Log        => W LogOp.create

224 |       Logistic   => W LogisticOp.create

225 |       Not        => W NotOp.create

226 |       Neg        => W NegOp.create

227 |       Sin        => W SineOp.create

228 |       Sqrt       => W SqrtOp.create

229 |       Tan        => W TanOp.create

230 |       Tanh       => W TanhOp.create

232 |       Acos       => W AcosOp.create

233 |       Acosh      => W AcoshOp.create

234 |       Asin       => W AsinOp.create

235 |       Asinh      => W AsinhOp.create

236 |       Atan       => W AtanOp.create

237 |       Atanh      => W AtanhOp.create

238 |       Cosh       => W CoshOp.create

239 |       Sinh       => W SinhOp.create

240 |       Erf        => W ErfOp.create

241 |       ErfInv     => W ErfInvOp.create

242 |       Square     => W SquareOp.create

243 |   iop (Convert dtype resultShape x) = do

244 |     resultType <- cast <$> RankedTensorType.get resultShape !(mlirType ctx dtype)

245 |     OpLike <$> ConvertOp.create builder uloc resultType !(ivalue x)

246 |   iop (BitCastConvert dtype resultShape x) = do

247 |     resultType <- cast <$> RankedTensorType.get resultShape !(mlirType ctx dtype)

248 |     OpLike <$> ConvertOp.create builder uloc resultType !(ivalue x)

249 |   iop (BinaryElementwise f lhs rhs) = do

250 |     let W mkop @{iface} = f.create

251 |     OpLike <$> mkop builder uloc !(ivalue lhs) !(ivalue rhs)

253 |     where

255 |     data Wrap : Type where

256 |       W : forall a . (OpBuilder -> Location -> Value.Value -> Value.Value -> ErrIO a) -> Op a => Wrap

258 |     (.create) : BinaryOp -> Wrap

259 |     (.create) = \case

260 |       Compare direction => W (\b, l, x, y => CompareOp.create b l x y (cast direction))

261 |       Add => W AddOp.create

262 |       Div => W DivOp.create

263 |       Max => W MaxOp.create

264 |       Min => W MinOp.create

265 |       Mul => W MulOp.create

266 |       Pow => W PowOp.create

267 |       Rem => W RemOp.create

268 |       Sub => W SubtractOp.create

269 |       And => W AndOp.create

270 |       Or  => W OrOp.create

271 |       ShiftRightLogical => W ShiftRightLogicalOp.create

272 |   iop (If resTy pred true false) = do

273 |     op <- IfOp.create builder uloc !(itype cache ctx resTy) !(ivalue pred)

274 |     bodyT <- emplaceBlock $ getTrueBranch op

275 |     bodyF <- emplaceBlock $ getFalseBranch op

276 |     addArguments true bodyT

277 |     addArguments false bodyF

278 |     interpretBody cache ctx uloc bodyT true StablehloOps.ReturnOp.create

279 |     interpretBody cache ctx uloc bodyF false StablehloOps.ReturnOp.create

280 |     pure $ OpLike op

281 |   iop (While cond body inits) = do

282 |     inits <- mkValueRange =<< traverse ivalue (toList inits)

283 |     op <- WhileOp.create builder uloc inits

284 |     cond' <- emplaceBlock $ getCond op

285 |     body' <- emplaceBlock $ getBody op

286 |     addArguments body body'

287 |     addArguments cond cond'

288 |     interpretBody cache ctx uloc cond' cond StablehloOps.ReturnOp.create

289 |     interpretBody cache ctx uloc body' body StablehloOps.ReturnOp.create

290 |     pure $ OpLike op

291 |   iop (Reduce body inits axes xs) = do

292 |     inits <- mkValueRange !(traverse ivalue $ toList inits)

293 |     xs <- mkValueRange !(traverse ivalue $ toList xs)

294 |     op <- ReduceOp.create builder uloc xs inits axes

295 |     body' <- emplaceBlock $ getBody op

296 |     addArguments body body'

297 |     interpretBody cache ctx uloc body' body StablehloOps.ReturnOp.create

298 |     pure $ OpLike op

299 |   iop (Slice starts stops strides x) = do

300 |     OpLike <$> SliceOp.create builder uloc !(ivalue x) starts stops strides

301 |   iop (DynamicSlice starts sizes x) = do

302 |     starts <- mkValueRange !(traverse ivalue starts)

303 |     OpLike <$> DynamicSliceOp.create builder uloc !(ivalue x) starts sizes

304 |   iop (Cholesky x) = OpLike <$> CholeskyOp.create builder uloc !(ivalue x) True

305 |   iop (Concat axis xs) = do

306 |     xs <- mkValueRange =<< traverse ivalue (toList xs)

307 |     OpLike <$> ConcatenateOp.create builder uloc xs axis

308 |   iop (Iota shape dtype dim) = do

309 |     resultType <- cast <$> RankedTensorType.get shape !(mlirType ctx dtype)

310 |     OpLike <$> IotaOp.create builder uloc resultType dim

311 |   iop (DotGeneral lb rb lc rc resultType lhs rhs) = do

312 |     ddn <- DotDimensionNumbersAttr.get ctx lb rb lc rc

313 |     resTy <- itype cache ctx resultType

314 |     OpLike <$> DotGeneralOp.create builder uloc resTy !(ivalue lhs) !(ivalue rhs) ddn

315 |   iop (Map f xs resTy dims) = do

316 |     xs <- mkValueRange =<< traverse ivalue (toList xs)

317 |     op <- MapOp.create builder uloc !(itype cache ctx resTy) xs dims

318 |     f' <- emplaceBlock $ getComputation op

319 |     addArguments f f'

320 |     interpretBody cache ctx uloc f' f StablehloOps.ReturnOp.create

321 |     pure $ OpLike op

322 |   iop (Reshape dtype to x) = do

323 |     resultType <- cast <$> RankedTensorType.get to !(mlirType ctx dtype)

324 |     OpLike <$> ReshapeOp.create builder uloc resultType !(ivalue x)

325 |   iop (Select pred true false) = do

326 |     OpLike <$> SelectOp.create builder uloc !(ivalue pred) !(ivalue true) !(ivalue false)

327 |   iop (Sort comp axis isStable x) = do

328 |     op <- SortOp.create builder uloc !(ivalue x) axis isStable

329 |     comp' <- emplaceBlock $ getComparator op

330 |     addArguments comp comp'

331 |     interpretBody cache ctx uloc comp' comp StablehloOps.ReturnOp.create

332 |     pure $ OpLike op

333 |   iop (Reverse axes x) = OpLike <$> ReverseOp.create builder uloc !(ivalue x) axes

334 |   iop (Transpose ordering x) = OpLike <$> TransposeOp.create builder uloc !(ivalue x) ordering

335 |   iop (TriangularSolve a b lower) =

336 |     OpLike <$> TriangularSolveOp.create

337 |       builder uloc !(ivalue a) !(ivalue b) True lower False NoTranspose

338 |   iop (Rng state resultType) = do

339 |     stateTy <- itype cache ctx $ TensorType [2] U64

340 |     OpLike <$> RngBitGeneratorOp.create

341 |       builder uloc stateTy !(itype cache ctx resultType) ThreeFry !(ivalue state)

343 | public export

344 | data TensorData : (Shape -> DType -> Type) -> Shape -> DType -> Type where

345 |   MkTensorData : {shape : _} -> {dtype : _} -> f shape dtype -> TensorData f shape dtype

347 | %hide Literal.All2.All2

349 | replace :

350 |   (xs : All2 (TensorData $ \_, _ => a) shapes dtypes) ->

351 |   Vect (length xs) b ->

352 |   All2 (TensorData $ \_, _ => b) shapes dtypes

353 | replace [] [] = []

354 | replace (MkTensorData _ :: xs) (y :: ys) = MkTensorData y :: replace xs ys

356 | export covering

357 | execute :

358 |   Device ->

359 |   Fn 0 ->

360 |   All2 (TensorData $ \_, _ => ()) shapes dtypes ->

361 |   ErrIO $ All2 (TensorData $ \_ => Array . DType.idrisType) shapes dtypes

362 | execute (MkDevice mlirCtx passManager api client) f types = do

363 |   uloc <- UnknownLoc.get mlirCtx

364 |   cache <- newArray $ cast $ counter f.env

365 |   moduleOp <- ModuleOp.create uloc "root"

366 |   ignore $ interpretFunc mlirCtx uloc cache "main" f

367 |   True <- run passManager !(getOperation moduleOp)

368 |     | _ => throwE $ MlirPassError "Failed to run MLIR passes"

369 |   code <- cppString

370 |   version <- toString !getCurrentVersion

371 |   True <- serializePortableArtifact moduleOp version !(rawStringOStream code)

372 |     | False => throwE $ InvalidHloError "Failed to serialize MLIR for version \{c_str version}"

373 |   bimapEitherT PjrtErr id $ do

374 |     executableBuildOptions <- mkExecutableBuildOptions

375 |     compileOptions <- serializeAsString !(mkCompileOptions executableBuildOptions)

376 |     program <- mkPjrtProgram code

377 |     loadedExec <- pjrtClientCompile api client program compileOptions

378 |     delete program

379 |     delete code

380 |     delete compileOptions

381 |     delete executableBuildOptions

383 |     buffers <- pjrtLoadedExecutableExecute api loadedExec (length types)

384 |     pjrtLoadedExecutableDestroy api loadedExec

386 |     forProperty (replace types buffers) $ \(MkTensorData {shape, dtype} buffer) => do

387 |       let at = dtypeIsArrayType dtype

388 |       arr <- mkArray (product shape)

389 |       pjrtBufferToHostBuffer api buffer arr

390 |       pjrtBufferDestroy api buffer

391 |       pure (MkTensorData arr)