untitled

#include <cuda_runtime.h>

__global__ void interleave_kernel(const float* A, const float* B, float* output, int N) {
    if (threadIdx.x == 0) {
        output[2*N-2] = A[N-1];
        output[2*N-1] = B[N-1];
    }
}

__global__ void interleave_kernel_vectorization(const float2* __restrict__ A2, const float2* __restrict__ B2, float4* __restrict__ output4, int N) {
    int idx = blockDim.x * blockIdx.x + threadIdx.x;

    if (idx < N) {
        float2 tmpA = A2[idx];
        float2 tmpB = B2[idx];
        output4[idx] = make_float4(tmpA.x, tmpB.x, tmpA.y, tmpB.y);
    } 
}

// A, B, output are device pointers (i.e. pointers to memory on the GPU)
extern "C" void solve(const float* A, const float* B, float* output, int N) {
    int N2 = N >> 1;
    int tail = N & 1;

    if (N2 > 0) {
        const float2* A2 = reinterpret_cast<const float2*> (A);
        const float2* B2 = reinterpret_cast<const float2*> (B);
        float4*  output4 = reinterpret_cast<float4*> (output);

        int threadsPerBlock = 256;
        int blocksPerGrid = (N2 + threadsPerBlock - 1) / threadsPerBlock;
        interleave_kernel_vectorization<<<blocksPerGrid, threadsPerBlock>>>(A2, B2, output4, N2);
    }

    if (tail) {
        interleave_kernel<<<1, 1>>>(A, B, output, N);
    }

    // interleave_kernel<<<blocksPerGrid, threadsPerBlock>>>(A, B, output, N);
    cudaDeviceSynchronize();
}