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.Stablehlo.Dialect.StablehloAttrs
18 |
19 | import Compiler.FFI
20 | import Compiler.Array
21 | import Compiler.MLIR.IR.MLIRContext
22 |
23 | ffi : String -> String
24 | ffi = libxla "c/stablehlo/dialect/StablehloAttrs.h"
25 |
26 | %foreign (ffi "DotDimensionNumbersAttr_delete")
27 | prim__deleteDotDimensionNumbersAttr : AnyPtr -> PrimIO ()
28 |
29 | %foreign (ffi "DotDimensionNumbersAttr_get")
30 | prim__dotDimensionNumbersAttrGet :
31 |   AnyPtr ->
32 |   GCAnyPtr -> Bits64 ->
33 |   GCAnyPtr -> Bits64 ->
34 |   GCAnyPtr -> Bits64 ->
35 |   GCAnyPtr -> Bits64 ->
36 |   PrimIO AnyPtr
37 |
38 | public export
39 | data DotDimensionNumbersAttr = MkDotDimensionNumbersAttr GCAnyPtr
40 |
41 | namespace DotDimensionNumbersAttr
42 |   export
43 |   get :
44 |     HasIO io =>
45 |     MLIRContext ->
46 |     List Nat ->
47 |     List Nat ->
48 |     List Nat ->
49 |     List Nat ->
50 |     io DotDimensionNumbersAttr
51 |   get
52 |     (MkMLIRContext ctx)
53 |     lhsBatchingDimensions
54 |     rhsBatchingDimensions
55 |     lhsContractingDimensions
56 |     rhsContractingDimensions = do
57 |       MkArray lb lbLen <- fromList $ cast {to = Int64} <$> lhsBatchingDimensions
58 |       MkArray rb rbLen <- fromList $ cast {to = Int64} <$> rhsBatchingDimensions
59 |       MkArray lc lcLen <- fromList $ cast {to = Int64} <$> lhsContractingDimensions
60 |       MkArray rc rcLen <- fromList $ cast {to = Int64} <$> rhsContractingDimensions
61 |       attr <- primIO $ prim__dotDimensionNumbersAttrGet ctx lb lbLen rb rbLen lc lcLen rc rcLen
62 |       attr <- onCollectAny' attr (primIO . prim__deleteDotDimensionNumbersAttr)
63 |       pure (MkDotDimensionNumbersAttr attr)
64 |