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.HLO.Translate.HloToMhlo.HloUtils

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

 37 | import Compiler.Xla.PJRT.PjrtExecutable

 38 | import Compiler.Xla.Shape

 39 | import Compiler.Xla.ShapeUtil

 40 | import Compiler.Xla.XlaData

 41 | import Compiler.DType

 42 | import Compiler.IR

 43 | import Compiler.FFI

 44 | import Compiler.LiteralRW

 45 | import DType

 46 | import Literal

 47 | import Util

 48 | import Device

 51 | data Err

 52 |   = OutOfBounds Nat Nat

 53 |   | ValueNotFound Nat

 54 |   | PjrtErr PjrtError

 55 |   | MlirPassError String

 56 |   | InvalidHloError String

 58 | data BoundSet : Type where

 59 |   Parameters : Block -> BoundSet

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

 63 | Show Err where

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

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

 66 |   show (PjrtErr err) = show err

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

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

 70 | public export 0

 71 | ErrIO : Type -> Type

 72 | ErrIO = EitherT Err IO

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

 75 | set cache idx x = do

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

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

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

 80 | get cache idx = do

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

 82 |   let max = cast (max cache)

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

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

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

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

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

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

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

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

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

 96 | 0 Finalizer : Type -> Type

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

100 | interpretBody :

101 |   {auto moduleOp : ModuleOp} ->

102 |   IOArray BoundSet ->

103 |   MLIRContext ->

104 |   Location ->

105 |   Block ->

106 |   Fn arity ->

107 |   Finalizer a ->

108 |   ErrIO ()

111 | interpretFunc :

112 |   {auto moduleOp : ModuleOp} ->

113 |   MLIRContext ->

114 |   Location ->

115 |   IOArray BoundSet ->

116 |   String ->

117 |   Fn n ->

118 |   ErrIO FuncOp

119 | interpretFunc ctx uloc cache name f = do

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

121 |   fn <- FuncOp.create uloc name fnType

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

123 |   pushBack moduleOp !(getOperation fn)

124 |   pure fn

126 | interpretBody cache ctx uloc block f finalizer = do

127 |   builder <- atBlockEnd block

128 |   set cache f.tag (Parameters block)

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

130 |   results <- traverse ivalue f.results

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

133 |   where

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

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

138 |   iopref (BoundSet k) = get cache k

139 |   iopref (Concrete x) = iop x

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

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

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

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

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

148 |   iop (CallByName tag resTys xs) = do

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

150 |     let name = "func\{show tag}"

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

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

153 |   iop (Grad shape f x) = do

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

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

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

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

158 |     body <- emplaceBlock $ getBody op

159 |     addArguments f body

160 |     interpretBody cache ctx uloc body f YieldOp.create

161 |     pure $ OpLike op

162 |   iop (MinValue dtype) = do

163 |     apInt <-

164 |       if isSigned dtype

165 |       then getSignedMinValue (numBits dtype)

166 |       else getMinValue (numBits dtype)

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

168 |     attr <- APInt.get type apInt

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

170 |   iop (MaxValue dtype) = do

171 |     apInt <-

172 |       if isSigned dtype

173 |       then getSignedMaxValue (numBits dtype)

174 |       else getMaxValue (numBits dtype)

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

176 |     attr <- APInt.get type apInt

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

178 |   iop MinFiniteFloat = do

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

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

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

182 |   iop MaxFiniteFloat = do

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

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

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

186 |   iop (Lit shape dtype lit) = do

187 |     attr <- createDenseElementsAttrFromLiteral !(write dtype lit) builder

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

189 |   iop (Broadcast dtype from to x) =

190 |     if elem 0 to && from /= to

191 |       then do

192 |         shape <- mkShape to dtype

193 |         literal <- allocLiteral shape

194 |         attr <- createDenseElementsAttrFromLiteral literal builder

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

196 |       else

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

198 |        in do

199 |         resTy <- itype cache ctx $ TensorType to dtype

200 |         OpLike <$> BroadcastInDimOp.create builder uloc resTy !(ivalue x) broadcastDims

201 |   iop (UnaryElementwise f x) = do

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

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

205 |     where

207 |     data Wrap : Type where

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

210 |     (.create) : UnaryOp -> Wrap

211 |     (.create) = \case

212 |       Abs        => W AbsOp.create

213 |       Ceil       => W CeilOp.create

214 |       Cos        => W CosineOp.create

215 |       Exp        => W ExpOp.create

216 |       Floor      => W FloorOp.create

217 |       Log        => W LogOp.create

218 |       Logistic   => W LogisticOp.create

219 |       Not        => W NotOp.create

220 |       Neg        => W NegOp.create

221 |       Sin        => W SineOp.create

222 |       Sqrt       => W SqrtOp.create

