Skip to content

Commit 40c7708

Browse files
github-actions[bot]enzyme-ci-bot[bot]mofeing
authored
Regenerate MLIR Bindings (#680)
* Regenerate MLIR Bindings * fix docs --------- Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com> Co-authored-by: Sergio Sánchez Ramírez <sergio.sanchez.ramirez+git@bsc.es>
1 parent fe8ba34 commit 40c7708

File tree

3 files changed

+250
-0
lines changed

3 files changed

+250
-0
lines changed

‎docs/src/.vitepress/config.mts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ export default defineConfig({
8787
{ text: "TPU", link: "/api/tpu" },
8888
{ text: "Triton", link: "/api/triton" },
8989
{ text: "Shardy", link: "/api/shardy" },
90+
{ text: "MPI", link: "/api/mpi" },
9091
],
9192
},
9293
{
@@ -147,6 +148,7 @@ export default defineConfig({
147148
{ text: "TPU", link: "/api/tpu" },
148149
{ text: "Triton", link: "/api/triton" },
149150
{ text: "Shardy", link: "/api/shardy" },
151+
{ text: "MPI", link: "/api/mpi" },
150152
],
151153
},
152154
{

‎docs/src/api/mpi.md‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
```@meta
2+
CollapsedDocStrings = true
3+
```
4+
5+
# MPI Dialect
6+
7+
Refer to the [official documentation](https://mlir.llvm.org/docs/Dialects/MPI/) for
8+
more details.
9+
10+
```@autodocs
11+
Modules = [Reactant.MLIR.Dialects.mpi]
12+
```

‎src/mlir/Dialects/MPI.jl‎

Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
module mpi
2+
using ...IR
3+
import ...IR:
4+
NamedAttribute,
5+
Value,
6+
Location,
7+
Block,
8+
Region,
9+
Attribute,
10+
create_operation,
11+
context,
12+
IndexType
13+
import ..Dialects: namedattribute, operandsegmentsizes
14+
import ...API
15+
16+
"""
17+
`comm_rank`
18+
19+
Communicators other than `MPI_COMM_WORLD` are not supported for now.
20+
21+
This operation can optionally return an `!mpi.retval` value that can be used
22+
to check for errors.
23+
"""
24+
function comm_rank(;
25+
retval=nothing::Union{Nothing,IR.Type}, rank::IR.Type, location=Location()
26+
)
27+
op_ty_results = IR.Type[rank,]
28+
operands = Value[]
29+
owned_regions = Region[]
30+
successors = Block[]
31+
attributes = NamedAttribute[]
32+
!isnothing(retval) && push!(op_ty_results, retval)
33+
34+
return create_operation(
35+
"mpi.comm_rank",
36+
location;
37+
operands,
38+
owned_regions,
39+
successors,
40+
attributes,
41+
results=op_ty_results,
42+
result_inference=false,
43+
)
44+
end
45+
46+
"""
47+
`error_class`
48+
49+
`MPI_Error_class` maps return values from MPI calls to a set of well-known
50+
MPI error classes.
51+
"""
52+
function error_class(val::Value; errclass::IR.Type, location=Location())
53+
op_ty_results = IR.Type[errclass,]
54+
operands = Value[val,]
55+
owned_regions = Region[]
56+
successors = Block[]
57+
attributes = NamedAttribute[]
58+
59+
return create_operation(
60+
"mpi.error_class",
61+
location;
62+
operands,
63+
owned_regions,
64+
successors,
65+
attributes,
66+
results=op_ty_results,
67+
result_inference=false,
68+
)
69+
end
70+
71+
"""
72+
`finalize`
73+
74+
This function cleans up the MPI state. Afterwards, no MPI methods may
75+
be invoked (excpet for MPI_Get_version, MPI_Initialized, and MPI_Finalized).
76+
Notably, MPI_Init cannot be called again in the same program.
77+
78+
This operation can optionally return an `!mpi.retval` value that can be used
79+
to check for errors.
80+
"""
81+
function finalize(; retval=nothing::Union{Nothing,IR.Type}, location=Location())
82+
op_ty_results = IR.Type[]
83+
operands = Value[]
84+
owned_regions = Region[]
85+
successors = Block[]
86+
attributes = NamedAttribute[]
87+
!isnothing(retval) && push!(op_ty_results, retval)
88+
89+
return create_operation(
90+
"mpi.finalize",
91+
location;
92+
operands,
93+
owned_regions,
94+
successors,
95+
attributes,
96+
results=op_ty_results,
97+
result_inference=false,
98+
)
99+
end
100+
101+
"""
102+
`init`
103+
104+
This operation must preceed most MPI calls (except for very few exceptions,
105+
please consult with the MPI specification on these).
106+
107+
Passing &argc, &argv is not supported currently.
108+
109+
This operation can optionally return an `!mpi.retval` value that can be used
110+
to check for errors.
111+
"""
112+
function init(; retval=nothing::Union{Nothing,IR.Type}, location=Location())
113+
op_ty_results = IR.Type[]
114+
operands = Value[]
115+
owned_regions = Region[]
116+
successors = Block[]
117+
attributes = NamedAttribute[]
118+
!isnothing(retval) && push!(op_ty_results, retval)
119+
120+
return create_operation(
121+
"mpi.init",
122+
location;
123+
operands,
124+
owned_regions,
125+
successors,
126+
attributes,
127+
results=op_ty_results,
128+
result_inference=false,
129+
)
130+
end
131+
132+
"""
133+
`recv`
134+
135+
MPI_Recv performs a blocking receive of `size` elements of type `dtype`
136+
from rank `dest`. The `tag` value and communicator enables the library to
137+
determine the matching of multiple sends and receives between the same
138+
ranks.
139+
140+
Communicators other than `MPI_COMM_WORLD` are not supprted for now.
141+
The MPI_Status is set to `MPI_STATUS_IGNORE`, as the status object
142+
is not yet ported to MLIR.
143+
144+
This operation can optionally return an `!mpi.retval` value that can be used
145+
to check for errors.
146+
"""
147+
function recv(
148+
ref::Value,
149+
tag::Value,
150+
rank::Value;
151+
retval=nothing::Union{Nothing,IR.Type},
152+
location=Location(),
153+
)
154+
op_ty_results = IR.Type[]
155+
operands = Value[ref, tag, rank]
156+
owned_regions = Region[]
157+
successors = Block[]
158+
attributes = NamedAttribute[]
159+
!isnothing(retval) && push!(op_ty_results, retval)
160+
161+
return create_operation(
162+
"mpi.recv",
163+
location;
164+
operands,
165+
owned_regions,
166+
successors,
167+
attributes,
168+
results=op_ty_results,
169+
result_inference=false,
170+
)
171+
end
172+
173+
"""
174+
`retval_check`
175+
176+
This operation compares MPI status codes to known error class
177+
constants such as `MPI_SUCCESS`, or `MPI_ERR_COMM`.
178+
"""
179+
function retval_check(val::Value; res::IR.Type, errclass, location=Location())
180+
op_ty_results = IR.Type[res,]
181+
operands = Value[val,]
182+
owned_regions = Region[]
183+
successors = Block[]
184+
attributes = NamedAttribute[namedattribute("errclass", errclass),]
185+
186+
return create_operation(
187+
"mpi.retval_check",
188+
location;
189+
operands,
190+
owned_regions,
191+
successors,
192+
attributes,
193+
results=op_ty_results,
194+
result_inference=false,
195+
)
196+
end
197+
198+
"""
199+
`send`
200+
201+
MPI_Send performs a blocking send of `size` elements of type `dtype` to rank
202+
`dest`. The `tag` value and communicator enables the library to determine
203+
the matching of multiple sends and receives between the same ranks.
204+
205+
Communicators other than `MPI_COMM_WORLD` are not supprted for now.
206+
207+
This operation can optionally return an `!mpi.retval` value that can be used
208+
to check for errors.
209+
"""
210+
function send(
211+
ref::Value,
212+
tag::Value,
213+
rank::Value;
214+
retval=nothing::Union{Nothing,IR.Type},
215+
location=Location(),
216+
)
217+
op_ty_results = IR.Type[]
218+
operands = Value[ref, tag, rank]
219+
owned_regions = Region[]
220+
successors = Block[]
221+
attributes = NamedAttribute[]
222+
!isnothing(retval) && push!(op_ty_results, retval)
223+
224+
return create_operation(
225+
"mpi.send",
226+
location;
227+
operands,
228+
owned_regions,
229+
successors,
230+
attributes,
231+
results=op_ty_results,
232+
result_inference=false,
233+
)
234+
end
235+
236+
end # mpi

0 commit comments

Comments
 (0)