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