IMB-MPI1: FP16/BF16 data_type support

Author: JuliaRS

src_cpp/MPI1/MPI1_suite.cpp

@@ -312,7 +312,7 @@ template <> bool BenchmarkSuite<BS_MPI1>::declare_args(args_parser &parser, std:
            set_description(
                 "The argument after -data_type is a one from possible strings,\n"
                 "Specifying that type will be used:\n"
-                "byte, char, int, float, double\n"
+                "byte, char, int, float, double, float16, bfloat16\n"
                 "\n"
                 "Example:\n"
                 "-data_type char\n"
@@ -323,7 +323,7 @@ template <> bool BenchmarkSuite<BS_MPI1>::declare_args(args_parser &parser, std:
            set_description(
                 "The argument after -red_data_type is a one from possible strings,\n"
                 "Specifying that type will be used:\n"
-                "char, int, float, double\n"
+                "char, int, float, double, float16, bfloat16\n"
                 "\n"
                 "Example:\n"
                 "-red_data_type int\n"
@@ -414,12 +414,24 @@ MPI_Op get_op(MPI_Datatype type) {
     MPI_Datatype mpi_int = MPI_INT;
     MPI_Datatype mpi_float = MPI_FLOAT;
     MPI_Datatype mpi_double = MPI_DOUBLE;
+#ifdef MPIX_C_FLOAT16
+    MPI_Datatype mpi_float16 = MPIX_C_FLOAT16;
+#endif
+#ifdef MPIX_C_BF16
+    MPI_Datatype mpi_bfloat16 = MPIX_C_BF16;
+#endif
     size_t type_size = sizeof(MPI_Datatype);
 
     if (!memcmp(&type, &mpi_char, type_size)) { MPI_Op_create(&(contig_sum<char>), 1, &op); }
     else if (!memcmp(&type, &mpi_int, type_size)) { MPI_Op_create(&(contig_sum<int>), 1, &op); }
     else if (!memcmp(&type, &mpi_float, type_size)) { MPI_Op_create(&(contig_sum<float>), 1, &op); }
     else if (!memcmp(&type, &mpi_double, type_size)) { MPI_Op_create(&(contig_sum<double>), 1, &op); }
+#ifdef MPIX_C_FLOAT16
+    else if (!memcmp(&type, &mpi_float16, type_size)) { op = MPI_OP_NULL; fprintf(stdout, "\nWarning: contig_type doesn't supported\n"); }
+#endif
+#ifdef MPIX_C_BF16
+    else if (!memcmp(&type, &mpi_bfloat16, type_size)) { op = MPI_OP_NULL; fprintf(stdout, "\nWarning: contig_type doesn't supported \n"); }
+#endif
 
     return op;
 }
@@ -431,13 +443,25 @@ string type_to_name(MPI_Datatype type) {
     MPI_Datatype mpi_int = MPI_INT;
     MPI_Datatype mpi_float = MPI_FLOAT;
     MPI_Datatype mpi_double = MPI_DOUBLE;
+#ifdef MPIX_C_FLOAT16
+    MPI_Datatype mpi_float16 = MPIX_C_FLOAT16;
+#endif
+#ifdef MPIX_C_BF16
+    MPI_Datatype mpi_bfloat16 = MPIX_C_BF16;
+#endif
     size_t type_size = sizeof(MPI_Datatype);
 
     if (!memcmp(&type, &mpi_byte, type_size)) { name = "MPI_BYTE"; }
     else if (!memcmp(&type, &mpi_char, type_size)) { name = "MPI_CHAR"; }
     else if (!memcmp(&type, &mpi_int, type_size)) { name = "MPI_INT"; }
     else if (!memcmp(&type, &mpi_float, type_size)) { name = "MPI_FLOAT"; }
     else if (!memcmp(&type, &mpi_double, type_size)) { name = "MPI_DOUBLE"; }
+#ifdef MPIX_C_FLOAT16
+    else if (!memcmp(&type, &mpi_float16, type_size)) { name = "MPIX_C_FLOAT16"; }
+#endif
+#ifdef MPIX_C_BF16
+    else if (!memcmp(&type, &mpi_bfloat16, type_size)) { name = "MPIX_C_BF16"; }
+#endif
 
     return name;
 }
@@ -602,6 +626,16 @@ template <> bool BenchmarkSuite<BS_MPI1>::prepare(const args_parser &parser, con
     } else if (given_data_type == "double") {
         c_info.s_data_type = MPI_DOUBLE;
         c_info.r_data_type = MPI_DOUBLE;
+#ifdef MPIX_C_FLOAT16
+    } else if (given_data_type == "float16") {
+        c_info.s_data_type = MPIX_C_FLOAT16;
+        c_info.r_data_type = MPIX_C_FLOAT16;
+#endif
+#ifdef MPIX_C_BF16
+    } else if (given_data_type == "bfloat16") {
+        c_info.s_data_type = MPIX_C_BF16;
+        c_info.r_data_type = MPIX_C_BF16;
+#endif
     } else {
         output << "Invalid data_type " << given_data_type << endl;
         output << "    Set data_type byte" << endl;
@@ -619,6 +653,14 @@ template <> bool BenchmarkSuite<BS_MPI1>::prepare(const args_parser &parser, con
         c_info.red_data_type = MPI_FLOAT;
     } else if (given_red_data_type == "double") {
         c_info.red_data_type = MPI_DOUBLE;
+#ifdef MPIX_C_FLOAT16
+    } else if (given_red_data_type == "float16") {
+        c_info.red_data_type = MPIX_C_FLOAT16;
+#endif
+#ifdef MPIX_C_BF16
+    } else if (given_red_data_type == "bfloat16") {
+        c_info.red_data_type = MPIX_C_BF16;
+#endif
     } else {
         output << "Invalid red_data_type " << given_red_data_type << endl;
         output << "    Set red_data_type float" << endl;