fix(nv): 改正 rearrange

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

src/ops/rearrange/cuda/rearrange.cc

@@ -7,7 +7,8 @@ infiniopStatus_t cudaCreateRearrangeDescriptor(CudaHandle_t handle,
                                                RearrangeCudaDescriptor_t *desc_ptr,
                                                infiniopTensorDescriptor_t dst,
                                                infiniopTensorDescriptor_t src) {
-    if (!dtype_eq(dst->dt, src->dt)) {
+    auto dt = dst->dt;
+    if (!dtype_eq(src->dt, dt)) {
         return STATUS_BAD_TENSOR_DTYPE;
     }
 
@@ -24,62 +25,43 @@ infiniopStatus_t cudaCreateRearrangeDescriptor(CudaHandle_t handle,
         return STATUS_BAD_TENSOR_STRIDES;
     }
 
-    if (ndim == 1) {
-        *desc_ptr = new RearrangeCudaDescriptor{
-            handle->device,
-            handle->device_id,
-            0, 0, 0, 0,
-            1, 1, 1,
-            static_cast<unsigned long>(dst->shape[0] * dst->dt.size)};
-        return STATUS_SUCCESS;
+    switch (ndim) {
+        case 1:
+            *desc_ptr = new RearrangeCudaDescriptor{
+                handle->device,
+                handle->device_id,
+                dt.size * dst->shape[0],
+                1, 1,
+                0, 0,
+                0, 0};
+            break;
+        case 2:
+            *desc_ptr = new RearrangeCudaDescriptor{
+                handle->device,
+                handle->device_id,
+                dt.size * dst->shape[1],
+                1, dst->shape[0],
+                0, dst->strides[0],
+                0, src->strides[0]};
+            break;
+        case 3:
+            *desc_ptr = new RearrangeCudaDescriptor{
+                handle->device,
+                handle->device_id,
+                dt.size * dst->shape[2],
+                dst->shape[0], dst->shape[1],
+                dst->strides[0], dst->strides[1],
+                src->strides[0], src->strides[1]};
+            break;
+        default:
+            return STATUS_BAD_TENSOR_SHAPE;
     }
 
-    unsigned int r = 0, c = 0, b = 0;
-    unsigned int rsa = 0, csa = 0, rsb = 0, csb = 0;
-    if (ndim == 2) {
-        c = dst->shape[0];
-        b = dst->shape[1];
-        csa = dst->strides[0];
-        csb = src->strides[0];
-    } else if (ndim == 3) {
-        r = dst->shape[0];
-        c = dst->shape[1];
-        b = dst->shape[2];
-        csa = dst->strides[1];
-        csb = src->strides[1];
-        rsa = dst->strides[0];
-        rsb = src->strides[0];
-    } else {
-        for (int i = ndim - 3; i >= 1; --i) {
-            if (dst->shape[i] * dst->strides[i] != dst->strides[i - 1] || src->shape[i] * src->strides[i] != src->strides[i - 1]) {
-                return STATUS_BAD_TENSOR_STRIDES;
-            }
-        }
-        r = std::accumulate(dst->shape, dst->shape + ndim - 2, 1, std::multiplies<unsigned int>());
-        c = dst->shape[ndim - 2];
-        b = dst->shape[ndim - 1];
-        csa = dst->strides[ndim - 2];
-        csb = src->strides[ndim - 2];
-        rsa = dst->strides[ndim - 3];
-        rsb = src->strides[ndim - 3];
-    }
-    auto contiguous_bytes = b * dst->dt.size;
-    if (contiguous_bytes % WARP_SIZE != 0) {
-        return STATUS_BAD_PARAM;
-    }
-    auto bytes_per_thread = contiguous_bytes / WARP_SIZE;
-    if (bytes_per_thread <= 0 || bytes_per_thread > 32 || (bytes_per_thread & (bytes_per_thread - 1)) != 0) {
-        return STATUS_BAD_PARAM;
-    }
-    *desc_ptr = new RearrangeCudaDescriptor{
-        handle->device,
-        handle->device_id,
-        rsa,
-        rsb,
-        csa,
-        csb,
-        r, c, b,
-        bytes_per_thread};
+    (*desc_ptr)->dst_rs *= dt.size;
+    (*desc_ptr)->dst_cs *= dt.size;
+    (*desc_ptr)->src_rs *= dt.size;
+    (*desc_ptr)->src_cs *= dt.size;
+
     return STATUS_SUCCESS;
 }
 infiniopStatus_t cudaDestroyRearrangeDescriptor(RearrangeCudaDescriptor_t desc) {

src/ops/rearrange/cuda/rearrange.cu

@@ -4,11 +4,11 @@
 template<class Tmem>
 static __global__ void rearrange(
     void *__restrict__ dst,
-    unsigned int const rsa,
-    unsigned int const csa,
+    int const rsa,
+    int const csa,
     void const *__restrict__ src,
-    unsigned int const rsb,
-    unsigned int const csb,
+    int const rsb,
+    int const csb,
     unsigned int const ncols) {
 
     auto row = blockIdx.y,
@@ -25,39 +25,42 @@ static __global__ void rearrange(
 
 void rearrange_nv_gpu(RearrangeCudaDescriptor_t desc, void *y, void const *x, void *stream) {
     auto cuda_stream = reinterpret_cast<cudaStream_t>(stream);
-    if (desc->r == 1 && desc->c == 1 && desc->b == 1) {
-        cudaMemcpyAsync(y, x, desc->bytes_per_thread, cudaMemcpyDeviceToDevice, cuda_stream);
+    auto unit = desc->unit,
+         r = desc->r, c = desc->c;
+    auto dst_rs = desc->dst_rs, dst_cs = desc->dst_cs,
+         src_rs = desc->src_rs, src_cs = desc->src_cs;
+
+    if (r == 1 && c == 1) {
+        cudaMemcpyAsync(y, x, unit, cudaMemcpyDeviceToDevice, cuda_stream);
         return;
     }
 
-    uint64_t rsa = desc->rsa, csa = desc->csa, rsb = desc->rsb, csb = desc->csb;
-    unsigned int r = desc->r, c = desc->c, b = desc->b, bytes_per_thread = desc->bytes_per_thread;
-    auto dst_ptr = static_cast<void *>(reinterpret_cast<uint8_t *>(y));
-    rsa /= b;
-    csa /= b;
-    auto src_ptr = static_cast<void const *>(reinterpret_cast<uint8_t const *>(x));
-    rsb /= b;
-    csb /= b;
-    dim3 grid_dims = dim3((c + MAX_WARP_PER_BLOCK - 1) / MAX_WARP_PER_BLOCK, r);
-    dim3 block_dims = dim3(WARP_SIZE, (c + grid_dims.x - 1) / grid_dims.x);
-    switch (bytes_per_thread) {
+    auto warps = 1024 / WARP_SIZE;
+    auto grid = dim3((c + warps - 1) / warps, r);
+    auto block = dim3(WARP_SIZE, (c + grid.x - 1) / grid.x);
+    dst_rs /= unit;
+    dst_cs /= unit;
+    src_rs /= unit;
+    src_cs /= unit;
+
+    switch (unit / WARP_SIZE) {
         case 1:
-            rearrange<uchar1><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<uchar1><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         case 2:
-            rearrange<uchar2><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<uchar2><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         case 4:
-            rearrange<float1><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<float1><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         case 8:
-            rearrange<float2><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<float2><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         case 16:
-            rearrange<float4><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<float4><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         case 32:
-            rearrange<double4><<<grid_dims, block_dims, 0, cuda_stream>>>(dst_ptr, rsa, csa, src_ptr, rsb, csb, c);
+            rearrange<double4><<<grid, block, 0, cuda_stream>>>(y, dst_rs, dst_cs, x, src_rs, src_cs, c);
             break;
         default:
             break;

src/ops/rearrange/cuda/rearrange.cuh

@@ -7,12 +7,8 @@
 struct RearrangeCudaDescriptor {
     Device device;
     int device_id;
-    uint64_t rsa;
-    uint64_t rsb;
-    uint64_t csa;
-    uint64_t csb;
-    uint64_t r, c, b;
-    uint64_t bytes_per_thread;
+    uint64_t unit, r, c;
+    int64_t dst_rs, dst_cs, src_rs, src_cs;
 };
 
 typedef struct RearrangeCudaDescriptor *RearrangeCudaDescriptor_t;