0 | {--
  1 | Copyright (C) 2025  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.Enzyme.MLIR.Dialect.Ops
 18 |
 19 | import Compiler.FFI
 20 | import Compiler.Array
 21 | import Compiler.MLIR.IR.Builders
 22 | import Compiler.MLIR.IR.Location
 23 | import Compiler.MLIR.IR.MLIRContext
 24 | import Compiler.MLIR.IR.Operation
 25 | import Compiler.MLIR.IR.OpDefinition
 26 | import Compiler.MLIR.IR.Region
 27 | import Compiler.MLIR.IR.TypeRange
 28 | import Compiler.MLIR.IR.ValueRange
 29 |
 30 | ffi : String -> String
 31 | ffi = libxla "c/Enzyme/MLIR/Dialect/Ops.h"
 32 |
 33 | %foreign (ffi "AutoDiffRegionOp_delete")
 34 | prim__deleteAutoDiffRegionOp : AnyPtr -> PrimIO ()
 35 |
 36 | %foreign (ffi "AutoDiffRegionOp_create")
 37 | prim__autoDiffRegionOpCreate :
 38 |   GCAnyPtr -> AnyPtr -> GCAnyPtr -> GCAnyPtr -> GCAnyPtr -> Int -> Int -> PrimIO AnyPtr
 39 |
 40 | public export
 41 | data AutoDiffRegionOp = MkAutoDiffRegionOp GCAnyPtr
 42 |
 43 | public export
 44 | data Activity
 45 |   = EnzymeActive
 46 |   | EnzymeDup
 47 |   | EnzymeConst
 48 |   | EnzymeDupnoneed
 49 |   | EnzymeActivenoneed
 50 |   | EnzymeConstnoneed
 51 |
 52 | Cast Activity Int where
 53 |   cast = \case
 54 |     EnzymeActive       => 0
 55 |     EnzymeDup          => 1
 56 |     EnzymeConst        => 2
 57 |     EnzymeDupnoneed    => 3
 58 |     EnzymeActivenoneed => 4
 59 |     EnzymeConstnoneed  => 5
 60 |
 61 | namespace AutoDiffRegionOp
 62 |   export
 63 |   create :
 64 |     HasIO io =>
 65 |     OpBuilder ->
 66 |     MLIRContext ->
 67 |     Location ->
 68 |     TypeRange ->
 69 |     ValueRange ->
 70 |     Activity ->
 71 |     Activity ->
 72 |     io AutoDiffRegionOp
 73 |   create
 74 |     (MkOpBuilder builder)
 75 |     (MkMLIRContext ctx)
 76 |     (MkLocation location)
 77 |     (MkTypeRange outputs)
 78 |     (MkValueRange inputs)
 79 |     activity
 80 |     retActivity = do
 81 |       op <- primIO $ prim__autoDiffRegionOpCreate
 82 |         builder ctx location outputs inputs (cast activity) (cast retActivity)
 83 |       op <- onCollectAny' op (primIO . prim__deleteAutoDiffRegionOp)
 84 |       pure (MkAutoDiffRegionOp op)
 85 |
 86 | %foreign (ffi "AutoDiffRegionOp_getBody")
 87 | prim__autoDiffRegionOpGetBody : GCAnyPtr -> AnyPtr
 88 |
 89 | export
 90 | getBody : AutoDiffRegionOp -> Region
 91 | getBody (MkAutoDiffRegionOp op) = MkRegion $ prim__autoDiffRegionOpGetBody op
 92 |
 93 | %foreign (ffi "AutoDiffRegionOp_getOperation")
 94 | prim__autoDiffRegionOpGetOperation : GCAnyPtr -> PrimIO AnyPtr
 95 |
 96 | export
 97 | Op AutoDiffRegionOp where
 98 |   getOperation (MkAutoDiffRegionOp op) =
 99 |     MkOperation <$> (primIO $ prim__autoDiffRegionOpGetOperation op)
100 |
101 | %foreign (ffi "YieldOp_delete")
102 | prim__deleteYieldOp : AnyPtr -> PrimIO ()
103 |
104 | %foreign (ffi "YieldOp_create")
105 | prim__yieldOpCreate : GCAnyPtr -> GCAnyPtr -> GCAnyPtr -> PrimIO AnyPtr
106 |
107 | public export
108 | data YieldOp = MkYieldOp GCAnyPtr
109 |
110 | namespace YieldOp
111 |   export
112 |   create : HasIO io => OpBuilder -> Location -> ValueRange -> io YieldOp
113 |   create (MkOpBuilder builder) (MkLocation location) (MkValueRange operands) = do
114 |     op <- primIO $ prim__yieldOpCreate builder location operands
115 |     op <- onCollectAny' op (primIO . prim__deleteYieldOp)
116 |     pure (MkYieldOp op)
117 |
118 | %foreign (ffi "BatchOp_delete")
119 | prim__deleteBatchOp : AnyPtr -> PrimIO ()
120 |
121 | %foreign (ffi "BatchOp_create")
122 | prim__batchOpCreate :
123 |   GCAnyPtr ->
124 |   AnyPtr ->
125 |   GCAnyPtr ->
126 |   GCAnyPtr ->
127 |   String ->
128 |   GCAnyPtr ->
129 |   GCAnyPtr ->
130 |   Bits64 ->
131 |   PrimIO AnyPtr
132 |
133 | public export
134 | data BatchOp = MkBatchOp GCAnyPtr
135 |
136 | namespace BatchOp
137 |   export
138 |   create :
139 |     HasIO io =>
140 |     OpBuilder ->
141 |     MLIRContext ->
142 |     Location ->
143 |     TypeRange ->
144 |     String ->
145 |     ValueRange ->
146 |     List Nat ->
147 |     io BatchOp
148 |   create
149 |     (MkOpBuilder builder)
150 |     (MkMLIRContext ctx)
151 |     (MkLocation location)
152 |     (MkTypeRange outputs)
153 |     fn
154 |     (MkValueRange inputs)
155 |     batchShape = do
156 |       MkArray batchShape batchShapeLen <- fromList $ cast {to = Int64} <$> batchShape
157 |       op <- primIO $ prim__batchOpCreate
158 |         builder ctx location outputs fn inputs batchShape batchShapeLen
159 |       op <- onCollectAny' op (primIO . prim__deleteBatchOp)
160 |       pure (MkBatchOp op)
161 |
162 | %foreign (ffi "BatchOp_getOperation")
163 | prim__batchOpGetOperation : GCAnyPtr -> PrimIO AnyPtr
164 |
165 | export
166 | Op BatchOp where
167 |   getOperation (MkBatchOp op) = MkOperation <$> (primIO $ prim__batchOpGetOperation op)
168 |