#define BLOCK 256

// Sum weights over out_features dimension to get a vector of size in_features
__global__ void weight_col_sum_kernel(const float* __restrict__ W,
                                      int out_features,
                                      int in_features,
                                      float* __restrict__ wsum) {
    int col = blockIdx.x;
    if (col >= in_features) return;

    int tid = threadIdx.x;
    double acc = 0.0;

    // W is [out_features, in_features], row-major
    for (int j = tid; j < out_features; j += blockDim.x) {
        acc += (double)W[j * in_features + col];
    }

    __shared__ double shm[BLOCK];
    shm[tid] = acc;
    __syncthreads();

    // Block reduction
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shm[tid] += shm[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        wsum[col] = (float)shm[0];
    }
}

// Compute sum_j (bias_j - subtract_j) into a single scalar
__global__ void bias_sub_sum_kernel(const float* __restrict__ bias,
                                    const float* __restrict__ subtract,
                                    int has_bias,
                                    int out_features,
                                    float* __restrict__ bs_sum) {
    int tid = threadIdx.x;
    double acc = 0.0;

    for (int j = tid; j < out_features; j += blockDim.x) {
        float b = has_bias ? bias[j] : 0.0f;
        acc += (double)b - (double)subtract[j];
    }

    __shared__ double shm[BLOCK];
    shm[tid] = acc;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shm[tid] += shm[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        bs_sum[0] = (float)shm[0];
    }
}

// Fused output kernel: computes per-batch-row scalar = GELU(((X * sum_j W_j) + sum_j(bias - subtract)) / out_features)
// Then adds this scalar to every element of the corresponding row of X (residual add)
__global__ void fused_output_kernel(const float* __restrict__ X,
                                    const float* __restrict__ wsum,
                                    const float* __restrict__ bs_sum_ptr,
                                    int B,
                                    int in_features,
                                    int out_features,
                                    float* __restrict__ Out) {
    int b = blockIdx.x;
    if (b >= B) return;

    int tid = threadIdx.x;

    // Dot product of X[b, :] with wsum[:]
    double acc = 0.0;
    for (int i = tid; i < in_features; i += blockDim.x) {
        acc += (double)X[b * in_features + i] * (double)wsum[i];
    }

    __shared__ double shm[BLOCK];
    __shared__ float gelu_scalar;
    shm[tid] = acc;
    __syncthreads();

    // Reduce within block
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) {
            shm[tid] += shm[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        float bs_sum = bs_sum_ptr[0];
        double mean = (shm[0] + (double)bs_sum) / (double)out_features;
        float m = (float)mean;
        // Exact GELU using erf, as in PyTorch default
        float v = 0.5f * m * (1.0f + erff(m * 0.70710678118654752440f));
        gelu_scalar = v;
    }
    __syncthreads();

    for (int i = tid; i < in_features; i += blockDim.x) {
        Out[b * in_features + i] = X[b * in_features + i] + gelu_scalar;
    }
}

torch::Tensor fused_forward_cuda(torch::Tensor x,
                                 torch::Tensor weight,
                                 torch::Tensor bias,
                                 torch::Tensor subtract,
                                 bool has_bias) {
// [PK]: Setup omitted for readabilty

    // Launch kernels
    dim3 block_wsum(BLOCK);
    dim3 grid_wsum(in_features);
    weight_col_sum_kernel<<<grid_wsum, block_wsum, 0, stream.stream()>>>(W_ptr, out_features, in_features, wsum_ptr);

    dim3 block_bsum(BLOCK);
    dim3 grid_bsum(1);
    bias_sub_sum_kernel<<<grid_bsum, block_bsum, 0, stream.stream()>>>(bias_ptr, sub_ptr, has_bias ? 1 : 0, out_features, bs_sum_ptr);

    dim3 block_out(BLOCK);
    dim3 grid_out(B);
    fused_output_kernel<<<grid_out, block_out, 0, stream.stream()>>>(
        X_ptr,
        wsum_ptr,
        bs_sum_ptr,
        B,
        in_features,
        out_features,
        Out_ptr
    );

    return out;
}

// Python/Pytorch related code omitted for readability

*Pseudocode:*
gemv(x, gemm_weights_summed_per_row)/out_features