MPI Ops and lowering passes

src/enzyme_ad/jax/Dialect/EnzymeXLAAttrs.td

@@ -133,4 +133,72 @@ def EnzymeXLA_GuaranteedAnalysisResult : I32EnumAttr<"GuaranteedAnalysisResult",
 def EnzymeXLA_GuaranteedAnalysisResultAttr : EnumAttr<EnzymeXLA_Dialect,
     EnzymeXLA_GuaranteedAnalysisResult, "guaranteed">;
 
+// MPI
+
+def EnzymeXLA_MPIDatatype : I32EnumAttr<"MPIDatatype",
+    "MPI Datatype",
+    [
+      I32EnumAttrCase<"MPI_DATATYPE_NULL", 0>,
+      I32EnumAttrCase<"MPI_INT8_T", 1>,
+      I32EnumAttrCase<"MPI_UINT8_T", 2>,
+      I32EnumAttrCase<"MPI_INT16_T", 3>,
+      I32EnumAttrCase<"MPI_UINT16_T", 4>,
+      I32EnumAttrCase<"MPI_INT32_T", 5>,
+      I32EnumAttrCase<"MPI_UINT32_T", 6>,
+      I32EnumAttrCase<"MPI_INT64_T", 7>,
+      I32EnumAttrCase<"MPI_UINT64_T", 8>,
+      I32EnumAttrCase<"MPI_BYTE", 9>,
+      I32EnumAttrCase<"MPI_SHORT", 10>,
+      I32EnumAttrCase<"MPI_UNSIGNED_SHORT", 11>,
+      I32EnumAttrCase<"MPI_INT", 12>,
+      I32EnumAttrCase<"MPI_UNSIGNED", 13>,
+      I32EnumAttrCase<"MPI_LONG", 14>,
+      I32EnumAttrCase<"MPI_UNSIGNED_LONG", 15>,
+      I32EnumAttrCase<"MPI_LONG_LONG_INT", 16>,
+      I32EnumAttrCase<"MPI_UNSIGNED_LONG_LONG", 17>,
+      I32EnumAttrCase<"MPI_CHAR", 18>,
+      I32EnumAttrCase<"MPI_SIGNED_CHAR", 19>,
+      I32EnumAttrCase<"MPI_UNSIGNED_CHAR", 20>,
+      I32EnumAttrCase<"MPI_WCHAR", 21>,
+      I32EnumAttrCase<"MPI_FLOAT", 22>,
+      I32EnumAttrCase<"MPI_DOUBLE", 23>,
+      I32EnumAttrCase<"MPI_C_FLOAT_COMPLEX", 24>,
+      I32EnumAttrCase<"MPI_C_DOUBLE_COMPLEX", 25>,
+      I32EnumAttrCase<"MPI_C_BOOL", 26>
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::enzymexla";
+}
+
+def EnzymeXLA_MPIDatatypeAttr : EnumAttr<EnzymeXLA_Dialect,
+    EnzymeXLA_MPIDatatype, "datatype"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
+def EnzymeXLA_MPIOp : I32EnumAttr<"MPIOp",
+    "MPI Operator",
+    [
+      I32EnumAttrCase<"MPI_OP_NULL", 0>,
+      I32EnumAttrCase<"MPI_BAND", 1>,
+      I32EnumAttrCase<"MPI_BOR", 2>,
+      I32EnumAttrCase<"MPI_BXOR", 3>,
+      I32EnumAttrCase<"MPI_LAND", 4>,
+      I32EnumAttrCase<"MPI_LOR", 5>,
+      I32EnumAttrCase<"MPI_LXOR", 6>,
+      I32EnumAttrCase<"MPI_MAX", 7>,
+      I32EnumAttrCase<"MPI_MIN", 8>,
+      I32EnumAttrCase<"MPI_PROD", 9>,
+      I32EnumAttrCase<"MPI_REPLACE", 10>,
+      I32EnumAttrCase<"MPI_SUM", 11>,
+      I32EnumAttrCase<"MPI_NO_OP", 12>
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::enzymexla";
+}
+
+def EnzymeXLA_MPIOpAttr : EnumAttr<EnzymeXLA_Dialect,
+    EnzymeXLA_MPIOp, "op"> {
+  let assemblyFormat = "`<` $value `>`";
+}
+
 #endif // ENZYMEXLA_ATTRS

src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td

@@ -1088,4 +1088,130 @@ def AffineStoreVar : EnzymeXLA_Op<"store_var", [Pure]> {
   let summary = "Fake store an SSA value for conversion to ISL";
 }
 
