From d650f7a50c3b330526c058aed0fab78c66aae10d Mon Sep 17 00:00:00 2001
From: Anusha GodavarthySurya <anusha.godavarthysurya@amd.com>
Date: Wed, 25 Jan 2023 08:56:32 +0000
Subject: [PATCH] SWDEV-345571 - Added support for half/half2 data types for
 warp shuffle functions

Change-Id: I276a55129d6527b73292d75eb18403acd2623c8a


[ROCm/clr commit: 32a9249d3edf83916a114b697932570598c5e1a9]
---
 .../include/hip/amd_detail/amd_hip_fp16.h     | 60 +++++++++++++++++++
 1 file changed, 60 insertions(+)

diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h
index ab6496d6d1..fa58c3e45d 100644
--- a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h
+++ b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_fp16.h
@@ -1710,6 +1710,66 @@ THE SOFTWARE.
             using half = __half;
             using half2 = __half2;
         #endif
+        #if !defined(__HIPCC_RTC__)
+        #include "amd_device_functions.h"
+        #include "amd_warp_functions.h"
+            __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
+            __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
+            __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
+            __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;
+            }
+        #endif // !defined(__HIPCC_RTC__)
     #endif // defined(__cplusplus)
 #elif defined(__GNUC__)
     #include "hip_fp16_gcc.h"