mpi: add MPIX_BFLOAT16

This serves as an example how we add a new builtin mpi datatype.
    1. define the constant in mpi.h.in
    2. (optional) define the internal datatype in mpir_datatype.h if there isn't
one already
        2a. add alignment in MPIR_Datatype_builtintype_alignment
        2b. add mapping in MPII_Typerep_get_yaksa_type
    3. define the mapping in configure.ac
    4. (optional) define case for the supported reduction op

Author: hzhou

configure.ac

@@ -3576,6 +3576,7 @@ AC_DEFINE_UNQUOTED([MPIR_LOGICAL2_INTERNAL],  [MPIR_FORTRAN_LOGICAL16], [Interna
 AC_DEFINE_UNQUOTED([MPIR_LOGICAL4_INTERNAL],  [MPIR_FORTRAN_LOGICAL32], [Internal type for MPI_LOGICAL4])
 AC_DEFINE_UNQUOTED([MPIR_LOGICAL8_INTERNAL],  [MPIR_FORTRAN_LOGICAL64], [Internal type for MPI_LOGICAL8])
 AC_DEFINE_UNQUOTED([MPIR_LOGICAL16_INTERNAL], [MPIR_FORTRAN_LOGICAL128],[Internal type for MPI_LOGICAL16])
+AC_DEFINE_UNQUOTED([MPIR_BFLOAT16_INTERNAL],  [MPIR_BFLOAT16],[Internal type for MPIX_BFLOAT16])
 AC_MSG_RESULT([done])
 
 # ------------------------------------------------------------------------

src/include/mpi.h.in

@@ -259,6 +259,8 @@ typedef int MPI_Datatype;
 #define MPI_LOGICAL4      ((MPI_Datatype)0x4c000449)
 #define MPI_LOGICAL8      ((MPI_Datatype)0x4c00084a)
 #define MPI_LOGICAL16     ((MPI_Datatype)0x4c00104b)
+/* other */
+#define MPIX_BFLOAT16     ((MPI_Datatype)0x4c00024c)
 
 /* Communicators */
 typedef int MPI_Comm;

src/include/mpir_datatype.h

@@ -72,6 +72,7 @@
 #define MPIR_COMPLEX32          ((MPI_Datatype)0x4c840800)
 #define MPIR_COMPLEX64          ((MPI_Datatype)0x4c841000)
 #define MPIR_COMPLEX128         ((MPI_Datatype)0x4c842000)
+#define MPIR_BFLOAT16           ((MPI_Datatype)0x4c850200)      /* bfloat16, use MPIR_TYPE_ALT_FLOAT */
 #define MPIR_ALT_FLOAT96        ((MPI_Datatype)0x4c850c00)      /* long double (80-bit extended precision) on i386 */
 #define MPIR_ALT_FLOAT128       ((MPI_Datatype)0x4c851000)      /* long double (80-bit extended precision) on x86-64 */
 #define MPIR_ALT_COMPLEX96      ((MPI_Datatype)0x4c861800)      /* long double complex on i386 */

src/include/mpir_objects.h

@@ -217,7 +217,7 @@ const char *MPIR_Handle_get_kind_str(int kind);
 #define MPIR_GROUP_PREALLOC 8
 #endif
 
-#define MPIR_DATATYPE_N_BUILTIN 76      /* 0x4c - must be in sync with mpi.h.in */
+#define MPIR_DATATYPE_N_BUILTIN 77      /* 0x4d - must be in sync with mpi.h.in */
 #ifdef MPID_DATATYPE_PREALLOC
 #define MPIR_DATATYPE_PREALLOC MPID_DATATYPE_PREALLOC
 #else

src/mpi/coll/op/op_fns.c

@@ -11,6 +11,8 @@
 
 #define MPIR_LSUM(a,b) ((a)+(b))
 
+static void bfloat16_sum(void *invec, void *inoutvec, MPI_Aint len);
+
 void MPIR_SUM(void *invec, void *inoutvec, MPI_Aint * Len, MPI_Datatype * type)
 {
     MPI_Aint i, len = *Len;
@@ -35,6 +37,9 @@ void MPIR_SUM(void *invec, void *inoutvec, MPI_Aint * Len, MPI_Datatype * type)
             break;                                         \
         }
                 MPIR_OP_TYPE_GROUP(COMPLEX)
+        case MPIR_BFLOAT16:
+            bfloat16_sum(invec, inoutvec, len);
+            break;
         default:
             MPIR_Assert(0);
             break;
@@ -442,3 +447,39 @@ void MPIR_REPLACE(void *invec, void *inoutvec, MPI_Aint * Len, MPI_Datatype * ty
   fn_fail:
     goto fn_exit;
 }
+
+/* -- internal static routines -- */
+
+/* BFloat16 - software arithemetics
+ * TODO: add hardware support, e.g. via AVX512 intrinsics
+ */
+static float bfloat16_load(void *p)
+{
+    uint32_t u = ((uint32_t) (*(uint16_t *) p) << 16);
+    float v;
+    memcpy(&v, &u, sizeof(float));
+    return v;
+}
+
+static void bfloat16_store(void *p, float v)
+{
+    uint32_t u;
+    memcpy(&u, &v, sizeof(float));
+    if (u & 0x8000) {
+        /* round up */
+        *(uint16_t *) p = (u >> 16) + 1;
+    } else {
+        /* truncation */
+        *(uint16_t *) p = (u >> 16);
+    }
+
+}
+
+static void bfloat16_sum(void *invec, void *inoutvec, MPI_Aint len)
+{
+    for (MPI_Aint i = 0; i < len * 2; i += 2) {
+        float a = bfloat16_load((char *) inoutvec + i);
+        float b = bfloat16_load((char *) invec + i);
+        bfloat16_store((char *) inoutvec + i, a + b);
+    }
+}

src/mpi/datatype/typerep/src/typerep_yaksa_init.c

@@ -52,6 +52,7 @@ yaksa_type_t MPII_Typerep_get_yaksa_type(MPI_Datatype type)
 
         case MPIR_FIXED16:
         case MPIR_FLOAT16:
+        case MPIR_BFLOAT16:
         case MPIR_COMPLEX8:
             yaksa_type = TYPEREP_YAKSA_TYPE__FIXED2;
             break;

src/mpi/datatype/typeutil.c

@@ -103,6 +103,7 @@ struct MPIR_Datatype_builtin_entry MPIR_Internal_types[] = {
     type_name_entry(LOGICAL4,           LOGICAL),         /* 0x49 */
     type_name_entry(LOGICAL8,           LOGICAL),         /* 0x4a */
     type_name_entry(LOGICAL16,          LOGICAL),         /* 0x4b */
+    type_name_x(BFLOAT16,               FLOATING_POINT),  /* 0x4c */
     /* *INDENT-ON* */
 };
 
@@ -169,6 +170,7 @@ int MPIR_Datatype_builtintype_alignment(MPI_Datatype type)
         case MPIR_INT16:
         case MPIR_UINT16:
         case MPIR_FLOAT16:
+        case MPIR_BFLOAT16:
             return ALIGNOF_INT16_T;
         case MPIR_FIXED32:
         case MPIR_INT32:
@@ -190,7 +192,7 @@ int MPIR_Datatype_builtintype_alignment(MPI_Datatype type)
         case MPIR_ALT_COMPLEX128:
             return ALIGNOF_LONG_DOUBLE;
         default:
-            /* handle error cases? */
+            MPIR_Assert(0);
             return 1;
     }
 }