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