
  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.ValueRange
 27 | import Compiler.FFI
 28 | 
 29 | ffi : String -> String
 30 | ffi = libxla "c/mlir/Dialect/Func/IR/FuncOps.h"
 31 | 
 32 | %foreign (ffi "MLIRContext_loadDialect_FuncDialect")
 33 | prim__loadDialectFuncDialect : AnyPtr -> PrimIO ()
 34 | 
 35 | export
 36 | loadDialectFuncDialect : HasIO io => MLIRContext -> io ()
 37 | loadDialectFuncDialect (MkMLIRContext ctx) = primIO $ prim__loadDialectFuncDialect ctx
 38 | 
 39 | public export
 40 | data FuncOp = MkFuncOp GCAnyPtr
 41 | 
 42 | %foreign (ffi "FuncOp_delete")
 43 | prim__deleteFuncOp : AnyPtr -> PrimIO ()
 44 | 
 45 | %foreign (ffi "FuncOp_create")
 46 | prim__funcOpCreate : GCAnyPtr -> String -> GCAnyPtr -> PrimIO AnyPtr
 47 | 
 48 | namespace FuncOp
 49 |   export
 50 |   create : HasIO io => Location -> String -> FunctionType -> io FuncOp
 51 |   create (MkLocation location) name (MkFunctionType type) = do
 52 |     op <- primIO $ prim__funcOpCreate location name type
 53 |     op <- onCollectAny' op (primIO . prim__deleteFuncOp)
 54 |     pure (MkFuncOp op)
 55 | 
 56 | %foreign (ffi "FuncOp_getOperation")
 57 | prim__funcOpGetOperation : GCAnyPtr -> PrimIO AnyPtr
 58 | 
 59 | export
 60 | Op FuncOp where
 61 |   getOperation (MkFuncOp op) = MkOperation <$> (primIO $ prim__funcOpGetOperation op)
 62 | 
 63 | %foreign (ffi "FuncOp_addEntryBlock")
 64 | prim__funcOpAddEntryBlock : GCAnyPtr -> PrimIO AnyPtr
 65 | 
 66 | export
 67 | addEntryBlock : HasIO io => FuncOp -> io Block
 68 | addEntryBlock (MkFuncOp op) = MkBlock <$> (primIO $ prim__funcOpAddEntryBlock op)
 69 | 
 70 | public export
 71 | data ReturnOp = MkReturnOp GCAnyPtr
 72 | 
 73 | %foreign (ffi "ReturnOp_delete")
 74 | prim__deleteReturnOp : AnyPtr -> PrimIO ()
 75 | 
 76 | %foreign (ffi "ReturnOp_create")
 77 | prim__returnOpCreate : GCAnyPtr -> GCAnyPtr -> GCAnyPtr -> PrimIO AnyPtr
 78 | 
 79 | namespace ReturnOp
 80 |   export
 81 |   create : HasIO io => OpBuilder -> Location -> ValueRange -> io ReturnOp
 82 |   create (MkOpBuilder builder) (MkLocation location) (MkValueRange results) = do
 83 |     op <- primIO $ prim__returnOpCreate builder location results
 84 |     op <- onCollectAny' op (primIO . prim__deleteReturnOp)
 85 |     pure (MkReturnOp op)
 86 |