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

Code before:

    if constexpr (prefix_keep) {
        for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
            dst_row[i0] = src_row[i0];
        }
        for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
            dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
        }
    } else {
        for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
            dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
        }
        for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
            dst_row[i0] = src_row[i0];
        }
    }
}

template<typename T>
static void tri_cuda(
    const T * src, T * dst,
    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 ggml_tri_type ttype,
    cudaStream_t stream) {

    dim3 block_dims(CUDA_TRI_BLOCK_SIZE, 1, 1);
    dim3 grid_dims(ne01, ne02, ne03);
    const size_t type_size = sizeof(T);

    const int add_to_split = (ttype == GGML_TRI_TYPE_LOWER_DIAG || ttype == GGML_TRI_TYPE_UPPER) ? 1 : 0;
    const bool prefix_keep = (ttype == GGML_TRI_TYPE_LOWER || ttype == GGML_TRI_TYPE_LOWER_DIAG);

    if (prefix_keep) {
        if (add_to_split == 0) {
            tri_kernel<T, true, 0><<<grid_dims, block_dims, 0, stream>>>(
                src, dst,
                ne00, ne01, ne02, ne03,
                nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
                nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
            );
        } else { // only 0 and 1 supported
            tri_kernel<T, true, 1><<<grid_dims, block_dims, 0, stream>>>(
                src, dst,
                ne00, ne01, ne02, ne03,
                nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
                nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
            );
        }
    } else {

Review comment (suggestion):
```suggestion
        const T * src, T * dst,
        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 ggml_tri_type ttype,
        cudaStream_t stream) {
```

Code after:
    const T * src_row = src + i1*nb01 + i2*nb02 + i3*nb03;
    T       * dst_row = dst + i1*nb1  + i2*nb2  + i3*nb3;

    if constexpr (prefix_keep) {
        for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
            dst_row[i0] = src_row[i0];
        }
        for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
            dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
        }
    } else {
        for (int64_t i0 = threadIdx.x; i0 < split_point; i0 += blockDim.x) {
            dst_row[i0] = ggml_cuda_cast<T, float>(0.0f);
        }
        for (int64_t i0 = threadIdx.x + split_point; i0 < ne00; i0 += blockDim.x) {
            dst_row[i0] = src_row[i0];
        }
    }
}

template<typename T>
static void tri_cuda(
        const T * src, T * dst,
        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 ggml_tri_type ttype,
        cudaStream_t stream) {

    dim3 block_dims(CUDA_TRI_BLOCK_SIZE, 1, 1);
    dim3 grid_dims(ne01, ne02, ne03);
    const size_t type_size = sizeof(T);

    const int add_to_split = (ttype == GGML_TRI_TYPE_LOWER_DIAG || ttype == GGML_TRI_TYPE_UPPER) ? 1 : 0;
    const bool prefix_keep = (ttype == GGML_TRI_TYPE_LOWER || ttype == GGML_TRI_TYPE_LOWER_DIAG);

    if (prefix_keep) {
        if (add_to_split == 0) {
            tri_kernel<T, true, 0><<<grid_dims, block_dims, 0, stream>>>(
                src, dst,
                ne00, ne01, ne02, ne03,
                nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
                nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
            );
        } else { // only 0 and 1 supported
            tri_kernel<T, true, 1><<<grid_dims, block_dims, 0, stream>>>(
                src, dst,
                ne00, ne01, ne02, ne03,
                nb00 / type_size, nb01 / type_size, nb02 / type_size, nb03 / type_size,
                nb0 / type_size, nb1 / type_size, nb2 / type_size, nb3 / type_size
            );