223 |       Tan        => W TanOp.create

224 |       Tanh       => W TanhOp.create

226 |       Acos       => W AcosOp.create

227 |       Acosh      => W AcoshOp.create

228 |       Asin       => W AsinOp.create

229 |       Asinh      => W AsinhOp.create

230 |       Atan       => W AtanOp.create

231 |       Atanh      => W AtanhOp.create

232 |       Cosh       => W CoshOp.create

233 |       Sinh       => W SinhOp.create

234 |       Erf        => W ErfOp.create

235 |       ErfInv     => W ErfInvOp.create

236 |       Square     => W SquareOp.create

237 |   iop (Convert dtype resultShape x) = do

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

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

240 |   iop (BitCastConvert dtype resultShape x) = do

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

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

243 |   iop (BinaryElementwise f lhs rhs) = do

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

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

247 |     where

249 |     data Wrap : Type where

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

252 |     (.create) : BinaryOp -> Wrap

253 |     (.create) = \case

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

255 |       Add => W AddOp.create

256 |       Div => W DivOp.create

257 |       Max => W MaxOp.create

258 |       Min => W MinOp.create

259 |       Mul => W MulOp.create

260 |       Pow => W PowOp.create

261 |       Rem => W RemOp.create

262 |       Sub => W SubtractOp.create

263 |       And => W AndOp.create

264 |       Or  => W OrOp.create

265 |       ShiftRightLogical => W ShiftRightLogicalOp.create

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

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

268 |     bodyT <- emplaceBlock $ getTrueBranch op

269 |     bodyF <- emplaceBlock $ getFalseBranch op

270 |     addArguments true bodyT

271 |     addArguments false bodyF

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

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

274 |     pure $ OpLike op

275 |   iop (While cond body inits) = do

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

277 |     op <- WhileOp.create builder uloc inits

278 |     cond' <- emplaceBlock $ getCond op

279 |     body' <- emplaceBlock $ getBody op

280 |     addArguments body body'

281 |     addArguments cond cond'

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

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

284 |     pure $ OpLike op

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

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

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

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

289 |     body' <- emplaceBlock $ getBody op

290 |     addArguments body body'

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

292 |     pure $ OpLike op

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

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

295 |   iop (DynamicSlice starts sizes x) = do

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

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

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

299 |   iop (Concat axis xs) = do

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

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

302 |   iop (Iota shape dtype dim) = do

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

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

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

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

307 |     resTy <- itype cache ctx resultType

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

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

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

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

312 |     f' <- emplaceBlock $ getComputation op

313 |     addArguments f f'

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

315 |     pure $ OpLike op

316 |   iop (Reshape dtype to x) = do

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

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

319 |   iop (Select pred true false) = do

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

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

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

323 |     comp' <- emplaceBlock $ getComparator op

324 |     addArguments comp comp'

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

326 |     pure $ OpLike op

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

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

329 |   iop (TriangularSolve a b lower) =

330 |     OpLike <$> TriangularSolveOp.create

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

332 |   iop (Rng state resultType) = do

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

334 |     OpLike <$> RngBitGeneratorOp.create

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

338 | export covering

339 | execute :

340 |   Device ->

341 |   Fn 0 ->

342 |   {outputs : _} ->

343 |   Vect outputs Xla.Shape ->

344 |   ErrIO $ Vect outputs Literal

345 | execute (MkDevice mlirCtx passManager api client) f shapes = do

346 |   uloc <- UnknownLoc.get mlirCtx

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

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

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

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

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

352 |   code <- cppString

353 |   version <- toString !getCurrentVersion

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

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

356 |   bimapEitherT PjrtErr id $ do

357 |     executableBuildOptions <- mkExecutableBuildOptions

358 |     compileOptions <- serializeAsString !(mkCompileOptions executableBuildOptions)

359 |     program <- mkPjrtProgram code

360 |     loadedExec <- pjrtClientCompile api client program compileOptions

361 |     delete program

362 |     delete code

363 |     delete compileOptions

364 |     delete executableBuildOptions

366 |     buffers <- pjrtLoadedExecutableExecute api loadedExec outputs

367 |     pjrtLoadedExecutableDestroy api loadedExec

369 |     for (zip buffers shapes) $ \(buffer, shape) => do

370 |       literal <- allocLiteral shape

371 |       -- is this pure?

372 |       -- note we can probably avoid the difficulties around async

373 |       -- by awaiting the event in pjrtBufferToHostBuffer, thus

374 |       -- making that function synchronous

375 |       event <- pjrtBufferToHostBuffer api buffer literal

376 |       pjrtEventAwait api event

377 |       pjrtEventDestroy api event

378 |       pjrtBufferDestroy api buffer

380 |       pure literal