+// MPI Ops
+
+def MPICommRankOp : EnzymeXLA_Op<"mpi.comm_rank", [Pure]> {
+  let summary = "Equivalent to " "`MPI_Comm_rank(MPI_COMM_WORLD, &rank)`";
+
+  let results = (
+    outs TensorOf<[I32]> : $rank
+  );
+
+  let assemblyFormat = "attr-dict `:` type(results)";
+}
+
+def MPICommSizeOp : EnzymeXLA_Op<"mpi.comm_size", [Pure]> {
+  let summary = "Equivalent to MPI_Comm_size(MPI_COMM_WORLD, &size)";
+
+  let results = (
+    outs TensorOf<[I32]> : $size
+  );
+
+  let assemblyFormat = "attr-dict `:` type(results)";
+}
+
+def MPIBarrierOp : EnzymeXLA_Op<"mpi.barrier", []> {
+  let summary = "Equivalent to MPI_Barrier(MPI_COMM_WORLD)";
+  let assemblyFormat = "attr-dict";
+}
+
+def MPISendOp : EnzymeXLA_Op<"mpi.send", []> {
+  let summary = "Equivalent to "
+                "`MPI_Send(&buf, count, datatype, dest, tag, comm)`";
+
+  let arguments = (
+    ins AnyTensor : $buf,
+    TensorOf<[I32]> : $count,
+    TensorOf<[I32]> : $dest,
+    TensorOf<[I32]> : $tag,
+    EnzymeXLA_MPIDatatypeAttr:$datatype
+  );
+
+  let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
+}
+
+def MPIRecvOp : EnzymeXLA_Op<"mpi.recv", []> {
+  let summary = "Equivalent to "
+                "`MPI_Recv(&buf, count, datatype, source, tag, MPI_COMM_WORLD, MPI_STATUS_IGNORE)`";
+
+  let arguments = (
+    ins AnyTensor : $inbuf,
+    TensorOf<[I32]> : $count,
+    TensorOf<[I32]> : $source,
+    TensorOf<[I32]> : $tag,
+    EnzymeXLA_MPIDatatypeAttr:$datatype
+  );
+
+  let results = (
+    outs AnyTensor : $outbuf
+  );
+
+  let assemblyFormat = "`(` operands `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def MPIIsendOp : EnzymeXLA_Op<"mpi.isend", []> {
+  let summary = "Equivalent to "
+                "`MPI_Isend(&buf, count, datatype, dest, tag, MPI_COMM_WORLD, &request)`";
+
+  let arguments = (
+    ins AnyTensor : $buf,
+    TensorOf<[I32]> : $count,
+    TensorOf<[I32]> : $dest,
+    TensorOf<[I32]> : $tag,
+    EnzymeXLA_MPIDatatypeAttr:$datatype
+  );
+
+  let results = (
+    outs TensorOf<[I64]> : $request
+  );
+
+  let assemblyFormat = "`(` operands `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def MPIIrecvOp : EnzymeXLA_Op<"mpi.irecv", []> {
+  let summary = "Equivalent to "
+                "`MPI_Irecv(&buf, count, datatype, source, tag, MPI_COMM_WORLD, &request)`";
+
+  let arguments = (
+    ins AnyTensor : $inbuf,
+    TensorOf<[I32]> : $count,
+    TensorOf<[I32]> : $source,
+    TensorOf<[I32]> : $tag,
+    EnzymeXLA_MPIDatatypeAttr:$datatype
+  );
+
+  let results = (
+    outs AnyTensor : $outbuf,
+    TensorOf<[I64]> : $request
+  );
+
+  let assemblyFormat = "`(` operands `)` attr-dict `:` functional-type(operands, results)";
+}
+
+def MPIWaitOp : EnzymeXLA_Op<"mpi.wait", []> {
+  let summary = "Equivalent to "
+                "`MPI_Wait(&request, &status)`";
+  let arguments = (ins TensorOf<[I64]> : $request);
+  let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)";
+}
+
+def MPIAllreduceOp : EnzymeXLA_Op<"mpi.allreduce", []> {
+  let summary = "Equivalent to "
+                "`MPI_Allreduce(&sendbuf, &recvbuf, count, datatype, op, MPI_COMM_WORLD)`";
+
+  let arguments = (
+    ins AnyTensor : $sendbuf,
+    AnyTensor : $inbuf,
+    TensorOf<[I32]> : $count,
+    EnzymeXLA_MPIDatatypeAttr:$datatype,
+    EnzymeXLA_MPIOpAttr:$op
+  );
+
+  let results = (
+    outs AnyTensor : $outbuf
+  );
+
+  let assemblyFormat = "`(` operands `)` attr-dict `:` functional-type(operands, results)";
+}
+
 #endif // ENZYMEXLA_OPS

src/enzyme_ad/jax/Integrations/c/EnzymeXLA.cpp

@@ -122,3 +122,141 @@ MlirAttribute enzymexlaGuaranteedAnalysisResultAttrGet(MlirContext ctx,
   return wrap(mlir::enzymexla::GuaranteedAnalysisResultAttr::get(unwrap(ctx),
                                                                  analysis));
 }
