opt

Author: BonsonW

src/nn_cuda.cu

@@ -60,9 +60,9 @@ void silu_mul_cuda(
     uint64_t K
 ) {
     auto threads = 1024;
-    auto blocks = ((K + threads - 1) / threads) * MN;
+    auto blocks = MN;
 
-    silu_mul<<<blocks, threads >>>(
+    silu_mul<<<blocks, threads>>>(
         (half *)x_gpu,
         (half *)o_gpu,
         K,

src/nn_hip.hip

@@ -65,9 +65,9 @@ void silu_mul_hip(
     hipError_t ret;
 
     auto threads = 1024;
-    auto blocks = ((K + threads - 1) / threads) * MN;
+    auto blocks = MN;
 
-    silu_mul<<<blocks, threads >>>(
+    silu_mul<<<blocks, threads>>>(
         (half *)x_gpu,
         (half *)o_gpu,
         K,

src/nn_kernel_cuda.h

@@ -74,20 +74,18 @@ __global__ void silu_mul(
     const int K,
     const int MN
 ) {
-    int i = blockIdx.x * blockDim.x + threadIdx.x;
-    int j = i / K;
-    int k = i - (j * K);
+    int j = blockIdx.x;
 
-    if (k < K && j < MN) {
-        k += j * (K * 2);
+    for (int k = threadIdx.x; k < K; k += blockDim.x) {
+        int i = k + j * (K * 2);
 
-        half y = x_gpu[k];
-        half gate = x_gpu[k + K];
+        half y = x_gpu[i];
+        half gate = x_gpu[i + K];
 
         float g = __half2float(gate);
         float silu = g / (1.0f + __expf(-g));
 
-        o_gpu[i] = __float2half(silu * __half2float(y));
+        o_gpu[k + j * K] = __float2half(silu * __half2float(y));
     }
 }
 

src/nn_kernel_hip.h

@@ -75,20 +75,18 @@ __global__ void silu_mul(
     const int K,
     const int MN
 ) {
-    int i = blockIdx.x * blockDim.x + threadIdx.x;
-    int j = i / K;
-    int k = i - (j * K);
+    int j = blockIdx.x;
 
-    if (k < K && j < MN) {
-        k += j * (K * 2);
+    for (int k = threadIdx.x; k < K; k += blockDim.x) {
+        int i = k + j * (K * 2);
 
-        half y = x_gpu[k];
-        half gate = x_gpu[k + K];
+        half y = x_gpu[i];
+        half gate = x_gpu[i + K];
 
         float g = __half2float(gate);
         float silu = g / (1.0f + __expf(-g));
 
-        o_gpu[i] = __float2half(silu * __half2float(y));
+        o_gpu[k + j * K] = __float2half(silu * __half2float(y));
     }
 }