From 609598cd152415bc0f580b6c93e8cf4d104086d1 Mon Sep 17 00:00:00 2001
From: Anusha GodavarthySurya <anusha.godavarthysurya@amd.com>
Date: Tue, 22 Nov 2022 11:51:56 +0000
Subject: [PATCH] SWDEV-345571 - Added support for half/half2 data types for
 warp shuffle functions

Change-Id: Ia0791760c1b0a35bbbf0b1f8435aa924af205169


[ROCm/clr commit: ad1fae7df4f029a31647f40712d3e600a9250c04]
---
 .../hip/amd_detail/amd_warp_functions.h       | 60 +++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h
index ebdc86374d..b671b3d0ca 100644
--- a/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h
+++ b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h
@@ -23,6 +23,10 @@ THE SOFTWARE.
 #ifndef HIP_INCLUDE_HIP_AMD_DETAIL_WARP_FUNCTIONS_H
 #define HIP_INCLUDE_HIP_AMD_DETAIL_WARP_FUNCTIONS_H
 
+#if !defined(__HIPCC_RTC__)
+#include "hip/hip_fp16.h"
+#endif
+
 static constexpr int warpSize = __AMDGCN_WAVEFRONT_SIZE;
 
 __device__
@@ -48,6 +52,20 @@ float __shfl(float var, int src_lane, int width = warpSize) {
 }
 __device__
 inline
+__half __shfl(__half var, int src_lane, int width = warpSize) {
+    union { int i; __half h; } tmp; tmp.h = var;
+    tmp.i = __shfl(tmp.i, src_lane, width);
+    return tmp.h;
+}
+__device__
+inline
+__half2 __shfl(__half2 var, int src_lane, int width = warpSize) {
+    union { int i; __half2 h; } tmp; tmp.h = var;
+    tmp.i = __shfl(tmp.i, src_lane, width);
+    return tmp.h;
+}
+__device__
+inline
 double __shfl(double var, int src_lane, int width = warpSize) {
     static_assert(sizeof(double) == 2 * sizeof(int), "");
     static_assert(sizeof(double) == sizeof(uint64_t), "");
@@ -153,6 +171,20 @@ float __shfl_up(float var, unsigned int lane_delta, int width = warpSize) {
 }
 __device__
 inline
+__half __shfl_up(__half var, unsigned int lane_delta, int width = warpSize) {
+    union { int i; __half h; } tmp; tmp.h = var;
+    tmp.i = __shfl_up(tmp.i, lane_delta, width);
+    return tmp.h;
+}
+__device__
+inline
+__half2 __shfl_up(__half2 var, unsigned int lane_delta, int width = warpSize) {
+    union { int i; __half2 h; } tmp; tmp.h = var;
+    tmp.i = __shfl_up(tmp.i, lane_delta, width);
+    return tmp.h;
+}
+__device__
+inline
 double __shfl_up(double var, unsigned int lane_delta, int width = warpSize) {
     static_assert(sizeof(double) == 2 * sizeof(int), "");
     static_assert(sizeof(double) == sizeof(uint64_t), "");
@@ -259,6 +291,20 @@ float __shfl_down(float var, unsigned int lane_delta, int width = warpSize) {
 }
 __device__
 inline
+__half __shfl_down(__half var, unsigned int lane_delta, int width = warpSize) {
+    union { int i; __half h; } tmp; tmp.h = var;
+    tmp.i = __shfl_down(tmp.i, lane_delta, width);
+    return tmp.h;
+}
+__device__
+inline
+__half2 __shfl_down(__half2 var, unsigned int lane_delta, int width = warpSize) {
+    union { int i; __half2 h; } tmp; tmp.h = var;
+    tmp.i = __shfl_down(tmp.i, lane_delta, width);
+    return tmp.h;
+}
+__device__
+inline
 double __shfl_down(double var, unsigned int lane_delta, int width = warpSize) {
     static_assert(sizeof(double) == 2 * sizeof(int), "");
     static_assert(sizeof(double) == sizeof(uint64_t), "");
@@ -362,6 +408,20 @@ float __shfl_xor(float var, int lane_mask, int width = warpSize) {
 }
 __device__
 inline
+__half __shfl_xor(__half var,  int lane_mask, int width = warpSize) {
+    union { int i; __half h; } tmp; tmp.h = var;
+    tmp.i = __shfl_xor(tmp.i, lane_mask, width);
+    return tmp.h;
+}
+__device__
+inline
+__half2 __shfl_xor(__half2 var,  int lane_mask, int width = warpSize) {
+    union { int i; __half2 h; } tmp; tmp.h = var;
+    tmp.i = __shfl_xor(tmp.i, lane_mask, width);
+    return tmp.h;
+}
+__device__
+inline
 double __shfl_xor(double var, int lane_mask, int width = warpSize) {
     static_assert(sizeof(double) == 2 * sizeof(int), "");
     static_assert(sizeof(double) == sizeof(uint64_t), "");