+
+MlirAttribute enzymexlaMPIDatatypeAttrGet(MlirContext ctx, int32_t mode) {
+  mlir::enzymexla::MPIDatatype datatype;
+  switch (mode) {
+  case 0:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_DATATYPE_NULL;
+    break;
+  case 1:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_INT8_T;
+    break;
+  case 2:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UINT8_T;
+    break;
+  case 3:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_INT16_T;
+    break;
+  case 4:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UINT16_T;
+    break;
+  case 5:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_INT32_T;
+    break;
+  case 6:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UINT32_T;
+    break;
+  case 7:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_INT64_T;
+    break;
+  case 8:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UINT64_T;
+    break;
+  case 9:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_BYTE;
+    break;
+  case 10:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_SHORT;
+    break;
+  case 11:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UNSIGNED_SHORT;
+    break;
+  case 12:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_INT;
+    break;
+  case 13:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UNSIGNED;
+    break;
+  case 14:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_LONG;
+    break;
+  case 15:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UNSIGNED_LONG;
+    break;
+  case 16:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_LONG_LONG_INT;
+    break;
+  case 17:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UNSIGNED_LONG_LONG;
+    break;
+  case 18:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_CHAR;
+    break;
+  case 19:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_SIGNED_CHAR;
+    break;
+  case 20:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_UNSIGNED_CHAR;
+    break;
+  case 21:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_WCHAR;
+    break;
+  case 22:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_FLOAT;
+    break;
+  case 23:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_DOUBLE;
+    break;
+  case 24:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_C_FLOAT_COMPLEX;
+    break;
+  case 25:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_C_DOUBLE_COMPLEX;
+    break;
+  case 26:
+    datatype = mlir::enzymexla::MPIDatatype::MPI_C_BOOL;
+    break;
+  default:
+    llvm_unreachable("Invalid MPI datatype mode");
+  }
+  return wrap(mlir::enzymexla::MPIDatatypeAttr::get(unwrap(ctx), datatype));
+}
+
+MlirAttribute enzymexlaMPIOpAttrGet(MlirContext ctx, int32_t mode) {
+  mlir::enzymexla::MPIOp op;
+  switch (mode) {
+  case 0:
+    op = mlir::enzymexla::MPIOp::MPI_OP_NULL;
+    break;
+  case 1:
+    op = mlir::enzymexla::MPIOp::MPI_BAND;
+    break;
+  case 2:
+    op = mlir::enzymexla::MPIOp::MPI_BOR;
+    break;
+  case 3:
+    op = mlir::enzymexla::MPIOp::MPI_BXOR;
+    break;
+  case 4:
+    op = mlir::enzymexla::MPIOp::MPI_LAND;
+    break;
+  case 5:
+    op = mlir::enzymexla::MPIOp::MPI_LOR;
+    break;
+  case 6:
+    op = mlir::enzymexla::MPIOp::MPI_LXOR;
+    break;
+  case 7:
+    op = mlir::enzymexla::MPIOp::MPI_MAX;
+    break;
+  case 8:
+    op = mlir::enzymexla::MPIOp::MPI_MIN;
+    break;
+  case 9:
+    op = mlir::enzymexla::MPIOp::MPI_PROD;
+    break;
+  case 10:
+    op = mlir::enzymexla::MPIOp::MPI_REPLACE;
+    break;
+  case 11:
+    op = mlir::enzymexla::MPIOp::MPI_SUM;
+    break;
+  case 12:
+    op = mlir::enzymexla::MPIOp::MPI_NO_OP;
+    break;
+  default:
+    llvm_unreachable("Invalid MPI op mode");
+  }
+  return wrap(mlir::enzymexla::MPIOpAttr::get(unwrap(ctx), op));
+}

src/enzyme_ad/jax/Integrations/c/EnzymeXLA.h

@@ -43,6 +43,16 @@ MLIR_CAPI_EXPORTED MlirAttribute enzymexlaSVDAlgorithmAttrGet(MlirContext ctx,
 MLIR_CAPI_EXPORTED MlirAttribute
 enzymexlaGeluApproximationAttrGet(MlirContext ctx, int32_t mode);
 
+//===----------------------------------------------------------------------===//
+// MPI Ops
+//===----------------------------------------------------------------------===//
+
+MLIR_CAPI_EXPORTED MlirAttribute enzymexlaMPIDatatypeAttrGet(MlirContext ctx,
+                                                             int32_t mode);
+
+MLIR_CAPI_EXPORTED MlirAttribute enzymexlaMPIOpAttrGet(MlirContext ctx,
+                                                       int32_t mode);
+
 //===----------------------------------------------------------------------===//
 // Other Ops / Attributes
 //===----------------------------------------------------------------------===//