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

Code before:
}

static __device__ __forceinline__ float op_mul(const float a, const float b) {
    return a * b;
}

static __device__ __forceinline__ float op_div(const float a, const float b) {
    return a / b;
}

template <float (*bin_op)(const float, const float),
          typename src0_t,
          typename src1_t,
          typename dst_t,
          typename... src1_ptrs>
static __global__ void k_bin_bcast(const src0_t *         src0,
                                   const src1_t *         src1,
                                   dst_t *                dst,
                                   const int              ne0,
                                   const int              ne1,
                                   const int              ne2,
                                   const uint3            ne3_fastdiv,
                                   const uint3            ne10_fastdiv,
                                   const uint3            ne11_fastdiv,
                                   const uint3            ne12_fastdiv,
                                   const uint3            ne13_fastdiv,
                                   /*int s0, */ const int s1,
                                   const int              s2,
                                   const int              s3,
                                   /*int s00,*/ const int s01,
                                   const int              s02,
                                   const int              s03,
                                   /*int s10,*/ const int s11,
                                   const int              s12,
                                   const int              s13,
                                   src1_ptrs... src1s) {
    const uint32_t i0s = blockDim.x * blockIdx.x + threadIdx.x;
    const uint32_t i1  = (blockDim.y * blockIdx.y + threadIdx.y);
    const uint32_t i2  = fastdiv((blockDim.z * blockIdx.z + threadIdx.z), ne3_fastdiv);
    const uint32_t i3  = (blockDim.z * blockIdx.z + threadIdx.z) - (i2 * ne3_fastdiv.z);

    if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3_fastdiv.z) {
        return;
    }

    const uint32_t i11 = fastmodulo(i1, ne11_fastdiv);
    const uint32_t i12 = fastmodulo(i2, ne12_fastdiv);
    const uint32_t i13 = fastmodulo(i3, ne13_fastdiv);

    const size_t i_src0 =  i3*s03 +  i2*s02 +  i1*s01;
    const size_t i_src1 = i13*s13 + i12*s12 + i11*s11;

Review comment (suggestion):
Preferably drop the `_fastdiv` prefix for the kernel arguments since there is no name conflict.

Code after:
}

static __device__ __forceinline__ float op_mul(const float a, const float b) {
    return a * b;
}

static __device__ __forceinline__ float op_div(const float a, const float b) {
    return a / b;
}

template <float (*bin_op)(const float, const float),
          typename src0_t,
          typename src1_t,
          typename dst_t,
          typename... src1_ptrs>
static __global__ void k_bin_bcast(const src0_t *         src0,
                                   const src1_t *         src1,
                                   dst_t *                dst,
                                   const int              ne0,
                                   const int              ne1,
                                   const int              ne2,
                                   const uint3            ne3,
                                   const uint3            ne10,
                                   const uint3            ne11,
                                   const uint3            ne12,
                                   const uint3            ne13,
                                   /*int s0, */ const int s1,
                                   const int              s2,
                                   const int              s3,
                                   /*int s00,*/ const int s01,
                                   const int              s02,
                                   const int              s03,
                                   /*int s10,*/ const int s11,
                                   const int              s12,
                                   const int              s13,
                                   src1_ptrs... src1s) {
    const uint32_t i0s = blockDim.x * blockIdx.x + threadIdx.x;
    const uint32_t i1  = (blockDim.y * blockIdx.y + threadIdx.y);
    const uint32_t i2  = fastdiv((blockDim.z * blockIdx.z + threadIdx.z), ne3);
    const uint32_t i3  = (blockDim.z * blockIdx.z + threadIdx.z) - (i2 * ne3.z);

    if (i0s >= ne0 || i1 >= ne1 || i2 >= ne2 || i3 >= ne3.z) {
        return;
    }

    const uint32_t i11 = fastmodulo(i1, ne11);
    const uint32_t i12 = fastmodulo(i2, ne12);
    const uint32_t i13 = fastmodulo(i3, ne13);

    const size_t i_src0 =  i3*s03 +  i2*s02 +  i1*s01;
    const size_t i_src1 = i13*s13 + i12*s12 + i11*s11;