
   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
  96 |   | ShiftRightLogical
  97 | 
  98 | %runElab derive "ComparisonDirection" [Show]
  99 | %runElab derive "BinaryOp" [Show]
 100 | 
 101 | public export
 102 | data UnaryOp =
 103 |     Not | Neg | Ceil | Floor | Abs | Log | Exp | Logistic | Sqrt | Sin | Cos | Tan | Tanh
 104 |   | Erf | ErfInv | Square | Asin | Acos | Atan | Sinh | Cosh | Asinh | Acosh | Atanh
 105 | 
 106 | %runElab derive "UnaryOp" [Show]
 107 | 
 108 | public export
 109 | data BroadcastShape = Explicit Shape | AddLeading (List Nat)
 110 | 
 111 | public export
 112 | data Op : Type where
 113 |   ||| Corresponds approximately to a FuncOp. The FuncOp's name is determined by the `Fn`s tag.
 114 |   ||| We use the tag to correspond to both the parameter set and the function itself. I think this
 115 |   ||| is OK because the function determines the parameter set (notably we don't introduce any
 116 |   ||| regions with multiple blocks, so each FuncOp only has one block).
 117 |   |||
 118 |   ||| The function is assumed **not** to capture variables, see `Passes.removeCaptures`.
 119 |   NamedFunc : Fn arity -> Op
 120 | 
 121 |   ||| Call a named function, by name.
 122 |   |||
 123 |   ||| The named functions must be listed in the `Env`, else it will not be interpreted.
 124 |   CallByName : (name : Nat) -> (resultTypes : List ValueType) -> List Value -> Op
 125 | 
 126 |   Lit : (shape : Shape) -> (dtype : DType) -> Literal shape (idrisType dtype) -> Op
 127 |   Grad : Shape -> Fn 1 -> Value -> Op
 128 |   Vectorize :
 129 |     (resultTypes : List ValueType) ->
 130 |     (batchShape : Shape) ->
 131 |     (targetFunc : Nat) ->
 132 |     (args : List Value) ->
 133 |     Op
 134 |   MinValue : DType -> Op
 135 |   MaxValue : DType -> Op
 136 |   MinFiniteFloat : Op
 137 |   MaxFiniteFloat : Op
 138 |   Iota : (shape : Shape) -> DType -> (axis : Nat) -> Op
 139 |   BitCastConvert : DType -> Shape -> Value -> Op
 140 |   Convert : DType -> Shape -> Value -> Op
 141 |   Reshape : DType -> Shape -> Value -> Op
 142 |   Slice : (starts, stops, strides : List Nat) -> Value -> Op
 143 |   DynamicSlice : (starts : List Value) -> (sizes : List Nat) -> Value -> Op
 144 |   DynamicUpdateSlice : (operand, update : Value) -> (starts : List Value) -> Op
 145 |   Concat : (axis : Nat) -> Vect (S n) Value -> Op
 146 |   Transpose : (ordering : List Nat) -> Value -> Op
 147 |   Broadcast : BroadcastShape -> Value -> Op
 148 |   Reduce : Fn (n + n) -> (inits : Vect n Value) -> (axes : List Nat) -> Vect n Value -> Op
 149 |   Sort : Fn 2 -> (axis : Nat) -> (isStable : Bool) -> Value -> Op
 150 |   Reverse : (axes : List Nat) -> Value -> Op
 151 |   BinaryElementwise : BinaryOp -> Value -> Value -> Op
 152 |   UnaryElementwise : UnaryOp -> Value -> Op
 153 |   Select : (predicate, onTrue, onFalse : Value) -> Op
 154 |   While : (condition, body : Fn n) -> (init : Vect n Value) -> Op
 155 |   If : (resultType : ValueType) -> (predicate : Value) -> (onTrue, onFalse : Fn 0) -> Op
 156 |   DotGeneral :
 157 |     (lBatch, lContract, rBatch, rContract: List Nat) ->
 158 |     (resultType : ValueType) ->
 159 |     Value ->
 160 |     Value ->
 161 |     Op
 162 |   Cholesky : Value -> Op
 163 |   TriangularSolve : Value -> Value -> (isLower : Bool) -> Op
 164 |   Rng : (state : Value) -> (resultType : ValueType) -> Op
 165 | 
 166 | public export
 167 | data OpRef = BoundSet Nat | Concrete Op
 168 | 
 169 | export
 170 | tagOpRef : Monad m => OpRef -> StateT Env m OpRef
 171 | tagOpRef (BoundSet x) = pure $ BoundSet x
 172 | tagOpRef (Concrete expr) = do
 173 |   MkEnv next env <- get
 174 |   put $ MkEnv (S next) ((next, expr) :: env)
 175 |   pure (BoundSet next)
 176 | 
 177 | export
 178 | reserve : State Env Nat
 179 | reserve = do
 180 |   MkEnv next env <- get
 181 |   put $ MkEnv (S next) env
 182 |   pure next
 183 | 
 184 | covering
 185 | showOp : Nat -> Op -> String
 186 | 
 187 | covering
 188 | showOpRef : Nat -> OpRef -> String
 189 | showOpRef indent (BoundSet k) = "Bound \{k}"
 190 | showOpRef indent (Concrete x) = showOp indent x
 191 | 
 192 | covering
 193 | showValue : Nat -> Value -> String
 194 | showValue indent (V idx op) = "(\{showOpRef indent op}):\{show idx}"
 195 | 
 196 | covering
 197 | showValueList : Traversable t => Nat -> t Value -> String
 198 | showValueList indent xs = "[" ++ joinBy ", " (toList $ map (showValue indent) xs) ++ "]"
 199 | 
 200 | covering
 201 | showEnv : Nat -> Env -> String
 202 | showEnv indent (MkEnv max env) = joinBy "\n" $ assert_total $ map fmt (reverse env)
 203 | 
 204 |   where
 205 | 
 206 |   fmt : (Nat, Op) -> String
 207 |   fmt (n, x) =
 208 |     let sep = replicate (4 + length (show max) `minus` length (show n)) ' '
 209 |      in "\{replicate indent ' '}\{show n}\{sep}\{showOp indent x}"
 210 | 
 211 | covering
 212 | showFn : Nat -> Fn arity -> String
 213 | showFn indent (MkFn parameterSetTag paramTypes resultTypes results env@(MkEnv _ env')) =
 214 |   let params = "\{show parameterSetTag} \{show paramTypes}"
 215 |       res = "\{showValueList (indent + 2) $ toList results}" in
 216 |   case env' of
 217 |     [] => "\{params} => \{res}"
 218 |     _  =>
 219 |       "\{params} => \{res} with vars {\n\{showEnv (indent + 4) env}\n\{replicate (indent + 2) ' '}}"
 220 | 
 221 | export Show (Fn arity) where show = assert_total $ showFn 0
 222 | 
 223 | showOp indent (NamedFunc f) = "NamedFunc \{showFn indent f}"
 224 | showOp indent (CallByName fTag _ xs) = "Call {targetFunc = \{show fTag}} \{showValueList indent xs}"
 225 | showOp indent (Lit shape dtype x) = "Lit \{shape} \{show dtype}"
 226 | showOp indent (Grad _ op x) = "Grad {op = \{showFn indent op}} \{showValue indent x}"
 227 | showOp indent (Vectorize _ batchShape fTag xs) =
 228 |   "Vectorize {batchShape = \{show batchShape}, targetFunc = \{show fTag}}"
 229 |     ++ " \{showValueList indent xs}"
 230 | showOp _      (MinValue dtype) = "MinValue \{show dtype}"
 231 | showOp _      (MaxValue dtype) = "MaxValue \{show dtype}"
 232 | showOp _      MinFiniteFloat = "MinFiniteFloat"
 233 | showOp _      MaxFiniteFloat = "MaxFiniteFloat"
 234 | showOp indent (Iota dtype shape axis) =
 235 |   "Iota {shape = \{show shape}, dtype = \{show dtype}, axis = \{axis}}"
 236 | showOp indent (Convert dtype shape x) =
 237 |   "Convert {dtype = \{show dtype}} \{showValue indent x}"
 238 | showOp indent (BitCastConvert dtype shape x) =
 239 |   "BitCastConvert {dtype = \{show dtype}} \{showValue indent x}"
 240 | showOp indent (Reshape _ to x) = "Reshape {to = \{to}} \{showValue indent x}"
 241 | showOp indent (Slice starts stops strides x) =
 242 |   "Slice {starts = \{starts}, stops = \{stops}, strides = \{strides}} \{showValue indent x}"
 243 | showOp indent (DynamicSlice starts sizes x) =
 244 |   "DynamicSlice {starts = \{showValueList indent starts}, sizes = \{sizes}} \{showValue indent x}"
 245 | showOp indent (DynamicUpdateSlice x update starts) =
 246 |   "DynamicSlice {update = \{showValue indent update}, starts = \{showValueList indent starts}}"
 247 |     ++ " \{showValue indent x}"
 248 | showOp indent (Concat axis xs) = "Concat {axis = \{axis}} \{showValueList indent $ toList xs}"
 249 | showOp indent (Transpose ordering x) = "Transpose {ordering = \{ordering}} \{showValue indent x}"
 250 | showOp indent (Broadcast bs x) =
 251 |   let bs : String = case bs of
 252 |         Explicit shape => "to = \{show shape}"
 253 |         AddLeading lead => "withLeading = \{show lead}"
 254 |    in "Broadcast {\{bs}} \{showValue indent x}"
 255 | showOp indent (Reduce op neutrals axes xs) =
 256 |   "Reduce {op = \{showFn indent op}, inits = \{showValueList indent $ toList neutrals}," ++
 257 |     " axes = \{axes}} \{showValueList indent $ toList xs}"
 258 | showOp indent (Sort f axis _ xs) =
 259 |   "Sort {f = \{showFn indent f}, axis = \{axis}} \{showValue indent xs}"
 260 | showOp indent (Reverse axes x) = "Reverse \{axes} \{showValue indent x}"
 261 | showOp indent (BinaryElementwise op x y) = "\{show op} \{showValue indent x} \{showValue indent y}"
 262 | showOp indent (UnaryElementwise op x) = "\{show op} \{showValue indent x}"
 263 | showOp indent (Select p t f) =
 264 |   "Select {predicate = \{showValue indent p}, onTrue = \{showValue indent t}," ++
 265 |     " onFalse = \{showValue indent f}}"
 266 | showOp indent (While c b is) =
 267 |   "While {condition = \{showFn indent c}, body = \{showFn indent b}," ++
 268 |     " initials = \{showValueList indent $ toList is}}"
 269 | showOp indent (If _ p ft ff) =
 270 |   "If {predicate = \{showValue indent p}, onTrue = \{showFn indent ft}," ++
 271 |     " onFalse = \{showFn indent ff}}"
 272 | showOp indent (DotGeneral lBatch lContract rBatch rContract _ x y) =
 273 |   "DotGeneral {lBatch = \{lBatch}, lContract = \{lContract}," ++
 274 |     " rBatch = \{rBatch}, rContract = \{rContract}} \{showValue indent x} \{showValue indent y}"
 275 | showOp indent (Cholesky x) = "Cholesky \{showValue indent x}"
 276 | showOp indent (TriangularSolve x y isLower) =
 277 |   "TriangularSolve {isLower = \{show isLower}} \{showValue indent x} \{showValue indent y}"
 278 | showOp indent (Rng state shape) = "Rng {state = \{showValue indent state}, shape = \{show shape}}"
 279 |