Add bf16 support in rccl-tests

2019-11-15 13:46:03 -08:00
parent 32399955af
commit 1cda2f52b6
4 changed files with 283 additions and 3 deletions
@@ -195,6 +195,10 @@ static size_t wordSize(ncclDataType_t type) {
    case ncclDouble:
    //case ncclFloat64:
      return 8;
+#if NCCL_MAJOR >= 2 && RCCL_BFLOAT16 == 1
+    case ncclBfloat16:
+      return 2;
+#endif
    default: return 0;
  }
 }