Packed BF16 cast with asm volatile.

include/ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp

@@ -278,6 +278,11 @@ struct PassThroughPack2
     template <typename Y, typename X>
     __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<ck::bhalf2_t, ck::float2_t>(x);
+    }
+
     __host__ __device__ constexpr void operator()(half2_t& y, const f8x2_t& x) const
     {
         auto t = type_convert<float2_t>(x);

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 <typename Y, typename X>
 __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<bhalf2_t, float2_t>(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<bhalf_t>(x[Number<0>{}]),
-        bf16_convert_rtn<bhalf_t>(x[Number<1>{}])
-    };
-#endif
-}
-
 // Declare a template function for bf16 conversion using RTN
 template <typename Y, typename X>
 __host__ __device__ constexpr Y bf16_convert_rtn(X x);
@@ -121,6 +97,27 @@ inline __host__ __device__ constexpr bhalf_t bf16_convert_rtn<bhalf_t, float>(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<bhalf2_t, float2_t>(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<bhalf_t>(x[0]),
+        bf16_convert_rtn<bhalf_t>(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<bhalf_t, half_t>(half_t x)

test/data_type/test_bhalf.cpp

@@ -51,12 +51,72 @@ TEST(BHALF_T, MantisaExpOverflow)
     ASSERT_TRUE(std::isnan(type_convert<float>(type_convert<bhalf_t>(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<ck::bhalf2_t, ck::float2_t>(x);
+    const float fval1 = type_convert<float>(bhalf2_val[0]);
+    const float fval2 = type_convert<float>(bhalf2_val[1]);
+    output->x = fval1;
+    output->y = fval2;
+}
+
 __global__ void cast(const float input, float* output)
 {
     const bhalf_t bhalf_val = type_convert<bhalf_t>(input);
     *output                 = type_convert<float>(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;

test_hip_devices.cpp

@@ -0,0 +1,23 @@
+#include <hip/hip_runtime.h>
+#include <iostream>
+
+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;
+}