0 | {--
  1 | Copyright (C) 2026  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.Passes
 18 |
 19 | import Control.Monad.State
 20 | import Data.SortedMap
 21 | import Data.Vect
 22 | import Data.List
 23 |
 24 | import Types
 25 | import DType
 26 | import Compiler.IR
 27 |
 28 | ||| In `(next, {(idx, tag) : idx', ...})`, `next` is the next available index into the function
 29 | ||| parameters, `(idx, tag)` point to the value in the unprocessed IR, and `idx'` indexes into the
 30 | ||| parameters of the function in the new IR
 31 | 0 Rebinds : Type
 32 | Rebinds = (Nat, SortedMap (Nat, Nat) Nat)
 33 |
 34 | parameters (fCurrent : Nat)
 35 |   covering
 36 |   rebindOp : Op -> State Rebinds Op
 37 |
 38 |   covering
 39 |   rebind : Value -> State Rebinds Value
 40 |   rebind $ V i (Concrete op) = V i <$> Concrete <$> rebindOp op
 41 |   rebind $ V i (BoundSet k) =
 42 |     if k >= fCurrent
 43 |     then pure (V i $ BoundSet k)
 44 |     else do
 45 |       (next, rebinds) <- get
 46 |       let op = BoundSet fCurrent
 47 |       case lookup (i, k) rebinds of
 48 |         Just i => pure $ V i op
 49 |         Nothing => do
 50 |           put (S next, insert (i, k) next rebinds)
 51 |           pure $ V next op
 52 |
 53 |   covering
 54 |   rebindFunc : Fn arity -> State Rebinds (Fn arity)
 55 |   rebindFunc (MkFn tag pt rt results env) = do
 56 |     results <- traverse rebind results
 57 |     ops <- traverse (\(t, x) => (t,) <$> rebindOp x) env.ops
 58 |     pure $ MkFn tag pt rt results ({ops := ops} env)
 59 |
 60 |   rebindOp (NamedFunc f) = pure $ NamedFunc f  -- NamedFunc assumes non-capturing
 61 |   rebindOp (CallByName tag resTy xs) = CallByName tag resTy <$> traverse rebind xs
 62 |   -- enzyme handles variable capture, so we don't remove captures from the differentiated function
 63 |   rebindOp (Grad shape f x) = Grad shape <$> rebindFunc f <*> rebind x
 64 |   rebindOp (Vectorize resTys shape name args) = Vectorize resTys shape name <$> traverse rebind args
 65 |   -- we don't rebind in HLO HOFs because capturing values is sometimes allowed, but capturing
 66 |   -- parameters is never allowed (except in While, but we don't mean to rebind While yet)
 67 |   rebindOp (Map f xs resultType ks) = pure $ Map f !(traverse rebind xs) resultType ks
 68 |   rebindOp (Reduce op inits axes xs) =
 69 |     pure $ Reduce op !(traverse rebind inits) axes !(traverse rebind xs)
 70 |   rebindOp (Sort comp axis isStable x) = pure $ Sort comp axis isStable !(rebind x)
 71 |   rebindOp (While condition body init) = While condition body <$> traverse rebind init
 72 |   rebindOp (If resultType p t f) = pure $ If resultType !(rebind p) t f
 73 |   rebindOp (BitCastConvert dtype ks x) = BitCastConvert dtype ks <$> rebind x
 74 |   rebindOp (Convert dtype ks x) = Convert dtype ks <$> rebind x
 75 |   rebindOp (Reshape dtype to x) = Reshape dtype to <$> rebind x
 76 |   rebindOp (Slice starts stops strides x) = Slice starts stops strides <$> rebind x
 77 |   rebindOp (DynamicSlice starts sizes x) =
 78 |     pure $ DynamicSlice !(traverse rebind starts) sizes !(rebind x)
 79 |   rebindOp (Concat axis xs) = Concat axis <$> traverse rebind xs
 80 |   rebindOp (Transpose ordering x) = Transpose ordering <$> rebind x
 81 |   rebindOp (Broadcast to x) = Broadcast to <$> rebind x
 82 |   rebindOp (Reverse axes x) = Reverse axes <$> rebind x
 83 |   rebindOp (BinaryElementwise op x y) = BinaryElementwise op <$> rebind x <*> rebind y
 84 |   rebindOp (UnaryElementwise op x) = UnaryElementwise op <$> rebind x
 85 |   rebindOp (Select p t f) = Select <$> rebind p <*> rebind t <*> rebind f
 86 |   rebindOp (DotGeneral lb lc rb rc resultType x y) =
 87 |     DotGeneral lb lc rb rc resultType <$> rebind x <*> rebind y
 88 |   rebindOp (Cholesky x) = Cholesky <$> rebind x
 89 |   rebindOp (TriangularSolve x y isLower) = pure $ TriangularSolve !(rebind x) !(rebind y) isLower
 90 |   rebindOp (Rng state resultType) = pure $ Rng !(rebind state) resultType
 91 |
 92 |   rebindOp x@(Lit _ _ _) = pure x
 93 |   rebindOp x@(MinValue _) = pure x
 94 |   rebindOp x@(MaxValue _) = pure x
 95 |   rebindOp x@MinFiniteFloat = pure x
 96 |   rebindOp x@MaxFiniteFloat = pure x
 97 |   rebindOp x@(Iota _ _ _) = pure x
 98 |
 99 | export covering
100 | removeCaptures : {n : _ } -> Fn n -> (m ** (Vect m Value, Fn (n + m)))
101 | removeCaptures (MkFn tag paramTypes resultTypes results (MkEnv count ops)) =
102 |   let ((_, rebinds), (results, ops)) = runState (n, empty) $ do
103 |         results <- traverse (rebind tag) results
104 |         ops <- traverse (\(tg, x) => (tg,) <$> rebindOp tag x) ops
105 |         pure (results, ops)
106 |
107 |       argify : (Nat, Nat) -> Value
108 |       argify (i, x) = V i $ BoundSet x
109 |
110 |       rebinds := the (Vect _ _) $ fst <$>
111 |         fromList (sortBy (compare `on` snd) (SortedMap.toList rebinds))
112 |       paramTypes := paramTypes ++ (uncurry TypeRef <$> rebinds)
113 |
114 |    in (_ ** (argify <$> rebinds, MkFn tag paramTypes resultTypes results (MkEnv count ops)))
115 |