From 08d0ac9e28a56a046c8f2505af51f0776f2de700 Mon Sep 17 00:00:00 2001 From: "enzyme-ci-bot[bot]" <78882869+enzyme-ci-bot[bot]@users.noreply.github.com> Date: Sun, 2 Feb 2025 11:08:11 +0000 Subject: [PATCH 1/2] Regenerate MLIR Bindings --- src/mlir/Dialects/MPI.jl | 236 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 236 insertions(+) create mode 100644 src/mlir/Dialects/MPI.jl diff --git a/src/mlir/Dialects/MPI.jl b/src/mlir/Dialects/MPI.jl new file mode 100644 index 0000000000..1ca0d84be8 --- /dev/null +++ b/src/mlir/Dialects/MPI.jl @@ -0,0 +1,236 @@ +module mpi +using ...IR +import ...IR: + NamedAttribute, + Value, + Location, + Block, + Region, + Attribute, + create_operation, + context, + IndexType +import ..Dialects: namedattribute, operandsegmentsizes +import ...API + +""" +`comm_rank` + +Communicators other than `MPI_COMM_WORLD` are not supported for now. + +This operation can optionally return an `!mpi.retval` value that can be used +to check for errors. +""" +function comm_rank(; + retval=nothing::Union{Nothing,IR.Type}, rank::IR.Type, location=Location() +) + op_ty_results = IR.Type[rank,] + operands = Value[] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + !isnothing(retval) && push!(op_ty_results, retval) + + return create_operation( + "mpi.comm_rank", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`error_class` + +`MPI_Error_class` maps return values from MPI calls to a set of well-known +MPI error classes. +""" +function error_class(val::Value; errclass::IR.Type, location=Location()) + op_ty_results = IR.Type[errclass,] + operands = Value[val,] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + + return create_operation( + "mpi.error_class", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`finalize` + +This function cleans up the MPI state. Afterwards, no MPI methods may +be invoked (excpet for MPI_Get_version, MPI_Initialized, and MPI_Finalized). +Notably, MPI_Init cannot be called again in the same program. + +This operation can optionally return an `!mpi.retval` value that can be used +to check for errors. +""" +function finalize(; retval=nothing::Union{Nothing,IR.Type}, location=Location()) + op_ty_results = IR.Type[] + operands = Value[] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + !isnothing(retval) && push!(op_ty_results, retval) + + return create_operation( + "mpi.finalize", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`init` + +This operation must preceed most MPI calls (except for very few exceptions, +please consult with the MPI specification on these). + +Passing &argc, &argv is not supported currently. + +This operation can optionally return an `!mpi.retval` value that can be used +to check for errors. +""" +function init(; retval=nothing::Union{Nothing,IR.Type}, location=Location()) + op_ty_results = IR.Type[] + operands = Value[] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + !isnothing(retval) && push!(op_ty_results, retval) + + return create_operation( + "mpi.init", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`recv` + +MPI_Recv performs a blocking receive of `size` elements of type `dtype` +from rank `dest`. The `tag` value and communicator enables the library to +determine the matching of multiple sends and receives between the same +ranks. + +Communicators other than `MPI_COMM_WORLD` are not supprted for now. +The MPI_Status is set to `MPI_STATUS_IGNORE`, as the status object +is not yet ported to MLIR. + +This operation can optionally return an `!mpi.retval` value that can be used +to check for errors. +""" +function recv( + ref::Value, + tag::Value, + rank::Value; + retval=nothing::Union{Nothing,IR.Type}, + location=Location(), +) + op_ty_results = IR.Type[] + operands = Value[ref, tag, rank] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + !isnothing(retval) && push!(op_ty_results, retval) + + return create_operation( + "mpi.recv", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`retval_check` + +This operation compares MPI status codes to known error class +constants such as `MPI_SUCCESS`, or `MPI_ERR_COMM`. +""" +function retval_check(val::Value; res::IR.Type, errclass, location=Location()) + op_ty_results = IR.Type[res,] + operands = Value[val,] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[namedattribute("errclass", errclass),] + + return create_operation( + "mpi.retval_check", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +""" +`send` + +MPI_Send performs a blocking send of `size` elements of type `dtype` to rank +`dest`. The `tag` value and communicator enables the library to determine +the matching of multiple sends and receives between the same ranks. + +Communicators other than `MPI_COMM_WORLD` are not supprted for now. + +This operation can optionally return an `!mpi.retval` value that can be used +to check for errors. +""" +function send( + ref::Value, + tag::Value, + rank::Value; + retval=nothing::Union{Nothing,IR.Type}, + location=Location(), +) + op_ty_results = IR.Type[] + operands = Value[ref, tag, rank] + owned_regions = Region[] + successors = Block[] + attributes = NamedAttribute[] + !isnothing(retval) && push!(op_ty_results, retval) + + return create_operation( + "mpi.send", + location; + operands, + owned_regions, + successors, + attributes, + results=op_ty_results, + result_inference=false, + ) +end + +end # mpi From a42798ffb652540a48aeb3afc8ca0bfc465993ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?= Date: Sun, 2 Feb 2025 20:58:43 +0100 Subject: [PATCH 2/2] fix docs --- docs/src/.vitepress/config.mts | 2 ++ docs/src/api/mpi.md | 12 ++++++++++++ 2 files changed, 14 insertions(+) create mode 100644 docs/src/api/mpi.md diff --git a/docs/src/.vitepress/config.mts b/docs/src/.vitepress/config.mts index 5fa5cb75d4..78099c1965 100644 --- a/docs/src/.vitepress/config.mts +++ b/docs/src/.vitepress/config.mts @@ -87,6 +87,7 @@ export default defineConfig({ { text: "TPU", link: "/api/tpu" }, { text: "Triton", link: "/api/triton" }, { text: "Shardy", link: "/api/shardy" }, + { text: "MPI", link: "/api/mpi" }, ], }, { @@ -147,6 +148,7 @@ export default defineConfig({ { text: "TPU", link: "/api/tpu" }, { text: "Triton", link: "/api/triton" }, { text: "Shardy", link: "/api/shardy" }, + { text: "MPI", link: "/api/mpi" }, ], }, { diff --git a/docs/src/api/mpi.md b/docs/src/api/mpi.md new file mode 100644 index 0000000000..5b0570714e --- /dev/null +++ b/docs/src/api/mpi.md @@ -0,0 +1,12 @@ +```@meta +CollapsedDocStrings = true +``` + +# MPI Dialect + +Refer to the [official documentation](https://mlir.llvm.org/docs/Dialects/MPI/) for +more details. + +```@autodocs +Modules = [Reactant.MLIR.Dialects.mpi] +```