
   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.IR
  18 | 
  19 | import Control.Monad.State
  20 | import Data.Primitives.Interpolation
  21 | import public Compiler.Stablehlo.Dialect.StablehloEnums
  22 | 
  23 | import Derive.Prelude
  24 | import Language.Reflection
  25 | 
  26 | import Compiler.LiteralRW
  27 | import Compiler.Xla.XlaData
  28 | import Literal
  29 | import DType
  30 | import Types
  31 | import Util
  32 | 
  33 | %language ElabReflection
  34 | 
  35 | Show a => Interpolation (List a) where
  36 |   interpolate = show
  37 | 
  38 | public export
  39 | data ValueType
  40 |   = ||| A concrete tensor type
  41 |     TensorType Shape DType
  42 | 
  43 |   | ||| Points to the type of the value at this index and tag
  44 |     TypeRef Nat Nat
  45 | 
  46 | Show ValueType where
  47 |   show (TensorType shape dtype) = "\{show shape} \{show dtype}"
  48 |   show (TypeRef idx tag) = "type of \{show idx} of \{show tag}"
  49 | 
  50 | public export data Op : Type
  51 | 
  52 | -- we use `List (Nat, Op)` for O(1) append (all we do when building the graph is append)
  53 | -- we can't use `(Nat, List Op)`, or even better `(n ** Vect n Op)`, because we don't handle
  54 | -- scoping properly so node pointers aren't contiguous and don't match list indices
  55 | public export
  56 | record Env where
  57 |   constructor MkEnv
  58 | 
  59 |   ||| The global counter
  60 |   counter : Nat
  61 | 
  62 |   ||| Local cached ops
  63 |   ops : List (Nat, Op)
  64 | 
  65 | export
  66 | empty : Env
  67 | empty = MkEnv 1 []  -- root function takes 0
  68 | 
  69 | export
  70 | emptyFrom : Env -> Env
  71 | emptyFrom (MkEnv n _) = MkEnv n []
  72 | 
  73 | export
  74 | updateCounterFrom : Env -> State Env ()
  75 | updateCounterFrom (MkEnv n _) = do
  76 |   MkEnv _ xs <- get
  77 |   put $ MkEnv n xs
  78 | 
  79 | public export data OpRef : Type
  80 | 
  81 | public export
  82 | data Value = V Nat OpRef
  83 | 
  84 | ||| An anonymous function. Approximates an MLIR `Region` or `Block` (these are somewhat synonymous
  85 | ||| in spidr since all regions have exactly one block). `tag` labels the parameter set.
  86 | public export
  87 | record Fn (arity : Nat) where
  88 |   constructor MkFn
  89 |   tag : Nat
  90 |   paramTypes : Vect arity ValueType
  91 |   resultTypes : Vect resultCount ValueType
  92 |   results : Vect resultCount Value
  93 |   env : Env
  94 | 
  95 | public export
  96 | data BinaryOp =
  97 |   Compare ComparisonDirection | And | Or | Add | Sub | Mul | Div | Rem | Pow | Min | Max
  98 |   | ShiftRightLogical
  99 | 
 100 | %runElab derive "ComparisonDirection" [Show]
 101 | %runElab derive "BinaryOp" [Show]
 102 | 
 103 | public export
 104 | data UnaryOp =
 105 |     Not | Neg | Ceil | Floor | Abs | Log | Exp | Logistic | Sqrt | Sin | Cos | Tan | Tanh
 106 |   | Erf | ErfInv | Square | Asin | Acos | Atan | Sinh | Cosh | Asinh | Acosh | Atanh
 107 | 
 108 | %runElab derive "UnaryOp" [Show]
 109 | 
 110 | public export
 111 | data Op : Type where
 112 |   ||| Corresponds approximately to a FuncOp. The FuncOp's name is determined by the `Fn`s tag.
 113 |   ||| We use the tag to correspond to both the parameter set and the function itself. I think this
 114 |   ||| is OK because the function determines the parameter set (notably we don't introduce any
 115 |   ||| regions with multiple blocks, so each FuncOp only has one block).
 116 |   |||
 117 |   ||| The function is assumed **not** to capture variables, see `Passes.removeCaptures`.
 118 |   NamedFunc : Fn arity -> Op
 119 | 
 120 |   ||| Call a named function, by name.
 121 |   |||
 122 |   ||| The named functions must be listed in the `Env`, else it will not be interpreted.
 123 |   CallByName : (name : Nat) -> (resultTypes : List ValueType) -> List Value -> Op
 124 | 
 125 |   Lit : (shape : Shape) -> (dtype : DType) -> Literal shape (idrisType dtype) -> Op
 126 |   Grad : Shape -> Fn 1 -> Value -> Op
 127 |   MinValue : DType -> Op
 128 |   MaxValue : DType -> Op
 129 |   MinFiniteFloat : Op
 130 |   MaxFiniteFloat : Op
 131 |   Iota : (shape : Shape) -> DType -> (axis : Nat) -> Op
 132 |   BitCastConvert : DType -> Shape -> Value -> Op
 133 |   Convert : DType -> Shape -> Value -> Op
 134 |   Reshape : DType -> Shape -> Value -> Op
 135 |   Slice : (starts, stops, strides : List Nat) -> Value -> Op
 136 |   DynamicSlice : (starts : List Value) -> (sizes : List Nat) -> Value -> Op
 137 |   Concat : (axis : Nat) -> Vect (S n) Value -> Op
 138 |   Transpose : (ordering : List Nat) -> Value -> Op
 139 |   Broadcast : DType -> (from, to : Shape) -> Value -> Op
 140 |   Map : Fn arity -> Vect arity Value -> (resultType : ValueType) -> Shape -> Op
 141 |   Reduce : Fn (n + n) -> (inits : Vect n Value) -> (axes : List Nat) -> Vect n Value -> Op
 142 |   Sort : Fn 2 -> (axis : Nat) -> (isStable : Bool) -> Value -> Op
 143 |   Reverse : (axes : List Nat) -> Value -> Op
 144 |   BinaryElementwise : BinaryOp -> Value -> Value -> Op
 145 |   UnaryElementwise : UnaryOp -> Value -> Op
 146 |   Select : (predicate, onTrue, onFalse : Value) -> Op
 147 |   While : (condition, body : Fn n) -> (init : Vect n Value) -> Op
 148 |   If : (resultType : ValueType) -> (predicate : Value) -> (onTrue, onFalse : Fn 0) -> Op
 149 |   DotGeneral :
 150 |     (lBatch, lContract, rBatch, rContract: List Nat) ->
 151 |     (resultType : ValueType) ->
 152 |     Value ->
 153 |     Value ->
 154 |     Op
 155 |   Cholesky : Value -> Op
 156 |   TriangularSolve : Value -> Value -> (isLower : Bool) -> Op
 157 |   Rng : (state : Value) -> (resultType : ValueType) -> Op
 158 | 
 159 | public export
 160 | data OpRef = BoundSet Nat | Concrete Op
 161 | 
 162 | export
 163 | tagOpRef : Monad m => OpRef -> StateT Env m OpRef
 164 | tagOpRef (BoundSet x) = pure $ BoundSet x
 165 | tagOpRef (Concrete expr) = do
 166 |   MkEnv next env <- get
 167 |   put $ MkEnv (S next) ((next, expr) :: env)
 168 |   pure (BoundSet next)
 169 | 
 170 | export
 171 | reserve : State Env Nat
 172 | reserve = do
 173 |   MkEnv next env <- get
 174 |   put $ MkEnv (S next) env
 175 |   pure next
 176 | 
 177 | covering
 178 | showOp : Nat -> Op -> String
 179 | 
 180 | covering
 181 | showOpRef : Nat -> OpRef -> String
 182 | showOpRef indent (BoundSet k) = "Bound \{k}"
 183 | showOpRef indent (Concrete x) = showOp indent x
 184 | 
 185 | covering
 186 | showValue : Nat -> Value -> String
 187 | showValue indent (V idx op) = "(\{showOpRef indent op}):\{show idx}"
 188 | 
 189 | covering
 190 | showValueList : Traversable t => Nat -> t Value -> String
 191 | showValueList indent xs = "[" ++ joinBy ", " (toList $ map (showValue indent) xs) ++ "]"
 192 | 
 193 | covering
 194 | showEnv : Nat -> Env -> String
 195 | showEnv indent (MkEnv max env) = joinBy "\n" $ assert_total $ map fmt (reverse env)
 196 | 
 197 |   where
 198 | 
 199 |   fmt : (Nat, Op) -> String
 200 |   fmt (n, x) =
 201 |     let sep = replicate (4 + length (show max) `minus` length (show n)) ' '
 202 |      in "\{replicate indent ' '}\{show n}\{sep}\{showOp indent x}"
 203 | 
 204 | covering
 205 | showFn : Nat -> Fn arity -> String
 206 | showFn indent (MkFn parameterSetTag paramTypes resultTypes results env@(MkEnv _ env')) =
 207 |   let params = "\{show parameterSetTag} \{show paramTypes}"
 208 |       res = "\{showValueList (indent + 2) $ toList results}" in
 209 |   case env' of
 210 |     [] => "\{params} => \{res}"
 211 |     _  =>
 212 |       "\{params} => \{res} with vars {\n\{showEnv (indent + 4) env}\n\{replicate (indent + 2) ' '}}"
 213 | 
 214 | export Show (Fn arity) where show = assert_total $ showFn 0
 215 | 
 216 | showOp indent (NamedFunc f) = "NamedFunc \{showFn indent f}"
 217 | showOp indent (CallByName fTag _ xs) = "Call {name = \{show fTag}} \{showValueList indent xs}"
 218 | showOp indent (Lit shape dtype x) = "Lit \{shape} \{show dtype}"
 219 | showOp indent (Grad _ op x) = "Grad {op = \{showFn indent op}} \{showValue indent x}"
 220 | showOp _      (MinValue dtype) = "MinValue \{show dtype}"
 221 | showOp _      (MaxValue dtype) = "MaxValue \{show dtype}"
 222 | showOp _      MinFiniteFloat = "MinFiniteFloat"
 223 | showOp _      MaxFiniteFloat = "MaxFiniteFloat"
 224 | showOp indent (Iota dtype shape axis) =
 225 |   "Iota {shape = \{show shape}, dtype = \{show dtype}, axis = \{axis}}"
 226 | showOp indent (Convert dtype shape x) =
 227 |   "Convert {dtype = \{show dtype}} \{showValue indent x}"
 228 | showOp indent (BitCastConvert dtype shape x) =
 229 |   "BitCastConvert {dtype = \{show dtype}} \{showValue indent x}"
 230 | showOp indent (Reshape _ to x) = "Reshape {to = \{to}} \{showValue indent x}"
 231 | showOp indent (Slice starts stops strides x) =
 232 |   "Slice {starts = \{starts}, stops = \{stops}, strides = \{strides}} \{showValue indent x}"
 233 | showOp indent (DynamicSlice starts sizes x) =
 234 |   "DynamicSlice {starts = \{showValueList indent starts}, sizes = \{sizes}} \{showValue indent x}"
 235 | showOp indent (Concat axis xs) = "Concat {axis = \{axis}} \{showValueList indent $ toList xs}"
 236 | showOp indent (Transpose ordering x) = "Transpose {ordering = \{ordering}} \{showValue indent x}"
 237 | showOp indent (Broadcast _ from to x) =
 238 |   "Broadcast {from = \{from}, to = \{to}} \{showValue indent x}"
 239 | showOp indent (Map f xs _ _) = "Map {f = \{showFn indent f}} \{showValueList indent $ toList xs}"
 240 | showOp indent (Reduce op neutrals axes xs) =
 241 |   "Reduce {op = \{showFn indent op}, inits = \{showValueList indent $ toList neutrals}," ++
 242 |     " axes = \{axes}} \{showValueList indent $ toList xs}"
 243 | showOp indent (Sort f axis _ xs) =
 244 |   "Sort {f = \{showFn indent f}, axis = \{axis}} \{showValue indent xs}"
 245 | showOp indent (Reverse axes x) = "Reverse \{axes} \{showValue indent x}"
 246 | showOp indent (BinaryElementwise op x y) = "\{show op} \{showValue indent x} \{showValue indent y}"
 247 | showOp indent (UnaryElementwise op x) = "\{show op} \{showValue indent x}"
 248 | showOp indent (Select p t f) =
 249 |   "Select {predicate = \{showValue indent p}, onTrue = \{showValue indent t}," ++
 250 |     " onFalse = \{showValue indent f}}"
 251 | showOp indent (While c b is) =
 252 |   "While {condition = \{showFn indent c}, body = \{showFn indent b}," ++
 253 |     " initials = \{showValueList indent $ toList is}}"
 254 | showOp indent (If _ p ft ff) =
 255 |   "If {predicate = \{showValue indent p}, onTrue = \{showFn indent ft}," ++
 256 |     " onFalse = \{showFn indent ff}}"
 257 | showOp indent (DotGeneral lBatch lContract rBatch rContract _ x y) =
 258 |   "DotGeneral {lBatch = \{lBatch}, lContract = \{lContract}," ++
 259 |     " rBatch = \{rBatch}, rContract = \{rContract}} \{showValue indent x} \{showValue indent y}"
 260 | showOp indent (Cholesky x) = "Cholesky \{showValue indent x}"
 261 | showOp indent (TriangularSolve x y isLower) =
 262 |   "TriangularSolve {isLower = \{show isLower}} \{showValue indent x} \{showValue indent y}"
 263 | showOp indent (Rng state shape) = "Rng {state = \{showValue indent state}, shape = \{show shape}}"
 264 |