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.MLIR.Dialect.Func.IR.FuncOps
 18 |
 19 | import Compiler.MLIR.IR.Block
 20 | import Compiler.MLIR.IR.Builders
 21 | import Compiler.MLIR.IR.BuiltinTypes
 22 | import Compiler.MLIR.IR.Location
 23 | import Compiler.MLIR.IR.MLIRContext
 24 | import Compiler.MLIR.IR.OpDefinition
 25 | import Compiler.MLIR.IR.Operation
 26 | import Compiler.MLIR.IR.TypeRange
 27 | import Compiler.MLIR.IR.ValueRange
 28 | import Compiler.FFI
 29 |
 30 | ffi : String -> String
 31 | ffi = libxla "c/mlir/Dialect/Func/IR/FuncOps.h"
 32 |
 33 | %foreign (ffi "MLIRContext_loadDialect_FuncDialect")
 34 | prim__loadDialectFuncDialect : AnyPtr -> PrimIO ()
 35 |
 36 | export
 37 | loadDialectFuncDialect : HasIO io => MLIRContext -> io ()
 38 | loadDialectFuncDialect (MkMLIRContext ctx) = primIO $ prim__loadDialectFuncDialect ctx
 39 |
 40 | public export
 41 | data CallOp = MkCallOp GCAnyPtr
 42 |
 43 | %foreign (ffi "CallOp_delete")
 44 | prim__deleteCallOp : AnyPtr -> PrimIO ()
 45 |
 46 | %foreign (ffi "CallOp_create")
 47 | prim__callOpCreate : GCAnyPtr -> GCAnyPtr -> String -> GCAnyPtr -> GCAnyPtr -> PrimIO AnyPtr
 48 |
 49 | namespace CallOp
 50 |   export
 51 |   create : HasIO io => OpBuilder -> Location -> String -> TypeRange -> ValueRange -> io CallOp
 52 |   create (MkOpBuilder builder) (MkLocation location) name (MkTypeRange results) (MkValueRange operands) = do
 53 |     op <- primIO $ prim__callOpCreate builder location name results operands
 54 |     op <- onCollectAny' op (primIO . prim__deleteCallOp)
 55 |     pure (MkCallOp op)
 56 |
 57 | %foreign (ffi "CallOp_getOperation")
 58 | prim__callOpGetOperation : GCAnyPtr -> PrimIO AnyPtr
 59 |
 60 | export
 61 | Op CallOp where
 62 |   getOperation (MkCallOp op) = MkOperation <$> (primIO $ prim__callOpGetOperation op)
 63 |
 64 | public export
 65 | data FuncOp = MkFuncOp GCAnyPtr
 66 |
 67 | %foreign (ffi "FuncOp_delete")
 68 | prim__deleteFuncOp : AnyPtr -> PrimIO ()
 69 |
 70 | %foreign (ffi "FuncOp_create")
 71 | prim__funcOpCreate : GCAnyPtr -> String -> GCAnyPtr -> PrimIO AnyPtr
 72 |
 73 | namespace FuncOp
 74 |   export
 75 |   create : HasIO io => Location -> String -> FunctionType -> io FuncOp
 76 |   create (MkLocation location) name (MkFunctionType type) = do
 77 |     op <- primIO $ prim__funcOpCreate location name type
 78 |     op <- onCollectAny' op (primIO . prim__deleteFuncOp)
 79 |     pure (MkFuncOp op)
 80 |
 81 | %foreign (ffi "FuncOp_getOperation")
 82 | prim__funcOpGetOperation : GCAnyPtr -> PrimIO AnyPtr
 83 |
 84 | export
 85 | Op FuncOp where
 86 |   getOperation (MkFuncOp op) = MkOperation <$> (primIO $ prim__funcOpGetOperation op)
 87 |
 88 | %foreign (ffi "FuncOp_addEntryBlock")
 89 | prim__funcOpAddEntryBlock : GCAnyPtr -> PrimIO AnyPtr
 90 |
 91 | export
 92 | addEntryBlock : HasIO io => FuncOp -> io Block
 93 | addEntryBlock (MkFuncOp op) = MkBlock <$> (primIO $ prim__funcOpAddEntryBlock op)
 94 |
 95 | public export
 96 | data ReturnOp = MkReturnOp GCAnyPtr
 97 |
 98 | %foreign (ffi "ReturnOp_delete")
 99 | prim__deleteReturnOp : AnyPtr -> PrimIO ()
100 |
101 | %foreign (ffi "ReturnOp_create")
102 | prim__returnOpCreate : GCAnyPtr -> GCAnyPtr -> GCAnyPtr -> PrimIO AnyPtr
103 |
104 | namespace ReturnOp
105 |   export
106 |   create : HasIO io => OpBuilder -> Location -> ValueRange -> io ReturnOp
107 |   create (MkOpBuilder builder) (MkLocation location) (MkValueRange results) = do
108 |     op <- primIO $ prim__returnOpCreate builder location results
109 |     op <- onCollectAny' op (primIO . prim__deleteReturnOp)
110 |     pure (MkReturnOp op)
111 |