Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/pad_reflect_1d.cu (CUDA)

Code before:
#include "pad_reflect_1d.cuh"

static __global__ void pad_reflect_1d_kernel_f32(
    const void * __restrict__ src0,
    void * __restrict__ dst,
    const int64_t ne0,
    const int64_t ne00,
    const int64_t ne01,
    const int64_t ne02,
    const int64_t ne03,
    const int64_t nb00,
    const int64_t nb01,
    const int64_t nb02,
    const int64_t nb03,
    const int64_t nb0,
    const int64_t nb1,
    const int64_t nb2,
    const int64_t nb3,
    const int p0,
    const int p1) {

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;

    const int64_t tile1  = blockIdx.x % ne01;     // i1
    const int64_t tile0  = blockIdx.x / ne01;     // nth i0 tile
    const int64_t i1     = tile1;
    const int64_t i0     = threadIdx.x + tile0 * blockDim.x;
    if ( i0 >= ne0 || i1 >= ne01 || i2 >= ne02 || i3 >= ne03 ) {
        return;
    }

    const char * src0_ptr = (const char *)src0 + i3*nb03 + i2*nb02 + i1*nb01;
    const char * dst_ptr = (const char *)dst + i3*nb3 + i2*nb2 + i1*nb1;

    float value;
    const int64_t j = i0 - p0;

    if ( j<0 ) {// i0<p0
        // Left padding - reflect
        value = *(const float *)(src0_ptr - j * nb00);
    } else if (j < ne00) { //i0 < ne0 - p1
        // Middle - copy
        value = *(const float *)(src0_ptr + j * nb00);
    } else  {
        // Right padding - reflect
        const int64_t src_idx = (ne0 - p1 - p0) - (p1 + 1 - (ne0 - i0)) - 1;
        value = *(const float *)(src0_ptr + src_idx * nb00);
    }
    *(float *)(dst_ptr + i0 * nb0) = value;
}

Review comment (suggestion):
Look into `fastdiv` defined in `common.cuh`. For such a small kernel this single division/modulo probably takes up a significant percentage of the runtime.

Code after:
#include "pad_reflect_1d.cuh"

static __global__ __launch_bounds__(CUDA_PAD_REFLECT_1D_BLOCK_SIZE, 1) void
    pad_reflect_1d_kernel_f32(
        const void * __restrict__ src0,
        void * __restrict__       dst,
        const int64_t             ne0,
        const int64_t             ne00,
        const uint3               ne01,
        const int64_t             ne02,
        const int64_t             ne03,
        const int64_t             nb00,
        const int64_t             nb01,
        const int64_t             nb02,
        const int64_t             nb03,
        const int64_t             nb0,
        const int64_t             nb1,
        const int64_t             nb2,
        const int64_t             nb3,
        const int                 p0,
        const int                 p1) {
    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;

    const uint2   div_mod_packed = fast_div_modulo(blockIdx.x, ne01);
    const int64_t tile1          = div_mod_packed.y;  // i1
    const int64_t tile0          = div_mod_packed.x;  // nth i0 tile
    const int64_t i1             = tile1;
    const int64_t i0             = threadIdx.x + tile0 * blockDim.x;

    // ne01.z is original value of unpacked ne01 (see init_fastdiv_values in common.cuh)
    if (i0 >= ne0 || i1 >= ne01.z || i2 >= ne02 || i3 >= ne03) {
        return;
    }

    const char * src0_ptr = (const char *) src0 + i3 * nb03 + i2 * nb02 + i1 * nb01;
    char *       dst_ptr  = (char *) dst + i3 * nb3 + i2 * nb2 + i1 * nb1;

    const int64_t rel_i0 = i0 - p0;  // relative i0 in src0
    int64_t src_idx;

    if (rel_i0 < 0) {
        // Left padding - reflect
        src_idx = -rel_i0;
    } else if (rel_i0 < ne00) {
        // Middle - copy
        src_idx = rel_i0;
    } else {
        // Right padding - reflect
        src_idx = 2 * ne00 - 2 - rel_i0;
    }