From 6777e65c1dcffb5e9bd2097f5f933a57aec1613e Mon Sep 17 00:00:00 2001 From: Andy li Date: Sat, 9 Mar 2024 07:17:53 +0800 Subject: [PATCH] Enable fp8 support (#1101) * initial checkin * resolve cr comments * resolve the build issue * fix the data correctless issue * update fp8 header file and update the unit test for fp8 support * remove fp16 from fp8 headers * fix ut issue and catch up the latest code from develop * udate according to cr comments * update ut according to cr comments * update num floats for each SumPostDiv from 4 to 6 * update fp8 header file name * fix the typo --- CMakeLists.txt | 1 + cmake/Generator.cmake | 4 +- ext-net/example/nccl/types.h | 2 + src/device/msccl_kernel_impl.h | 4 +- src/device/onerank.cu | 9 +- src/device/reduce_kernel.h | 74 ++ src/enqueue.cc | 25 +- src/include/collectives.h | 9 +- src/include/device.h | 22 +- src/include/msccl/msccl_kernel.h | 4 +- src/include/nccl_common.h | 6 +- src/include/rccl_float8.h | 1021 ++++++++++++++++++++++++++ src/misc/msccl/msccl_setup.cc | 20 +- src/nccl.h.in | 7 + test/AllGatherTests.cpp | 2 +- test/AllReduceTests.cpp | 4 +- test/AllToAllTests.cpp | 2 +- test/BroadcastTests.cpp | 2 +- test/ReduceScatterTests.cpp | 2 +- test/ReduceTests.cpp | 4 +- test/ScatterTests.cpp | 2 +- test/common/CollectiveArgs.cpp | 2 + test/common/CollectiveArgs.hpp | 4 +- test/common/EnvVars.cpp | 2 + test/common/PtrUnion.cpp | 42 +- test/common/PtrUnion.hpp | 3 + tools/ib-test/include/nccl.h | 5 +- tools/rccl_replayer/rcclReplayer.hpp | 2 + tools/topo_expl/include/nccl.h | 5 +- 29 files changed, 1243 insertions(+), 48 deletions(-) create mode 100644 src/include/rccl_float8.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 0e27118a94..d0977e3cfe 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -374,6 +374,7 @@ set(SRC_FILES src/include/proxy.h src/include/rccl_bfloat16.h src/include/rccl_vars.h + src/include/rccl_float8.h src/include/rocm_smi_wrap.h src/include/rocmwrap.h src/include/roctx.h diff --git a/cmake/Generator.cmake b/cmake/Generator.cmake index 147d92f0aa..91fc54d98a 100644 --- a/cmake/Generator.cmake +++ b/cmake/Generator.cmake @@ -25,9 +25,9 @@ set(ALL_COLLS "AllGather" "AllReduce" "AllToAllPivot" "Broadcast" "Reduce" "Redu set(ALL_ALGOS "TREE" "RING" "COLLNET_DIRECT" "COLLNET_CHAIN") set(ALL_PROTOS "LL" "LL128" "SIMPLE") set(ALL_REDOPS "Sum" "Prod" "MinMax" "PreMulSum" "SumPostDiv") -set(ALL_TYPES "int8_t" "uint8_t" "int32_t" "uint32_t" "int64_t" "uint64_t" "half" "float" "double" "rccl_bfloat16") +set(ALL_TYPES "int8_t" "uint8_t" "int32_t" "uint32_t" "int64_t" "uint64_t" "half" "float" "double" "rccl_bfloat16" "rccl_float8" "rccl_bfloat8") -set(FLOATS_LIST "half" "float" "double" "rccl_bfloat16") +set(FLOATS_LIST "half" "float" "double" "rccl_bfloat16" "rccl_float8" "rccl_bfloat8") ################################################################################ # The command line argument is used as a regex to filter the functions diff --git a/ext-net/example/nccl/types.h b/ext-net/example/nccl/types.h index 0a5d83788d..c78f7cae0b 100644 --- a/ext-net/example/nccl/types.h +++ b/ext-net/example/nccl/types.h @@ -16,6 +16,8 @@ typedef enum { ncclInt8 = 0, ncclChar = 0, ncclFloat32 = 7, ncclFloat = 7, ncclFloat64 = 8, ncclDouble = 8, ncclBfloat16 = 9, + ncclFp8E4M3 = 10, + ncclFp8E5M2 = 11, } ncclDataType_t; #endif diff --git a/src/device/msccl_kernel_impl.h b/src/device/msccl_kernel_impl.h index f3e172c9cc..3baf4be35c 100644 --- a/src/device/msccl_kernel_impl.h +++ b/src/device/msccl_kernel_impl.h @@ -414,7 +414,9 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, Simple, fullOps)(struct MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, half, fullOps) \ MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, float, fullOps) \ MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, double, fullOps) \ - MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps) + MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps) \ + MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_float8, fullOps) \ + MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat8, fullOps) #define MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(devredop, fullOps) \ MSCCL_IMPL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \ diff --git a/src/device/onerank.cu b/src/device/onerank.cu index 92ea13246a..62f4540a2a 100644 --- a/src/device/onerank.cu +++ b/src/device/onerank.cu @@ -1,5 +1,6 @@ /************************************************************************* * Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. + * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License. * * See LICENSE.txt for license information ************************************************************************/ @@ -63,9 +64,13 @@ ncclResult_t ncclLaunchOneRank(void* dst, void const* src, size_t nElts, struct case ncclInt64: kernel = (void const*)&oneRankReduce>; break; case ncclUint64: kernel = (void const*)&oneRankReduce>; break; case ncclFloat16: kernel = (void const*)&oneRankReduce>; break; - #if defined(RCCL_BFLOAT16) +#if defined(RCCL_BFLOAT16) case ncclBfloat16: kernel = (void const*)&oneRankReduce>; break; - #endif +#endif +#if defined(RCCL_FLOAT8) + case ncclFp8E4M3: kernel = (void const*)&oneRankReduce>; break; + case ncclFp8E5M2: kernel = (void const*)&oneRankReduce>; break; +#endif case ncclFloat32: kernel = (void const*)&oneRankReduce>; break; case ncclFloat64: kernel = (void const*)&oneRankReduce>; break; default: return ncclInvalidArgument; diff --git a/src/device/reduce_kernel.h b/src/device/reduce_kernel.h index d3fd848371..09d9314a9f 100644 --- a/src/device/reduce_kernel.h +++ b/src/device/reduce_kernel.h @@ -1,6 +1,7 @@ /************************************************************************* * Copyright (c) 2015-2021, NVIDIA CORPORATION. All rights reserved. * Modifications Copyright (c) 2019-2021 Advanced Micro Devices, Inc. All rights reserved. + * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License. * * See LICENSE.txt for license information ************************************************************************/ @@ -13,6 +14,8 @@ #include #include +#include "rccl_float8.h" + template struct IsFloatingPoint: std::false_type {}; template<> @@ -21,6 +24,12 @@ struct IsFloatingPoint: std::true_type {}; template<> struct IsFloatingPoint: std::true_type {}; #endif +#if defined(RCCL_FLOAT8) +template<> +struct IsFloatingPoint: std::true_type {}; +template<> +struct IsFloatingPoint: std::true_type {}; +#endif template<> struct IsFloatingPoint: std::true_type {}; template<> @@ -254,6 +263,15 @@ SPECIALIZE_REDUCE(FuncMinMax, half, 1, half, fn.isMinNotMax ? __hmin(x, y) : __h #endif #endif +#if defined(RCCL_FLOAT8) + SPECIALIZE_REDUCE(FuncSum, rccl_float8, 1, rccl_float8, rccl_float8(float(x) + float(y))) + SPECIALIZE_REDUCE(FuncProd, rccl_float8, 1, rccl_float8, rccl_float8(float(x) * float(y))) + SPECIALIZE_REDUCE(FuncMinMax, rccl_float8, 1, rccl_float8, rccl_float8(fn.isMinNotMax ? fminf(float(x), float(y)) : fmaxf(float(x), float(y)))) + SPECIALIZE_REDUCE(FuncSum, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(float(x) + float(y))) + SPECIALIZE_REDUCE(FuncProd, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(float(x) * float(y))) + SPECIALIZE_REDUCE(FuncMinMax, rccl_bfloat8, 1, rccl_bfloat8, rccl_bfloat8(fn.isMinNotMax ? fminf(float(x), float(y)) : fmaxf(float(x), float(y)))) +#endif + #undef SPECIALIZE_REDUCE //////////////////////////////////////////////////////////////////////////////// @@ -389,6 +407,38 @@ struct FuncPreMulSum { }; #endif +#if defined(RCCL_FLOAT8) + template<> + struct FuncPreMulSum { + // Change these to switch between all prescale, all postscale, or both by sqrt(N). + // Obviously, the only invalid combination is both true. An improvement would be + // make this parameterized as a build time setting and passed here through + // preprocessor definitions. + using EltType = rccl_float8; + float scalar; + __device__ FuncPreMulSum(uint64_t opArg=0) { + union { uint64_t u64; rccl_float8 val; }; + u64 = opArg; + scalar = (float)(val); + } + }; + + template<> + struct FuncPreMulSum { + // Change these to switch between all prescale, all postscale, or both by sqrt(N). + // Obviously, the only invalid combination is both true. An improvement would be + // make this parameterized as a build time setting and passed here through + // preprocessor definitions. + using EltType = rccl_bfloat8; + float scalar; + __device__ FuncPreMulSum(uint64_t opArg=0) { + union { uint64_t u64; rccl_bfloat8 val; }; + u64 = opArg; + scalar = (float)(val); + } + }; +#endif + template struct Apply_Reduce, /*EltPerPack=*/1> { __device__ static BytePack reduce(FuncPreMulSum fn, BytePack a, BytePack b) { @@ -456,6 +506,30 @@ struct Apply_PreOp, /*EltPerPack=*/1> { #endif #endif +#if defined(RCCL_FLOAT8) + template<> + struct Apply_PreOp, /*EltPerPack=*/1> { + static constexpr bool IsIdentity = false; + + __device__ static BytePack preOp( + FuncPreMulSum fn, BytePack a + ) { + return toPack(rccl_float8(float(fromPack(a)) * float(fn.scalar))); + } + }; + + template<> + struct Apply_PreOp, /*EltPerPack=*/1> { + static constexpr bool IsIdentity = false; + + __device__ static BytePack preOp( + FuncPreMulSum fn, BytePack a + ) { + return toPack(rccl_bfloat8(float(fromPack(a)) * float(fn.scalar))); + } + }; +#endif + //////////////////////////////////////////////////////////////////////////////// // FuncSumPostDiv diff --git a/src/enqueue.cc b/src/enqueue.cc index 746da0e1c7..5c501796bf 100644 --- a/src/enqueue.cc +++ b/src/enqueue.cc @@ -1,6 +1,7 @@ /************************************************************************* * Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved. * Modifications Copyright (c) 2019-2023 Advanced Micro Devices, Inc. All rights reserved. + * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License. * * See LICENSE.txt for license information ************************************************************************/ @@ -1556,9 +1557,13 @@ static ncclResult_t hostToDevRedOp( half f16; float f32; double f64; - #if defined(RCCL_BFLOAT16) - rccl_bfloat16 bf16; - #endif +#if defined(RCCL_BFLOAT16) + rccl_bfloat16 bf16; +#endif +#if defined(RCCL_FLOAT8) + rccl_float8 fp8_e4m3; + rccl_bfloat8 fp8_e5m2; +#endif void *ptr; }; u64 = 0; @@ -1594,12 +1599,22 @@ static ncclResult_t hostToDevRedOp( opFull->op = ncclDevPreMulSum; f16 = __float2half(float(1.0/comm->nRanks)); // __double2half not supported pre CUDA 11.x break; - #if defined(RCCL_BFLOAT16) +#if defined(RCCL_BFLOAT16) case ncclBfloat16: opFull->op = ncclDevPreMulSum; bf16 = (rccl_bfloat16)(float(1.0/comm->nRanks)); break; - #endif +#endif +#if defined(RCCL_FLOAT8) + case ncclFp8E4M3: + opFull->op = ncclDevPreMulSum; + fp8_e4m3 = static_cast(float(1.0/comm->nRanks)); + break; + case ncclFp8E5M2: + opFull->op = ncclDevPreMulSum; + fp8_e5m2 = static_cast(float(1.0/comm->nRanks)); + break; +#endif case ncclFloat32: opFull->op = ncclDevPreMulSum; f32 = float(1.0/comm->nRanks); diff --git a/src/include/collectives.h b/src/include/collectives.h index e698fdb38a..6af4dfc418 100644 --- a/src/include/collectives.h +++ b/src/include/collectives.h @@ -1,6 +1,7 @@ /************************************************************************* * Copyright (c) 2017-2022, NVIDIA CORPORATION. All rights reserved. * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved. + * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License. * * See LICENSE.txt for license information ************************************************************************/ @@ -29,11 +30,15 @@ inline int ncclTypeSize(ncclDataType_t type) { switch (type) { case ncclInt8: case ncclUint8: +#if defined(RCCL_FLOAT8) + case ncclFp8E4M3: + case ncclFp8E5M2: +#endif return 1; case ncclFloat16: - #if defined(RCCL_BFLOAT16) +#if defined(RCCL_BFLOAT16) case ncclBfloat16: - #endif +#endif return 2; case ncclInt32: case ncclUint32: diff --git a/src/include/device.h b/src/include/device.h index c9441e9ce0..0fc2acacfd 100644 --- a/src/include/device.h +++ b/src/include/device.h @@ -1,6 +1,7 @@ /************************************************************************* * Copyright (c) 2015-2022, NVIDIA CORPORATION. All rights reserved. * Modifications Copyright (c) 2019-2022 Advanced Micro Devices, Inc. All rights reserved. + * Modifications Copyright (c) Microsoft Corporation. Licensed under the MIT License. * * See LICENSE.txt for license information ************************************************************************/ @@ -9,6 +10,7 @@ #define NCCL_DEVICE_H_ #include "nccl.h" +#include "rccl_float8.h" #include "rccl_bfloat16.h" #include "nccl_common.h" #include "align.h" @@ -502,9 +504,13 @@ inline bool ncclNvlsSupported(int devRedOp, int type) { case ncclInt64: case ncclUint64: case ncclFloat16: - #if defined(RCCL_BFLOAT16) +#if defined(RCCL_BFLOAT16) case ncclBfloat16: - #endif +#endif +#if defined(RCCL_FLOAT8) + case ncclFp8E4M3: + case ncclFp8E5M2: +#endif return devRedOp == ncclDevSum || devRedOp == ncclDevMinMax; case ncclFloat: case ncclDouble: @@ -530,10 +536,10 @@ inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto) // / / / if (coll == ncclFuncAllReduce) { - row += (((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 4 * (algo * NCCL_NUM_PROTOCOLS + proto); + row += (((algo * NCCL_NUM_PROTOCOLS + proto) * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 6 * (algo * NCCL_NUM_PROTOCOLS + proto); goto have_row; } - row += (NCCL_NUM_ALGORITHMS - 2) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 4); + row += (NCCL_NUM_ALGORITHMS - 2) * NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 6); // RING / SIMPLE / Sum / int8_t if (coll == ncclFuncAllToAllPivot) goto have_row; @@ -548,17 +554,17 @@ inline int ncclDevFuncId(int coll, int devRedOp, int type, int algo, int proto) // RING / / / if (coll == ncclFuncReduce) { - row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 4 * proto; + row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 6 * proto; goto have_row; } - row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 4); + row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 6); // RING / / / if (coll == ncclFuncReduceScatter) { - row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 4 * proto; + row += ((proto * ncclNumDevRedOps + devRedOp) * ncclNumTypes + type) - /*floats for each SumPostDiv*/ 6 * proto; goto have_row; } - row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 4); + row += NCCL_NUM_PROTOCOLS * (ncclNumDevRedOps * ncclNumTypes - /*floats for each SumPostDiv*/ 6); // RING / SIMPLE / Sum / int8_t if (coll == ncclFuncSendRecv) goto have_row; diff --git a/src/include/msccl/msccl_kernel.h b/src/include/msccl/msccl_kernel.h index 7fe496539d..be343a264f 100644 --- a/src/include/msccl/msccl_kernel.h +++ b/src/include/msccl/msccl_kernel.h @@ -26,7 +26,9 @@ __global__ void MSCCL_KERNEL_ENTRY_NAME(devredop, type, proto, fullOps)(struct n MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, half, fullOps) \ MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, float, fullOps) \ MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, double, fullOps) \ - MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps) + MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps) \ + MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_float8, fullOps) \ + MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, rccl_bfloat8, fullOps) #define MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_NOFLOAT(devredop, fullOps) \ MSCCL_DECL_KERNEL_ENTRY_FUNC_DEVREDOP_TYPE(devredop, int8_t, fullOps) \ diff --git a/src/include/nccl_common.h b/src/include/nccl_common.h index 00b582d199..e895bbcf2f 100644 --- a/src/include/nccl_common.h +++ b/src/include/nccl_common.h @@ -15,9 +15,9 @@ typedef void (*ncclDebugLogger_t)(ncclDebugLogLevel level, unsigned long flags, #define NCCL_NUM_FUNCTIONS 5 // Send/Recv and AllToAllPivot not included for now typedef enum { ncclFuncBroadcast, ncclFuncReduce, ncclFuncAllGather, ncclFuncReduceScatter, ncclFuncAllReduce, ncclFuncSendRecv, ncclFuncSend, ncclFuncRecv, ncclFuncAllToAllPivot, ncclNumFuncs} ncclFunc_t; -#define FUNC_INDEX_P2P 835 -#define FUNC_INDEX_ALLTOALL_PIVOT 555 -#define FUNC_INDEX_TOTAL 846 +#define FUNC_INDEX_P2P 979 +#define FUNC_INDEX_ALLTOALL_PIVOT 651 +#define FUNC_INDEX_TOTAL 992 #define NCCL_NUM_ALGORITHMS 6 // Tree/Ring/CollNet* #define NCCL_ALGO_UNDEF -1 diff --git a/src/include/rccl_float8.h b/src/include/rccl_float8.h new file mode 100644 index 0000000000..01cab41f71 --- /dev/null +++ b/src/include/rccl_float8.h @@ -0,0 +1,1021 @@ +/* ************************************************************************ + * Copyright (C) 2016-2023 Advanced Micro Devices, Inc. All rights reserved. + * + * Permission is hereby granted, free of charge, to any person obtaining a copy + * of this software and associated documentation files (the "Software"), to deal + * in the Software without restriction, including without limitation the rights + * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell cop- + * ies of the Software, and to permit persons to whom the Software is furnished + * to do so, subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IM- + * PLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER + * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNE- + * CTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + * + * ************************************************************************ */ + +#ifndef ROCBLAS_FLOAT8_H +#define ROCBLAS_FLOAT8_H + +#include + +#if __cplusplus < 201103L || (!defined(__HCC__) && !defined(__HIPCC__)) +/*! \brief Struct to represent a 8 bit floating-point number. */ + +typedef struct +{ + uint8_t data; +} rccl_float8; + +typedef struct +{ + uint8_t data; +} rccl_bfloat8; + +#else // __cplusplus < 201103L || (!defined(__HCC__) && !defined(__HIPCC__)) + +#define HIP_HOST_DEVICE __host__ __device__ +#define HIP_HOST __host__ +#define HIP_DEVICE __device__ + +// We are clipping in down conversion by default +#define rccl_float8_downcast_clipping 1 + +namespace rocblas_hip_f8_impl +{ + __host__ inline int clz(uint32_t x) + { + return __builtin_clz(x); + } + __device__ inline int clz(uint32_t x) + { + return __clz(x); + } + + template + HIP_HOST_DEVICE uint8_t cast_to_f8(T _x, bool stoch = false, uint32_t rng = 0) + { + constexpr bool is_half = std::is_same::value; + constexpr bool is_float = std::is_same::value; + static_assert(wm + we == 7, "wm+we==7"); + static_assert(is_half || is_float, "Only half and float can be cast to f8"); + + const int mfmt = (sizeof(T) == 4) ? 23 : 10; + uint32_t x; + if(sizeof(T) == 4) + x = reinterpret_cast(_x); + else + x = reinterpret_cast(_x); + + uint32_t y, head, mantissa; + int exponent, bias; + uint32_t sign; + + if(sizeof(T) == 4) + { + head = x & 0xFF800000; + mantissa = x & 0x7FFFFF; + exponent = (head >> 23) & 0xFF; + sign = head >> 31; + bias = 127; + } + else + { + head = x & 0xFC00; + mantissa = x & 0x3FF; + exponent = (head >> 10) & 0x1F; + sign = head >> 15; + bias = 15; + } + + uint32_t signed_inf = (sign << 7) + (((1 << we) - 1) << wm); + + // Deal with inf and NaNs + if(negative_zero_nan) + { + if(sizeof(T) == 4) + { + if((x & 0x7F800000) == 0x7F800000) + return 0x80; + } + else + { + //if(__hisinf(x) || __hisnan(x)) + if((x & 0x7C00) == 0x7C00) + return 0x80; + } + } + else + { + if(sizeof(T) == 4) + { + if((x & 0x7F800000) == 0x7F800000) + return signed_inf + (mantissa != 0 ? 1 : 0); + } + else + { + if((x & 0x7C00) == 0x7C00) + return signed_inf + (mantissa != 0 ? 1 : 0); + } + } + if(x == 0) + return 0; + + // First need to check if it is normal or denorm as there is a difference of implict 1 + // Then need to adjust the exponent to align with the F8 exponent, in the meanwhile, shift + // The mantissa. Then for stochastic rounding, add rng to mantissa and truncate. And for + // RNE, no need to add rng. Then probably need to check whether there is carry and adjust + // exponent and mantissa again + + // For IEEE bias mode, the bias is 2^(k-1) -1 where k is the width of exponent bits + const int f8_bias = (1 << (we - 1)) - 1 + (negative_zero_nan ? 1 : 0); + const int f8_denormal_act_exponent = 1 - f8_bias; //actual exponent of f8 denormal + // act_exponent is the actual exponent of fp32/fp16 (after subtracting bias) + // f8_exponent is the converted f8 exponent with bias encoding + // exponent_diff is the diff between fp32/fp16 exponent and f8 exponent, + // the difference needs to be adjusted and mantissa shifted + int act_exponent, f8_exponent, exponent_diff; + + if(exponent == 0) + { // fp32/fp16 is in denormal. + /* fp32 denormal is below 2^-127 so it is usually not a concern here, we mostly concern fp16 here. + In this case, f8 is usually in denormal. But there could be exceptions. + fp16 denormal has exponent bias 15 while bf8 with NANOO has exponent bias 16. + It means that there are some numbers in fp16 denormal but they are bf8 (NANOO) normals - smallest bf8 (NANOO) normal is 2^-15. + fp16 numbers where exponent==0 (actual exponent -14) and highest bit of mantissa is 1 are bf8 (NANOO) normal. + In this case, the fp16 mantissa should be shift left by 1 */ + act_exponent = exponent - bias + 1; + exponent_diff = f8_denormal_act_exponent + - act_exponent; // actual exponent is exponent-bias+1 as it is denormal + } + else + { // fp32/fp16 is normal with implicit 1 + act_exponent = exponent - bias; + if(act_exponent <= f8_denormal_act_exponent) + { + /* This is the case where fp32/fp16 is normal but it is in f8 denormal range. + For example fp8 nanoo mode, denormal exponent is -7, but if the fp32/fp16 + actual exponent is -7, it is actually larger due to the implict 1, + Therefore it needs to be adjust to -6 and mantissa shift right by 1. + So for fp32/fp16, exponent -8 is the cut point to convert to fp8 nanoo */ + exponent_diff = f8_denormal_act_exponent - act_exponent; + } + else + { //both fp32/fp16 and f8 are in normal range + exponent_diff + = 0; // exponent_diff=0 does not mean there is no difference for this case, + //act_exponent could be larger. Just that it does not need shift mantissa + } + mantissa += (1 << mfmt); //Add the implicit 1 into mantissa + } + + bool midpoint = (mantissa & ((1 << (mfmt - wm + exponent_diff)) - 1)) + == (1 << (mfmt - wm + exponent_diff - 1)); + /* This part is a bit tricky. The judgment of whether it is a tie needs to be done before we shift right + as shift right could rip off some residual part and make something not midpoint look like midpoint. + For example, the fp16 number 0x1002 (0 00100 0000000010), it is larger than midpoint, + but after shift right by 4 bits, it would look like midpoint. + */ + + if(exponent_diff > 0) + mantissa >>= exponent_diff; + else if(exponent_diff == -1) + mantissa <<= -exponent_diff; + bool implicit_one = mantissa & (1 << mfmt); + //if there is no implict 1, it means the f8 is denormal and need to adjust to denorm exponent + f8_exponent = (act_exponent + exponent_diff) /*actual f8 exponent*/ + f8_bias + - (implicit_one ? 0 : 1); + + //Now we have the exponent and mantissa adjusted + uint32_t drop_mask = (1 << (mfmt - wm)) - 1; + bool odd = mantissa + & (1 << (mfmt - wm)); // if the least significant bit that is not truncated is 1 + mantissa + += (stoch ? rng : (midpoint ? (odd ? mantissa : mantissa - 1) : mantissa)) & drop_mask; + + //Now we deal with overflow + if(f8_exponent == 0) + { + if((1 << mfmt) & mantissa) + { + f8_exponent = 1; //denormal overflow to become normal, promote exponent + } + } + else + { + if((1 << (mfmt + 1)) & mantissa) + { + mantissa >>= 1; + f8_exponent++; + } + } + + mantissa >>= (mfmt - wm); + + // above range: quantize to maximum possible float of the same sign + const int max_exp = (1 << we) - (negative_zero_nan ? 1 : 2); + if(f8_exponent > max_exp) + { + if(clip) + { + mantissa = (1 << wm) - 1; + f8_exponent = max_exp; + } + else + { + return signed_inf; + } + } + + if(f8_exponent == 0 && mantissa == 0) + return negative_zero_nan ? 0 : (sign << 7); + mantissa &= (1 << wm) - 1; + return (sign << 7) | (f8_exponent << wm) | mantissa; + } + + template + HIP_HOST_DEVICE T cast_from_f8(uint8_t x) + { + constexpr bool is_half = std::is_same::value; + constexpr bool is_float = std::is_same::value; + static_assert(is_half || is_float, "only half and float are supported"); + + constexpr int weo = is_half ? 5 : 8; + constexpr int wmo = is_half ? 10 : (is_float ? 23 : 7); + + T fInf, fNegInf, fNaN, fNeg0; + if(is_half) + { + const uint16_t ihInf = 0x7C00; + const uint16_t ihNegInf = 0xFC00; + const uint16_t ihNaN = 0x7C01; + const uint16_t ihNeg0 = 0x8000; + fInf = reinterpret_cast(ihInf); + fNegInf = reinterpret_cast(ihNegInf); + fNaN = reinterpret_cast(ihNaN); + fNeg0 = reinterpret_cast(ihNeg0); + } + else if(is_float) + { + const uint32_t ifInf = 0x7F800000; + const uint32_t ifNegInf = 0xFF800000; + const uint32_t ifNaN = 0x7F800001; + const uint32_t ifNeg0 = 0x80000000; + fInf = reinterpret_cast(ifInf); + fNegInf = reinterpret_cast(ifNegInf); + fNaN = reinterpret_cast(ifNaN); + fNeg0 = reinterpret_cast(ifNeg0); + } + + if(x == 0) + return 0; + + uint32_t sign = x >> 7; + uint32_t mantissa = x & ((1 << wm) - 1); + int exponent = (x & 0x7F) >> wm; + if(negative_zero_nan) + { + if(x == 0x80) + return fNaN; + } + else + { + if(x == 0x80) + return fNeg0; + if(exponent == ((1 << we) - 1)) + return (mantissa == 0) ? (sign ? fNegInf : fInf) : fNaN; + } + typename std::conditional::type retval; + if(we == 5 && is_half && !negative_zero_nan) + { + retval = x << 8; + return reinterpret_cast(retval); + } + + const int exp_low_cutoff + = (1 << (weo - 1)) - (1 << (we - 1)) + 1 - (negative_zero_nan ? 1 : 0); + + //subnormal input + if(exponent == 0) + { + //guaranteed mantissa!=0 since cases 0x0 and 0x80 are handled above + int sh = 1 + clz(mantissa) - (32 - wm); + mantissa <<= sh; + exponent += 1 - sh; + mantissa &= ((1 << wm) - 1); + } + exponent += exp_low_cutoff - 1; + mantissa <<= wmo - wm; + + // subnormal output (occurs when T=half, we=5, negative_zero_nan=true) + if(exponent <= 0) + { + mantissa |= 1 << wmo; + mantissa >>= 1 - exponent; + exponent = 0; + } + + if(sizeof(T) == 2) + retval = (sign << 15) | (exponent << 10) | mantissa; + else + retval = (sign << 31) | (exponent << 23) | mantissa; + return reinterpret_cast(retval); + } +} // namespace rocblas_hip_f8_impl + +static __device__ bool rocblas_hip_f8_bias_mode_bit_device = true; +static bool rocblas_hip_f8_bias_mode_bit_host = true; + +struct rccl_float8 +{ + uint8_t data; + enum class rocblas_hip_f8_rounding_mode + { + standard, + stochastic + }; + + // default constructor + HIP_HOST_DEVICE rccl_float8() = default; + +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + // device specific optimized F8 down-conversion code + + template + static HIP_DEVICE uint8_t cast_to_f8_from_f32(float v, uint32_t rng = 0) + { + uint8_t i8data; + union + { + float fval; + uint32_t i32val; + uint8_t i8val[4]; // NOTE: not endian independent + } val; + + uint32_t ival = 0; + val.fval = v; + +#ifdef rccl_float8_downcast_clipping + if((val.i32val & 0x7F800000) != 0x7F800000) /// propagate NAN/INF, no clipping + val.fval = __builtin_amdgcn_fmed3f(val.fval, 240.0, -240.0); +#endif + if(stochastic_rounding) + { + ival = __builtin_amdgcn_cvt_sr_fp8_f32(val.fval, rng, ival, 0); // 0 pos + val.i32val = ival; + i8data = val.i8val[0]; // little endian + } + else // RNE CVT + { + ival = __builtin_amdgcn_cvt_pk_fp8_f32( + val.fval, val.fval, ival, false); // false -> WORD0 + val.i32val = ival; + i8data = val.i8val[0]; + } + return i8data; + } + +#endif // __gfx940__ + + // constructor from float +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + + // NOTE: ON-DEVICE... always optimal bias + explicit HIP_DEVICE rccl_float8(float v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + { + // runtime branch, use cast_to_f8_from_f32 if want to avoid it + if(rm == rocblas_hip_f8_rounding_mode::stochastic) + data = cast_to_f8_from_f32(v, rng); + else + data = cast_to_f8_from_f32(v); + } + + // Host only implementation using s/w simulation + explicit HIP_HOST +#else + // both Host and DEVICE for non-gfx940 using s/w simulation + explicit HIP_HOST_DEVICE +#endif + rccl_float8(float v, + rocblas_hip_f8_rounding_mode rm = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + { +#ifdef rccl_float8_downcast_clipping + data = rocblas_hip_f8_impl:: + cast_to_f8<3, 4, float, true /*negative_zero_nan*/, true /*clip*/>( + v, (rm == rocblas_hip_f8_rounding_mode::stochastic), rng); +#else // rccl_float8_downcast_clipping + data = rocblas_hip_f8_impl:: + cast_to_f8<3, 4, float, true /*negative_zero_nan*/, false /*clip*/>( + v, (rm == rocblas_hip_f8_rounding_mode::stochastic), rng); +#endif // rccl_float8_downcast_clipping + } + + // Constructor from half + explicit HIP_HOST_DEVICE rccl_float8(_Float16 v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_float8((float)v, rm, rng) + { + } + // constructor from int + explicit HIP_HOST_DEVICE rccl_float8(int v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_float8((float)v, rm, rng) + { + } + // constructor from double + explicit HIP_HOST_DEVICE rccl_float8(double v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_float8((float)v, rm, rng) + { + } + + // convert to float +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + // upcast using device specific intrinsic + explicit inline HIP_DEVICE operator float() const + { + float fval; + uint32_t i32val = static_cast(data); + + // upcast + asm volatile("v_cvt_f32_fp8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); + + return fval; + } + + explicit inline HIP_HOST operator float() const +#else // non gfx940 + explicit inline HIP_HOST_DEVICE operator float() const +#endif + { + return rocblas_hip_f8_impl::cast_from_f8<3, 4, float, true /*negative_zero_nan*/>(data); + } + + // convert to half + explicit inline HIP_HOST_DEVICE operator _Float16() const + { + return _Float16(float(*this)); // convert to float, then convert to f16 + } + + // check for zero + inline HIP_HOST_DEVICE bool is_zero() const + { + return data == 0x00; + } + + // check for nan + inline HIP_HOST_DEVICE bool is_nan() const + { + return data == 0x80; + } + + // check for inf + inline HIP_HOST_DEVICE bool is_inf() const + { + return data == 0x80; + } + + // assignment overloading only from the same F8 types + inline __host__ __device__ rccl_float8& operator=(const rccl_float8& a) + { + data = a.data; + return *this; + } +}; + +struct rccl_bfloat8 +{ + uint8_t data; + enum class rocblas_hip_f8_rounding_mode + { + standard, + stochastic + }; + + // default constructor + HIP_HOST_DEVICE rccl_bfloat8() = default; + +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + // device specific optimized F8 down-conversion code + + template + static HIP_DEVICE uint8_t cast_to_bf8_from_f32(float v, uint32_t rng = 0) + { + uint8_t i8data; + union + { + float fval; + uint32_t i32val; + uint8_t i8val[4]; // NOTE: not endian independent + } val; + + uint32_t ival = 0; + val.fval = v; + +#ifdef rccl_float8_downcast_clipping + if((val.i32val & 0x7F800000) != 0x7F800000) // propagate NAN/INF, no clipping + val.fval = __builtin_amdgcn_fmed3f(val.fval, 57344.0, -57344.0); +#endif + if(stochastic_rounding) + { + ival = __builtin_amdgcn_cvt_sr_bf8_f32(val.fval, rng, ival, 0); // 0 pos + val.i32val = ival; + i8data = val.i8val[0]; // little endian + } + else // RNE CVT + { + ival = __builtin_amdgcn_cvt_pk_bf8_f32( + val.fval, val.fval, ival, false); // false -> WORD0 + val.i32val = ival; + i8data = val.i8val[0]; + } + return i8data; + } + +#endif // __gfx940__ + + // constructor from float +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + + // NOTE: ON-DEVICE... always optimal bias + explicit HIP_DEVICE rccl_bfloat8(float v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + { + // runtime branch, use cast_to_f8_from_f32 if want to avoid it + if(rm == rocblas_hip_f8_rounding_mode::stochastic) + data = cast_to_bf8_from_f32(v, rng); + else + data = cast_to_bf8_from_f32(v); + } + + // Host only implementation using s/w simulation + explicit HIP_HOST +#else + // both Host and DEVICE for non-gfx940 using s/w simulation + explicit HIP_HOST_DEVICE +#endif + rccl_bfloat8(float v, + rocblas_hip_f8_rounding_mode rm = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + { +#ifdef rccl_float8_downcast_clipping + data = rocblas_hip_f8_impl:: + cast_to_f8<2, 5, float, true /*negative_zero_nan*/, true /*clip*/>( + v, (rm == rocblas_hip_f8_rounding_mode::stochastic), rng); +#else + data = rocblas_hip_f8_impl:: + cast_to_f8<2, 5, float, true /*negative_zero_nan*/, false /*clip*/>( + v, (rm == rocblas_hip_f8_rounding_mode::stochastic), rng); +#endif // rccl_float8_downcast_clipping + } + + // Constructor from half + explicit HIP_HOST_DEVICE rccl_bfloat8(_Float16 v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_bfloat8((float)v, rm, rng) + { + } + // constructor from int + explicit HIP_HOST_DEVICE rccl_bfloat8(int v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_bfloat8((float)v, rm, rng) + { + } + // constructor from double + explicit HIP_HOST_DEVICE rccl_bfloat8(double v, + rocblas_hip_f8_rounding_mode rm + = rocblas_hip_f8_rounding_mode::standard, + uint32_t rng = 0) + : rccl_bfloat8((float)v, rm, rng) + { + } + + // convert to float +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + // upcast using device specific intrinsic + explicit inline HIP_DEVICE operator float() const + { + float fval; + uint32_t i32val = static_cast(data); + + // upcast + asm volatile("v_cvt_f32_bf8 %0, %1 src0_sel:BYTE_0" : "=v"(fval) : "v"(i32val)); + + return fval; + } + + explicit inline HIP_HOST operator float() const +#else // non gfx940 + explicit inline HIP_HOST_DEVICE operator float() const +#endif + { + return rocblas_hip_f8_impl::cast_from_f8<2, 5, float, true /*negative_zero_nan*/>(data); + } + + // convert to half + explicit inline HIP_HOST_DEVICE operator _Float16() const + { + return _Float16(float(*this)); // convert to float, then convert to f16 + } + + // check for zero + inline HIP_HOST_DEVICE bool is_zero() const + { + return data == 0x00; + } + + // check for nan + inline HIP_HOST_DEVICE bool is_nan() const + { + return data == 0x80; + } + + // check for inf + inline HIP_HOST_DEVICE bool is_inf() const + { + return data == 0x80; + } + + // assignment overloading only from the same F8 types + inline __host__ __device__ rccl_bfloat8& operator=(const rccl_bfloat8& a) + { + data = a.data; + return *this; + } +}; + +namespace std +{ + inline rccl_float8 sin(rccl_float8 a) + { + return rccl_float8(sinf(float(a))); + } + inline rccl_float8 cos(rccl_float8 a) + { + return rccl_float8(cosf(float(a))); + } + inline rccl_bfloat8 sin(rccl_bfloat8 a) + { + return rccl_bfloat8(sinf(float(a))); + } + inline rccl_bfloat8 cos(rccl_bfloat8 a) + { + return rccl_bfloat8(cosf(float(a))); + } + __device__ __host__ constexpr rccl_float8 real(const rccl_float8& a) + { + return a; + } + __device__ __host__ constexpr rccl_bfloat8 real(const rccl_bfloat8& a) + { + return a; + } +} + +// Special operator overloading +inline std::ostream& operator<<(std::ostream& os, const rccl_float8& f8) +{ + return os << float(f8); +} + +inline std::ostream& operator<<(std::ostream& os, const rccl_bfloat8& bf8) +{ + return os << float(bf8); +} + +// all + operator overloading with mixed types +// mixed types, always converts to f32, does computation in f32, and returns float +inline __host__ __device__ float operator+(const float fa, rccl_float8 b) +{ + return (fa + float(b)); +} + +inline __host__ __device__ float operator+(const float fa, rccl_bfloat8 b) +{ + return (fa + float(b)); +} + +inline __host__ __device__ float operator+(rccl_float8 a, const float fb) +{ + return (float(a) + fb); +} + +inline __host__ __device__ float operator+(rccl_bfloat8 a, const float fb) +{ + return (float(a) + fb); +} + +inline __host__ __device__ float operator+(rccl_float8 a, rccl_bfloat8 b) +{ + return (float(a) + float(b)); +} + +inline __host__ __device__ float operator+(rccl_bfloat8 a, rccl_float8 b) +{ + return (float(a) + float(b)); +} + +inline __host__ __device__ rccl_float8 operator+(rccl_float8 a, rccl_float8 b) +{ + return rccl_float8(float(a) + float(b)); +} + +inline __host__ __device__ rccl_bfloat8 operator+(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return rccl_bfloat8(float(a) + float(b)); +} + +inline __host__ __device__ rccl_float8& operator+=(rccl_float8& a, rccl_float8 b) +{ + return a = rccl_float8(float(a) + float(b)); +} + +inline __host__ __device__ rccl_bfloat8& operator+=(rccl_bfloat8& a, rccl_bfloat8 b) +{ + return a = rccl_bfloat8(float(a) + float(b)); +} + +// overloading multiplication, always returns float, +inline __host__ __device__ float operator*(rccl_float8 a, rccl_float8 b) +{ + return float(a) * float(b); +} + +inline __host__ __device__ float operator*(float a, rccl_float8 b) +{ + return (a * float(b)); +} + +inline __host__ __device__ float operator*(rccl_float8 a, float b) +{ + return (float(a) * b); +} + +inline __host__ __device__ float operator*(int32_t a, rccl_float8 b) +{ + return ((float)a * float(b)); +} + +inline __host__ __device__ float operator*(double a, rccl_float8 b) +{ + return ((float)a * float(b)); +} + +inline __host__ __device__ float operator*(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return float(a) * float(b); +} + +inline __host__ __device__ float operator*(float a, rccl_bfloat8 b) +{ + return (a * float(b)); +} + +inline __host__ __device__ float operator*(rccl_bfloat8 a, float b) +{ + return (float(a) * b); +} + +inline __host__ __device__ float operator*(int32_t a, rccl_bfloat8 b) +{ + return ((float)a * float(b)); +} + +inline __host__ __device__ float operator*(double a, rccl_bfloat8 b) +{ + return ((float)a * float(b)); +} + +// overloading for mixed f8 and bf8 types +inline __host__ __device__ float operator*(rccl_float8 a, rccl_bfloat8 b) +{ + return float(a) * float(b); +} + +inline __host__ __device__ float operator*(rccl_bfloat8 a, rccl_float8 b) +{ + return float(a) * float(b); +} + +// all - operator overloading with mixed types +// mixed types, always converts to f32, does computation in f32, and returns float +inline __host__ __device__ float operator-(const float fa, rccl_float8 b) +{ + return (fa - float(b)); +} + +inline __host__ __device__ float operator-(const float fa, rccl_bfloat8 b) +{ + return (fa - float(b)); +} + +inline __host__ __device__ float operator-(rccl_float8 a, const float fb) +{ + return (float(a) - fb); +} + +inline __host__ __device__ float operator-(rccl_bfloat8 a, const float fb) +{ + return (float(a) - fb); +} + +inline __host__ __device__ float operator-(rccl_float8 a, rccl_bfloat8 b) +{ + return (float(a) - float(b)); +} + +inline __host__ __device__ float operator-(rccl_bfloat8 a, rccl_float8 b) +{ + return (float(a) - float(b)); +} + +inline __host__ __device__ rccl_float8 operator-(rccl_float8 a, rccl_float8 b) +{ + return rccl_float8(float(a) - float(b)); +} + +inline __host__ __device__ rccl_bfloat8 operator-(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return rccl_bfloat8(float(a) - float(b)); +} + +inline __host__ __device__ rccl_float8& operator-=(rccl_float8& a, rccl_float8 b) +{ + return a = rccl_float8(float(a) - float(b)); +} + +inline __host__ __device__ rccl_bfloat8& operator-=(rccl_bfloat8& a, rccl_bfloat8 b) +{ + return a = rccl_bfloat8(float(a) - float(b)); +} + +// overloading division, always returns float, +inline __host__ __device__ float operator/(rccl_float8 a, rccl_float8 b) +{ + return float(a) / float(b); +} + +inline __host__ __device__ float operator/(float a, rccl_float8 b) +{ + return (a / float(b)); +} + +inline __host__ __device__ float operator/(rccl_float8 a, float b) +{ + return (float(a) / b); +} + +inline __host__ __device__ float operator/(int32_t a, rccl_float8 b) +{ + return ((float)a / float(b)); +} + +inline __host__ __device__ float operator/(double a, rccl_float8 b) +{ + return ((float)a / float(b)); +} + +inline __host__ __device__ float operator/(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return float(a) / float(b); +} + +inline __host__ __device__ float operator/(float a, rccl_bfloat8 b) +{ + return (a / float(b)); +} + +inline __host__ __device__ float operator/(rccl_bfloat8 a, float b) +{ + return (float(a) / b); +} + +inline __host__ __device__ float operator/(int32_t a, rccl_bfloat8 b) +{ + return ((float)a / float(b)); +} + +inline __host__ __device__ float operator/(double a, rccl_bfloat8 b) +{ + return ((float)a / float(b)); +} + +// overloading for mixed f8 and bf8 types +inline __host__ __device__ float operator/(rccl_float8 a, rccl_bfloat8 b) +{ + return float(a) / float(b); +} + +inline __host__ __device__ float operator/(rccl_bfloat8 a, rccl_float8 b) +{ + return float(a) / float(b); +} + +// overloading for compare +inline __host__ __device__ bool operator==(rccl_float8 a, rccl_float8 b) +{ + return (a.data == b.data); +} + +inline __host__ __device__ bool operator==(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return (a.data == b.data); +} + +inline __host__ __device__ bool operator!=(rccl_float8 a, rccl_float8 b) +{ + return (a.data != b.data); +} + +inline __host__ __device__ bool operator!=(rccl_bfloat8 a, rccl_bfloat8 b) +{ + return (a.data != b.data); +} + +// ================ Explicit downcasting to support different rounding (RNE, SR) =============== +// NOTE: we going to remove all assignment operator overloading from other types and enforce +// this explicit_downcast function to make any roudning behavior default +// We have to explicitly call this function with SR flag + +template {}, int>::type = 0> +inline __host__ __device__ T explicit_downcast(Ta a, uint32_t rng = 0) +{ + // same type, no conversion + return a; +} + +// Use h/w intrinsic and optimized version when __gfx940__ +template < + typename T, + typename Ta, + bool stochastic_rounding, + typename std::enable_if<(!(std::is_same{}) + && (std::is_same{} || std::is_same{})), + int>::type + = 0> +inline __host__ __device__ T explicit_downcast(Ta a, uint32_t rng) +{ +#if defined(__gfx940__) || defined(__gfx941__) || defined(__gfx942__) + // NOTE: we are directly calling cast_to_f8_from_f32 instead of constructor to optimize away one runtime branch + T val; + if(std::is_same::value) + val.data = rccl_float8::cast_to_f8_from_f32(float(a), rng); + else + val.data = rccl_bfloat8::cast_to_bf8_from_f32(float(a), rng); + return val; +#else // non gfx940 + return T(float(a), + stochastic_rounding ? T::rocblas_hip_f8_rounding_mode::stochastic + : T::rocblas_hip_f8_rounding_mode::standard, + rng); +#endif // __gfx940__ +} + +// NOTE NOTE: The above code is good if we don't consider HIP-GEMM code and only consider the quantization +// However, if we need HIP-GEMM for fall-back, we would need explicit_cast handles Tacc=f32 to To=f16/bf16 conversion +template < + typename T, + typename Ta, + bool stochastic_rounding, + typename std::enable_if<(!(std::is_same{}) + && !(std::is_same{} || std::is_same{})), + int>::type + = 0> +inline __host__ __device__ T explicit_downcast(Ta a, uint32_t rng) +{ + // the return type is not a F8 types, no SR for those types + // not sure if we have direct conversion, so converting to float first + // no effect if the input type is float + return T(float(a)); +} + +// ================================================================================================= + +#endif // __cplusplus < 201103L || (!defined(__HCC__) && !defined(__HIPCC__)) + +#endif // ROCBLAS_FLOAT8_H diff --git a/src/misc/msccl/msccl_setup.cc b/src/misc/msccl/msccl_setup.cc index f8ee6be403..88b379f988 100644 --- a/src/misc/msccl/msccl_setup.cc +++ b/src/misc/msccl/msccl_setup.cc @@ -226,6 +226,10 @@ static ncclResult_t hostToDevRedOp( #if defined(RCCL_BFLOAT16) rccl_bfloat16 bf16; #endif + #if defined(RCCL_FLOAT8) + rccl_float8 fp8_e4m3; + rccl_bfloat8 fp8_e5m2; + #endif float f32; double f64; void *ptr; @@ -269,6 +273,16 @@ static ncclResult_t hostToDevRedOp( bf16 = (rccl_bfloat16)(float(1.0/comm->nRanks)); break; #endif + #if defined(RCCL_FLOAT8) + case ncclFp8E4M3: + opFull->op = ncclDevPreMulSum; + fp8_e4m3 = (rccl_float8)(float(1.0/comm->nRanks)); + break; + case ncclFp8E5M2: + opFull->op = ncclDevPreMulSum; + fp8_e5m2 = (rccl_bfloat8)(float(1.0/comm->nRanks)); + break; + #endif case ncclFloat32: opFull->op = ncclDevPreMulSum; f32 = float(1.0/comm->nRanks); @@ -315,7 +329,9 @@ static ncclResult_t hostToDevRedOp( MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, half, fullOps), \ MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, float, fullOps), \ MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, double, fullOps), \ - MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps) + MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, rccl_bfloat16, fullOps), \ + MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, rccl_float8, fullOps), \ + MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, rccl_bfloat8, fullOps) #define MSCCL_KERNEL_ENTRY_DEVREDOP_NOFLOAT(devredop, fullOps) \ MSCCL_KERNEL_ENTRY_DEVREDOP_TYPE(devredop, int8_t, fullOps), \ @@ -327,6 +343,8 @@ static ncclResult_t hostToDevRedOp( MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \ MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \ MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \ + MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \ + MSCCL_KERNEL_ENTRY_DEVREDOP_NULL(), \ MSCCL_KERNEL_ENTRY_DEVREDOP_NULL() #define MSCCL_KERNEL_ENTRY() \ diff --git a/src/nccl.h.in b/src/nccl.h.in index f9be58cb22..63d6d8ff89 100644 --- a/src/nccl.h.in +++ b/src/nccl.h.in @@ -21,6 +21,7 @@ #define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z)) #define RCCL_BFLOAT16 1 +#define RCCL_FLOAT8 1 #define RCCL_GATHER_SCATTER 1 #define RCCL_ALLTOALLV 1 @@ -362,7 +363,13 @@ typedef enum { ncclInt8 = 0, ncclChar = 0, ncclFloat32 = 7, ncclFloat = 7, ncclFloat64 = 8, ncclDouble = 8, ncclBfloat16 = 9, +#if defined(RCCL_FLOAT8) + ncclFp8E4M3 = 10, + ncclFp8E5M2 = 11, + ncclNumTypes = 12 } ncclDataType_t; +#else ncclNumTypes = 10 } ncclDataType_t; +#endif /*! @} */ /*! @defgroup rccl_api_custom_redop Custom Reduction Operator diff --git a/test/AllGatherTests.cpp b/test/AllGatherTests.cpp index 1be1a7f42e..067aa59002 100644 --- a/test/AllGatherTests.cpp +++ b/test/AllGatherTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollAllGather}; - std::vector const dataTypes = {ncclBfloat16, ncclFloat64}; + std::vector const dataTypes = {ncclBfloat16, ncclFloat64, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclSum}; std::vector const roots = {0}; std::vector const numElements = {586}; diff --git a/test/AllReduceTests.cpp b/test/AllReduceTests.cpp index aeb02fb4f3..d0ef4c6188 100644 --- a/test/AllReduceTests.cpp +++ b/test/AllReduceTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollAllReduce}; - std::vector const dataTypes = {ncclFloat16, ncclFloat64}; + std::vector const dataTypes = {ncclFloat16, ncclFloat64, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclMin}; std::vector const roots = {0}; std::vector const numElements = {12888}; @@ -70,7 +70,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollAllReduce}; - std::vector const dataTypes = {ncclInt32}; + std::vector const dataTypes = {ncclInt32, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclMax}; std::vector const roots = {0}; std::vector const numElements = {393216, 12888, 384}; diff --git a/test/AllToAllTests.cpp b/test/AllToAllTests.cpp index 5c8ea0b420..9ceaa94c3c 100644 --- a/test/AllToAllTests.cpp +++ b/test/AllToAllTests.cpp @@ -35,7 +35,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollAllToAll}; - std::vector const dataTypes = {ncclFloat64, ncclBfloat16}; + std::vector const dataTypes = {ncclFloat64, ncclBfloat16, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclSum}; std::vector const roots = {0}; std::vector const numElements = {5685}; diff --git a/test/BroadcastTests.cpp b/test/BroadcastTests.cpp index d925ef4d68..32511c39c9 100644 --- a/test/BroadcastTests.cpp +++ b/test/BroadcastTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollBroadcast}; - std::vector const dataTypes = {ncclBfloat16, ncclFloat64}; + std::vector const dataTypes = {ncclBfloat16, ncclFloat64, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclSum}; std::vector const roots = {0}; std::vector const numElements = {586}; diff --git a/test/ReduceScatterTests.cpp b/test/ReduceScatterTests.cpp index ef7f8f701b..23b0f56289 100644 --- a/test/ReduceScatterTests.cpp +++ b/test/ReduceScatterTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollReduceScatter}; - std::vector const dataTypes = {ncclFloat64, ncclBfloat16}; + std::vector const dataTypes = {ncclFloat64, ncclBfloat16, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclMax}; std::vector const roots = {0}; std::vector const numElements = {1048576}; diff --git a/test/ReduceTests.cpp b/test/ReduceTests.cpp index 5cc8b5f755..8fab6942d9 100644 --- a/test/ReduceTests.cpp +++ b/test/ReduceTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollReduce}; - std::vector const dataTypes = {ncclFloat16, ncclFloat64}; + std::vector const dataTypes = {ncclFloat16, ncclFloat64, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclMin}; std::vector const roots = {0}; std::vector const numElements = {393216}; @@ -70,7 +70,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollReduce}; - std::vector const dataTypes = {ncclBfloat16}; + std::vector const dataTypes = {ncclBfloat16, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclMax}; std::vector const roots = {0}; std::vector const numElements = {393216}; diff --git a/test/ScatterTests.cpp b/test/ScatterTests.cpp index c0b906475b..d7bd7ab083 100644 --- a/test/ScatterTests.cpp +++ b/test/ScatterTests.cpp @@ -32,7 +32,7 @@ namespace RcclUnitTesting // Configuration std::vector const funcTypes = {ncclCollScatter}; - std::vector const dataTypes = {ncclFloat64, ncclBfloat16}; + std::vector const dataTypes = {ncclFloat64, ncclBfloat16, ncclFp8E4M3, ncclFp8E5M2}; std::vector const redOps = {ncclSum}; std::vector const roots = {1}; std::vector const numElements = {24658}; diff --git a/test/common/CollectiveArgs.cpp b/test/common/CollectiveArgs.cpp index 36c9b45e40..6908f45e95 100644 --- a/test/common/CollectiveArgs.cpp +++ b/test/common/CollectiveArgs.cpp @@ -199,8 +199,10 @@ namespace RcclUnitTesting case ncclUint32: ss << scalarsPerRank.U4[this->globalRank]; break; case ncclInt64: ss << scalarsPerRank.I8[this->globalRank]; break; case ncclUint64: ss << scalarsPerRank.U8[this->globalRank]; break; + case ncclFp8E4M3: ss << scalarsPerRank.F1[this->globalRank]; break; case ncclFloat32: ss << scalarsPerRank.F4[this->globalRank]; break; case ncclFloat64: ss << scalarsPerRank.F8[this->globalRank]; break; + case ncclFp8E5M2: ss << scalarsPerRank.B1[this->globalRank]; break; case ncclBfloat16: ss << scalarsPerRank.B2[this->globalRank]; break; default: ss << "(UNKNOWN)"; } diff --git a/test/common/CollectiveArgs.hpp b/test/common/CollectiveArgs.hpp index 7bdd005f93..61347d44ee 100644 --- a/test/common/CollectiveArgs.hpp +++ b/test/common/CollectiveArgs.hpp @@ -53,7 +53,9 @@ namespace RcclUnitTesting "ncclFloat16", "ncclFloat32", "ncclFloat64", - "ncclBfloat16" + "ncclBfloat16", + "ncclFp8E4M3", + "ncclFp8E5M2" }; char const ncclRedOpNames[ncclNumOps][32] = diff --git a/test/common/EnvVars.cpp b/test/common/EnvVars.cpp index 07a9561d7b..db161b8160 100644 --- a/test/common/EnvVars.cpp +++ b/test/common/EnvVars.cpp @@ -114,6 +114,8 @@ namespace RcclUnitTesting dataTypes.push_back(ncclFloat32); dataTypes.push_back(ncclFloat64); dataTypes.push_back(ncclBfloat16); + dataTypes.push_back(ncclFp8E4M3); + dataTypes.push_back(ncclFp8E5M2); } // Build list of possible # GPU ranks based on env vars diff --git a/test/common/PtrUnion.cpp b/test/common/PtrUnion.cpp index 7602ea2e7d..ea3aaf922d 100644 --- a/test/common/PtrUnion.cpp +++ b/test/common/PtrUnion.cpp @@ -14,6 +14,8 @@ namespace RcclUnitTesting { case ncclInt8: return 1; case ncclUint8: return 1; + case ncclFp8E4M3:return 1; + case ncclFp8E5M2:return 1; case ncclInt32: return 4; case ncclUint32: return 4; case ncclInt64: return 8; @@ -160,9 +162,11 @@ namespace RcclUnitTesting case ncclUint32: U4[idx] = valueI; break; case ncclInt64: I8[idx] = valueI; break; case ncclUint64: U8[idx] = valueI; break; + case ncclFp8E4M3: F1[idx] = rccl_float8(valueF); break; case ncclFloat16: F2[idx] = __float2half(static_cast(valueF)); break; case ncclFloat32: F4[idx] = valueF; break; case ncclFloat64: F8[idx] = valueF; break; + case ncclFp8E5M2: B1[idx] = rccl_bfloat8(valueF); break; case ncclBfloat16: B2[idx] = rccl_bfloat16(static_cast(valueF)); break; default: ERROR("Unsupported datatype\n"); @@ -181,9 +185,11 @@ namespace RcclUnitTesting case ncclUint32: valueI = U4[idx]; break; case ncclInt64: valueI = I8[idx]; break; case ncclUint64: valueI = U8[idx]; break; + case ncclFp8E4M3: valueF = float(F1[idx]); break; case ncclFloat16: valueF = __half2float(F2[idx]); break; case ncclFloat32: valueF = F4[idx]; break; case ncclFloat64: valueF = F8[idx]; break; + case ncclFp8E5M2: valueF = float(B1[idx]); break; case ncclBfloat16: valueF = B2[idx]; break; default: ERROR("Unsupported datatype\n"); @@ -211,9 +217,11 @@ namespace RcclUnitTesting case ncclUint32: U4[idx] *= scalarsPerRank.U4[rank]; break; case ncclInt64: I8[idx] *= scalarsPerRank.I8[rank]; break; case ncclUint64: U8[idx] *= scalarsPerRank.U8[rank]; break; + case ncclFp8E4M3: F1[idx] = rccl_float8(F1[idx] * scalarsPerRank.F1[rank]); break; case ncclFloat16: F2[idx] = __float2half(__half2float(F2[idx]) * __half2float(scalarsPerRank.F2[rank])); break; case ncclFloat32: F4[idx] *= scalarsPerRank.F4[rank]; break; case ncclFloat64: F8[idx] *= scalarsPerRank.F8[rank]; break; + case ncclFp8E5M2: B1[idx] = rccl_bfloat8(B1[idx] * scalarsPerRank.B1[rank]); break; case ncclBfloat16: B2[idx] *= scalarsPerRank.B2[rank]; break; default: ERROR("Unsupported datatype\n"); @@ -244,9 +252,11 @@ namespace RcclUnitTesting case ncclUint32: U4[idx] = ReduceOp(op, U4[idx], inputCpu.U4[idx]); break; case ncclInt64: I8[idx] = ReduceOp(op, I8[idx], inputCpu.I8[idx]); break; case ncclUint64: U8[idx] = ReduceOp(op, U8[idx], inputCpu.U8[idx]); break; + case ncclFp8E4M3: F1[idx] = rccl_float8(ReduceOp(op, float(F1[idx]), float(inputCpu.F1[idx]))); break; case ncclFloat16: F2[idx] = __float2half(ReduceOp(op, __half2float(F2[idx]), __half2float(inputCpu.F2[idx]))); break; case ncclFloat32: F4[idx] = ReduceOp(op, F4[idx], inputCpu.F4[idx]); break; case ncclFloat64: F8[idx] = ReduceOp(op, F8[idx], inputCpu.F8[idx]); break; + case ncclFp8E5M2: B1[idx] = rccl_bfloat8(ReduceOp(op, float(B1[idx]), float(inputCpu.B1[idx]))); break; case ncclBfloat16: B2[idx] = ReduceOp(op, B2[idx], inputCpu.B2[idx]); break; default: ERROR("Unsupported datatype\n"); @@ -271,10 +281,12 @@ namespace RcclUnitTesting case ncclUint32: U4[idx] /= divisor; break; case ncclInt64: I8[idx] /= divisor; break; case ncclUint64: U8[idx] /= divisor; break; - case ncclFloat16: F2[idx] = __float2half(__half2float(F2[idx])/divisor); break; + case ncclFp8E4M3: F1[idx] = (rccl_float8((float)(F1[idx]) / divisor)); break; + case ncclFloat16: F2[idx] = __float2half(__half2float(F2[idx])/divisor); break; case ncclFloat32: F4[idx] /= divisor; break; case ncclFloat64: F8[idx] /= divisor; break; - case ncclBfloat16: B2[idx] = (rccl_bfloat16((float)(B2[idx]) / divisor)); break; + case ncclFp8E5M2: B1[idx] = (rccl_bfloat8((float)(B1[idx]) / divisor)); break; + case ncclBfloat16: B2[idx] = (rccl_bfloat16((float)(B2[idx]) / divisor)); break; default: ERROR("Unsupported datatype\n"); return TEST_FAIL; @@ -295,15 +307,17 @@ namespace RcclUnitTesting { switch (dataType) { - case ncclInt8: isMatch = (I1[idx] == expected.I1[idx]); break; - case ncclUint8: isMatch = (U1[idx] == expected.U1[idx]); break; - case ncclInt32: isMatch = (I4[idx] == expected.I4[idx]); break; - case ncclUint32: isMatch = (U4[idx] == expected.U4[idx]); break; - case ncclInt64: isMatch = (I8[idx] == expected.I8[idx]); break; - case ncclUint64: isMatch = (U8[idx] == expected.U8[idx]); break; - case ncclFloat16: isMatch = (fabs(__half2float(F2[idx]) - __half2float(expected.F2[idx])) < 9e-2); break; - case ncclFloat32: isMatch = (fabs(F4[idx] - expected.F4[idx]) < 1e-5); break; - case ncclFloat64: isMatch = (fabs(F8[idx] - expected.F8[idx]) < 1e-12); break; + case ncclInt8: isMatch = (I1[idx] == expected.I1[idx]); break; + case ncclUint8: isMatch = (U1[idx] == expected.U1[idx]); break; + case ncclInt32: isMatch = (I4[idx] == expected.I4[idx]); break; + case ncclUint32: isMatch = (U4[idx] == expected.U4[idx]); break; + case ncclInt64: isMatch = (I8[idx] == expected.I8[idx]); break; + case ncclUint64: isMatch = (U8[idx] == expected.U8[idx]); break; + case ncclFp8E4M3: isMatch = (fabs(float(F1[idx]) - float(expected.F1[idx])) < 9e-2); break; + case ncclFloat16: isMatch = (fabs(__half2float(F2[idx]) - __half2float(expected.F2[idx])) < 9e-2); break; + case ncclFloat32: isMatch = (fabs(F4[idx] - expected.F4[idx]) < 1e-5); break; + case ncclFloat64: isMatch = (fabs(F8[idx] - expected.F8[idx]) < 1e-12); break; + case ncclFp8E5M2: isMatch = (fabs(float(B1[idx]) - float(expected.B1[idx])) < 9e-2); break; case ncclBfloat16: isMatch = (fabs((float)B2[idx] - (float)expected.B2[idx]) < 9e-2); break; default: ERROR("Unsupported datatype\n"); @@ -328,12 +342,16 @@ namespace RcclUnitTesting ERROR("Expected output: %ld. Actual output: %ld at index %lu\n", expected.I8[idx], I8[idx], idx); break; case ncclUint64: ERROR("Expected output: %lu. Actual output: %lu at index %lu\n", expected.U8[idx], U8[idx], idx); break; + case ncclFp8E4M3: + ERROR("Expected output: %f. Actual output: %f at index %lu\n", (float)expected.F1[idx], (float)F1[idx], idx); break; case ncclFloat16: ERROR("Expected output: %f. Actual output: %f at index %lu\n", __half2float(expected.F2[idx]), __half2float(F2[idx]), idx); break; case ncclFloat32: ERROR("Expected output: %f. Actual output: %f at index %lu\n", expected.F4[idx], F4[idx], idx); break; case ncclFloat64: ERROR("Expected output: %lf. Actual output: %lf at index %lu\n", expected.F8[idx], F8[idx], idx); break; + case ncclFp8E5M2: + ERROR("Expected output: %f. Actual output: %f at index %lu\n", (float)expected.B1[idx], (float)B1[idx], idx); break; case ncclBfloat16: ERROR("Expected output: %f. Actual output: %f at index %lu\n", (float)expected.B2[idx], (float)B2[idx], idx); break; default: @@ -358,9 +376,11 @@ namespace RcclUnitTesting case ncclUint32: ss << U4[i]; break; case ncclInt64: ss << I8[i]; break; case ncclUint64: ss << U8[i]; break; + case ncclFp8E4M3: ss << (float)F1[i]; break; case ncclFloat16: ss << __half2float(F2[i]); break; case ncclFloat32: ss << F4[i]; break; case ncclFloat64: ss << F8[i]; break; + case ncclFp8E5M2: ss << (float)B1[i]; break; case ncclBfloat16: ss << (float)B2[i]; break; default: break; } diff --git a/test/common/PtrUnion.hpp b/test/common/PtrUnion.hpp index d905648942..467ef5a53f 100644 --- a/test/common/PtrUnion.hpp +++ b/test/common/PtrUnion.hpp @@ -7,6 +7,7 @@ #pragma once #include "ErrCode.hpp" #include "rccl/rccl.h" +#include "rccl_float8.h" #include "rccl_bfloat16.h" #include "hip/hip_fp16.h" @@ -43,8 +44,10 @@ namespace RcclUnitTesting int64_t* I8; // ncclInt64 uint64_t* U8; // ncclUint64 __half* F2; // ncclFloat16 + rccl_float8* F1; // ncclFp8E4M3 float* F4; // ncclFloat32 double* F8; // ncclFloat64 + rccl_bfloat8* B1; // ncclFp8E5M2 rccl_bfloat16* B2; // ncclBfloat16 constexpr PtrUnion() : ptr(nullptr) {} diff --git a/tools/ib-test/include/nccl.h b/tools/ib-test/include/nccl.h index 76da780694..2c86c33269 100755 --- a/tools/ib-test/include/nccl.h +++ b/tools/ib-test/include/nccl.h @@ -20,6 +20,7 @@ #define NCCL_VERSION(X,Y,Z) ((X) * 1000 + (Y) * 100 + (Z)) #define RCCL_BFLOAT16 1 +#define RCCL_BFLOAT8 1 #define RCCL_GATHER_SCATTER 1 #ifdef __cplusplus @@ -120,7 +121,9 @@ typedef enum { ncclInt8 = 0, ncclChar = 0, ncclFloat32 = 7, ncclFloat = 7, ncclFloat64 = 8, ncclDouble = 8, ncclBfloat16 = 9, - ncclNumTypes = 10 } ncclDataType_t; + ncclFp8E4M3 = 10, + ncclFp8E5M2 = 11, + ncclNumTypes = 12 } ncclDataType_t; /* * Collective communication operations diff --git a/tools/rccl_replayer/rcclReplayer.hpp b/tools/rccl_replayer/rcclReplayer.hpp index 3d8f1dc2fc..f031f84a8e 100644 --- a/tools/rccl_replayer/rcclReplayer.hpp +++ b/tools/rccl_replayer/rcclReplayer.hpp @@ -130,6 +130,8 @@ size_t DataTypeToBytes(ncclDataType_t const dataType) case ncclFloat32: return 4; case ncclFloat64: return 8; case ncclBfloat16: return 2; + case ncclFp8E4M3: return 1; + case ncclFp8E5M2: return 1; default: printf("Unsupported datatype (%d)\n", dataType); exit(0); diff --git a/tools/topo_expl/include/nccl.h b/tools/topo_expl/include/nccl.h index 27737e2231..729561bd8b 100644 --- a/tools/topo_expl/include/nccl.h +++ b/tools/topo_expl/include/nccl.h @@ -21,6 +21,7 @@ #define NCCL_VERSION(X,Y,Z) (((X) <= 2 && (Y) <= 8) ? (X) * 1000 + (Y) * 100 + (Z) : (X) * 10000 + (Y) * 100 + (Z)) #define RCCL_BFLOAT16 1 +#define RCCL_BFLOAT8 1 #define RCCL_GATHER_SCATTER 1 #define RCCL_ALLTOALLV 1 @@ -253,7 +254,9 @@ typedef enum { ncclInt8 = 0, ncclChar = 0, ncclFloat32 = 7, ncclFloat = 7, ncclFloat64 = 8, ncclDouble = 8, ncclBfloat16 = 9, - ncclNumTypes = 10 } ncclDataType_t; + ncclFp8E4M3 = 10, + ncclFp8E5M2 = 11, + ncclNumTypes = 12 } ncclDataType_t; /*! @brief ncclScalarResidence_t: Location and dereferencing logic for scalar arguments. */ typedef enum {