From 9769fa68a74d78c5b34e7638271d882d42a93020 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ville=20Pietil=C3=A4?= Date: Fri, 1 Aug 2025 13:22:41 +0000 Subject: [PATCH] Packed BF16 cast with asm volatile. --- .../element/unary_element_wise_operation.hpp | 5 ++ include/ck/utility/type_convert.hpp | 55 ++++++++--------- test/data_type/test_bhalf.cpp | 60 +++++++++++++++++++ test_hip_devices.cpp | 23 +++++++ 4 files changed, 114 insertions(+), 29 deletions(-) create mode 100644 test_hip_devices.cpp diff --git a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp index 4a87e8a277..0cf81038f8 100644 --- a/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp +++ b/include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp @@ -278,6 +278,11 @@ struct PassThroughPack2 template __host__ __device__ void operator()(Y& y, const X& x) const; + __host__ __device__ constexpr void operator()(ck::bhalf2_t& y, const ck::float2_t& x) const + { + y = bf16x2_convert_rne(x); + } + __host__ __device__ constexpr void operator()(half2_t& y, const f8x2_t& x) const { auto t = type_convert(x); diff --git a/include/ck/utility/type_convert.hpp b/include/ck/utility/type_convert.hpp index 13900a1575..4edbd21395 100644 --- a/include/ck/utility/type_convert.hpp +++ b/include/ck/utility/type_convert.hpp @@ -55,14 +55,11 @@ inline __device__ bhalf_t static_cast_float_to_bf16(float x) #if defined(__gfx950__) inline __device__ bhalf2_t static_cast_float2_to_bhalf2_rne(float2_t x) { - union - { - bhalf2_t bf16x2val; - uint32_t i32val; - } value; - // Use the intrinsic to convert float2 to bfloat2 with RNE - value.i32val = __builtin_amdgcn_cvt_pk_bf16_f32(x[Number<0>{}], x[Number<1>{}]); - return value.bf16x2val; + uint32_t result; + asm volatile("v_cvt_pk_bf16_f32 %0, %1, %2" + : "=v"(result) + : "v"(x[0]), "v"(x[1])); + return result; } #endif @@ -70,27 +67,6 @@ inline __device__ bhalf2_t static_cast_float2_to_bhalf2_rne(float2_t x) template __host__ __device__ constexpr Y bf16x2_convert_rne(X x); -/** - * @brief Converts a vector of 2 float (float2_t) to a vector of 2 16-bit bfloat types (bhalf2_t) using - * rounding to nearest/even (RNE). - * - * @param x The input vector of 2 float. - * @return The converted vector of 2 bhalf_t. - */ -template<> -inline __host__ __device__ constexpr bhalf2_t bf16x2_convert_rne(float2_t x) -{ -#if defined(__gfx950__) - return static_cast_float2_to_bhalf2_rne(x); -#else - // TODO: Perform real RNE conversion for bfloat16 - return { - bf16_convert_rtn(x[Number<0>{}]), - bf16_convert_rtn(x[Number<1>{}]) - }; -#endif -} - // Declare a template function for bf16 conversion using RTN template __host__ __device__ constexpr Y bf16_convert_rtn(X x); @@ -121,6 +97,27 @@ inline __host__ __device__ constexpr bhalf_t bf16_convert_rtn(fl #endif } +/** + * @brief Converts a vector of 2 float (float2_t) to a vector of 2 16-bit bfloat types (bhalf2_t) using + * rounding to nearest/even (RNE). + * + * @param x The input vector of 2 float. + * @return The converted vector of 2 bhalf_t. + */ +template<> +inline __host__ __device__ constexpr bhalf2_t bf16x2_convert_rne(float2_t x) +{ +#if defined(__gfx950__) + return static_cast_float2_to_bhalf2_rne(x); +#else + // TODO: Perform real RNE conversion for bfloat16 + return { + bf16_convert_rtn(x[0]), + bf16_convert_rtn(x[1]) + }; +#endif +} + // convert fp16 to bfp16 via fp32 with RTN if higher precision is needed template <> inline __host__ __device__ constexpr bhalf_t bf16_convert_rtn(half_t x) diff --git a/test/data_type/test_bhalf.cpp b/test/data_type/test_bhalf.cpp index ad31e194b8..ff0e62a43c 100644 --- a/test/data_type/test_bhalf.cpp +++ b/test/data_type/test_bhalf.cpp @@ -51,12 +51,72 @@ TEST(BHALF_T, MantisaExpOverflow) ASSERT_TRUE(std::isnan(type_convert(type_convert(float_val)))); } +__global__ void cast2(const float2 input, float2* output) +{ + ck::float2_t x {input.x, input.y}; + const ck::bhalf2_t bhalf2_val = ck::bf16x2_convert_rne(x); + const float fval1 = type_convert(bhalf2_val[0]); + const float fval2 = type_convert(bhalf2_val[1]); + output->x = fval1; + output->y = fval2; +} + __global__ void cast(const float input, float* output) { const bhalf_t bhalf_val = type_convert(input); *output = type_convert(bhalf_val); } +// Since the type_convert is almost at the bottom of the dependency tree, +// any changes there require rebuilding a large number of components. +// To build only this test, run 'make test_bhalf' to test the BF16 packed cast. +TEST(BHALF_T, PackedCastRoundtrip) +{ + constexpr int num_vals = 11; + const float abs_tol = std::pow(2, -7); + float float_vals[num_vals] = {0.5, 0.875, 1.5, 1, 2, 4, 8, 16, 32, 64, 128}; + + float2* float_val_after_cast_dev; + float2 float_val_after_cast_host; + hip_check_error(hipMalloc(&float_val_after_cast_dev, sizeof(float2))); + + // Positive + for(int i = 0; i < num_vals; i++) + { + for (int j = 0; j < num_vals; j++) + { + cast2<<<1, 1>>>(float2{float_vals[i], float_vals[j]}, float_val_after_cast_dev); + + hip_check_error(hipMemcpy(&float_val_after_cast_host, + float_val_after_cast_dev, + sizeof(float2), + hipMemcpyDeviceToHost)); + + ASSERT_NEAR(float_val_after_cast_host.x, float_vals[i], abs_tol); + ASSERT_NEAR(float_val_after_cast_host.y, float_vals[j], abs_tol); + } + } + + // Negative + for(int i = 0; i < num_vals; i++) + { + for (int j = 0; j < num_vals; j++) + { + cast2<<<1, 1>>>(float2{-float_vals[i], -float_vals[j]}, float_val_after_cast_dev); + + hip_check_error(hipMemcpy(&float_val_after_cast_host, + float_val_after_cast_dev, + sizeof(float2), + hipMemcpyDeviceToHost)); + + ASSERT_NEAR(float_val_after_cast_host.x, -float_vals[i], abs_tol); + ASSERT_NEAR(float_val_after_cast_host.y, -float_vals[j], abs_tol); + } + } + + hip_check_error(hipFree(float_val_after_cast_dev)); +} + TEST(BHALF_T, CastOnDevice) { constexpr int num_vals = 11; diff --git a/test_hip_devices.cpp b/test_hip_devices.cpp new file mode 100644 index 0000000000..18b0af7950 --- /dev/null +++ b/test_hip_devices.cpp @@ -0,0 +1,23 @@ +#include +#include + +int main() { + int deviceCount = 0; + hipError_t error = hipGetDeviceCount(&deviceCount); + + if (error != hipSuccess) { + std::cout << "HIP Error: " << hipGetErrorString(error) << std::endl; + return 1; + } + + std::cout << "Number of HIP devices: " << deviceCount << std::endl; + + for (int i = 0; i < deviceCount; i++) { + hipDeviceProp_t prop; + hipGetDeviceProperties(&prop, i); + std::cout << "Device " << i << ": " << prop.name + << " (Compute Capability: " << prop.major << "." << prop.minor << ")" << std::endl; + } + + return 0; +}