AI/rocm-systems
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SWDEV-345571 - Added support for half/half2 data types for warp shuffle functions

Browse Source 
Change-Id: Ia0791760c1b0a35bbbf0b1f8435aa924af205169


[ROCm/clr commit: ad1fae7df4]
This commit is contained in:
Anusha GodavarthySurya
2022-11-22 11:51:56 +00:00
committed by Anusha Godavarthy Surya
parent 49a3b0c6ba
commit 609598cd15
1 changed files with 60 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h
+60
View File
@@ -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), "");