Implement long / long long shuffles (#1829)

Implement additional data-types for shuffles (long and long long).
Based upon the double implementation.

hipamd/include/hip/hcc_detail/device_functions.h

@@ -319,6 +319,36 @@ double __shfl(double var, int src_lane, int width = warpSize) {
     double tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
     return tmp1;
 }
+__device__
+inline
+long __shfl(long var, int src_lane, int width = warpSize)
+{
+    static_assert(sizeof(long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long) == sizeof(uint64_t), "");
+
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl(tmp[0], src_lane, width);
+    tmp[1] = __shfl(tmp[1], src_lane, width);
+
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
+__device__
+inline
+long long __shfl(long long var, int src_lane, int width = warpSize)
+{
+    static_assert(sizeof(long long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long long) == sizeof(uint64_t), "");
+
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl(tmp[0], src_lane, width);
+    tmp[1] = __shfl(tmp[1], src_lane, width);
+
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
 
  __device__
 inline
@@ -356,6 +386,34 @@ double __shfl_up(double var, unsigned int lane_delta, int width = warpSize) {
     double tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
     return tmp1;
 }
+__device__
+inline
+long __shfl_up(long var, unsigned int lane_delta, int width = warpSize)
+{
+    static_assert(sizeof(long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long) == sizeof(uint64_t), "");
+
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_up(tmp[0], lane_delta, width);
+    tmp[1] = __shfl_up(tmp[1], lane_delta, width);
+
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
+__device__
+inline
+long long __shfl_up(long long var, unsigned int lane_delta, int width = warpSize)
+{
+    static_assert(sizeof(long long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long long) == sizeof(uint64_t), "");
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_up(tmp[0], lane_delta, width);
+    tmp[1] = __shfl_up(tmp[1], lane_delta, width);
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
 
 __device__
 inline
@@ -393,6 +451,34 @@ double __shfl_down(double var, unsigned int lane_delta, int width = warpSize) {
     double tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
     return tmp1;
 }
+__device__
+inline
+long __shfl_down(long var, unsigned int lane_delta, int width = warpSize)
+{
+    static_assert(sizeof(long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long) == sizeof(uint64_t), "");
+
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_down(tmp[0], lane_delta, width);
+    tmp[1] = __shfl_down(tmp[1], lane_delta, width);
+
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
+__device__
+inline
+long long __shfl_down(long long var, unsigned int lane_delta, int width = warpSize)
+{
+    static_assert(sizeof(long long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long long) == sizeof(uint64_t), "");
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_down(tmp[0], lane_delta, width);
+    tmp[1] = __shfl_down(tmp[1], lane_delta, width);
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
 
 __device__
 inline
@@ -430,6 +516,34 @@ double __shfl_xor(double var, int lane_mask, int width = warpSize) {
     double tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
     return tmp1;
 }
+__device__
+inline
+long __shfl_xor(long var, int lane_mask, int width = warpSize)
+{
+    static_assert(sizeof(long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long) == sizeof(uint64_t), "");
+
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_xor(tmp[0], lane_mask, width);
+    tmp[1] = __shfl_xor(tmp[1], lane_mask, width);
+
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
+__device__
+inline
+long long __shfl_xor(long long var, int lane_mask, int width = warpSize)
+{
+    static_assert(sizeof(long long) == 2 * sizeof(int), "");
+    static_assert(sizeof(long long) == sizeof(uint64_t), "");
+    int tmp[2]; __builtin_memcpy(tmp, &var, sizeof(tmp));
+    tmp[0] = __shfl_xor(tmp[0], lane_mask, width);
+    tmp[1] = __shfl_xor(tmp[1], lane_mask, width);
+    uint64_t tmp0 = (static_cast<uint64_t>(tmp[1]) << 32ull) | static_cast<uint32_t>(tmp[0]);
+    long long tmp1;  __builtin_memcpy(&tmp1, &tmp0, sizeof(tmp0));
+    return tmp1;
+}
 
 #define MASK1 0x00ff00ff
 #define MASK2 0xff00ff00

hipamd/tests/src/kernel/hipShflTests.cpp

@@ -0,0 +1,130 @@
+/*
+Copyright (c) 2015-present Advanced Micro Devices, Inc. All rights reserved.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+*/
+/* HIT_START
+ * BUILD: %t %s ../test_common.cpp
+ * TEST: %t
+ * HIT_END
+ */
+
+#include <iostream>
+#include <hip/hip_runtime.h>
+#include "test_common.h"
+
+#define WIDTH 4
+
+#define NUM (WIDTH * WIDTH)
+
+#define THREADS_PER_BLOCK_X 4
+#define THREADS_PER_BLOCK_Y 4
+#define THREADS_PER_BLOCK_Z 1
+
+// Device (Kernel) function, it must be void
+template <typename T>
+__global__ void matrixTranspose(T* out, T* in, const int width) {
+    int x = hipBlockDim_x * hipBlockIdx_x + hipThreadIdx_x;
+    T val = in[x];
+    for (int i = 0; i < width; i++) {
+        for (int j = 0; j < width; j++) out[i * width + j] = __shfl(val, j * width + i);
+    }
+}
+
+// CPU implementation of matrix transpose
+template <typename T>
+void matrixTransposeCPUReference(T* output, T* input, const unsigned int width) {
+    for (unsigned int j = 0; j < width; j++) {
+        for (unsigned int i = 0; i < width; i++) {
+            output[i * width + j] = input[j * width + i];
+        }
+    }
+}
+
+template<typename T>
+void runTest() {
+    T* Matrix;
+    T* TransposeMatrix;
+    T* cpuTransposeMatrix;
+
+    T* gpuMatrix;
+    T* gpuTransposeMatrix;
+
+    hipDeviceProp_t devProp;
+    hipGetDeviceProperties(&devProp, 0);
+
+    int i;
+    int errors;
+
+    Matrix = (T*)malloc(NUM * sizeof(T));
+    TransposeMatrix = (T*)malloc(NUM * sizeof(T));
+    cpuTransposeMatrix = (T*)malloc(NUM * sizeof(T));
+
+    // initialize the input data
+    for (i = 0; i < NUM; i++) {
+        Matrix[i] = (T)i * 10l;
+    }
+
+    // allocate the memory on the device side
+    hipMalloc((void**)&gpuMatrix, NUM * sizeof(T));
+    hipMalloc((void**)&gpuTransposeMatrix, NUM * sizeof(T));
+
+    // Memory transfer from host to device
+    hipMemcpy(gpuMatrix, Matrix, NUM * sizeof(T), hipMemcpyHostToDevice);
+
+    // Lauching kernel from host
+    hipLaunchKernelGGL(matrixTranspose<T>, dim3(1), dim3(THREADS_PER_BLOCK_X * THREADS_PER_BLOCK_Y), 0, 0,
+                    gpuTransposeMatrix, gpuMatrix, WIDTH);
+
+    // Memory transfer from device to host
+    hipMemcpy(TransposeMatrix, gpuTransposeMatrix, NUM * sizeof(T), hipMemcpyDeviceToHost);
+
+    // CPU MatrixTranspose computation
+    matrixTransposeCPUReference(cpuTransposeMatrix, Matrix, WIDTH);
+
+    // verify the results
+    errors = 0;
+    double eps = 1.0E-6;
+    for (i = 0; i < NUM; i++) {
+        if (TransposeMatrix[i] != cpuTransposeMatrix[i]) {
+            errors++;
+        }
+    }
+
+    // free the resources on device side
+    hipFree(gpuMatrix);
+    hipFree(gpuTransposeMatrix);
+
+    // free the resources on host side
+    free(Matrix);
+    free(TransposeMatrix);
+    free(cpuTransposeMatrix);
+
+    if (errors != 0) {
+        failed("Mismatch present");
+    }
+}
+
+int main() {
+    runTest<int>();
+    runTest<float>();
+    runTest<long>();
+    runTest<long long>();
+    passed();
+}