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

Code before:
            if (i0 < N * N) {
                sA[i0] = A_batch[i0];
            }
        }

#pragma unroll
        for (int i = 0; i < N * K; i += K * WARP_SIZE) {
            int i0 = i + threadIdx.x + threadIdx.y * blockDim.x;
            if (i0 < N * K) {
                sX[i0] = B_batch[i0];
            }
        }
    }

    __syncthreads();

    for (int row = 0; row < max(n, N); ++row) {
        float sum = 0.0f;

        if constexpr (K == 0) {
            for (int j = lane; j < row; j += WARP_SIZE) {
                sum += sA[row * n + j] * sX[j * k + col_idx];
            }
        } else {
            for (int j = lane; j < row; j += WARP_SIZE) {
                sum += sA[row * N + j] * sX[j * K + col_idx];
            }
        }

        sum = warp_reduce_sum(sum);

        if (lane == 0) {
            if constexpr (K == 0) {
                const float b_val  = sX[row * k + col_idx];  // Value from B
                const float a_diag = sA[row * n + row];
                if (a_diag != 0.0f) {
                    sX[row * k + col_idx] = (b_val - sum) / a_diag;
                } else {
                    sX[row * k + col_idx] = 0.0f;  // Avoid division by zero
                }
            } else {
                const float b_val  = sX[row * K + col_idx];  // Value from B
                const float a_diag = sA[row * N + row];
                if (a_diag != 0.0f) {
                    sX[row * K + col_idx] = (b_val - sum) / a_diag;
                } else {
                    sX[row * K + col_idx] = 0.0f;  // Avoid division by zero
                }
            }
        }
        __syncthreads();

Review comment (suggestion):
Without having looked at NSight Compute, the `sX` memory accesses here should be suffering massively from shared memory bank conflicts. In the worst case, if `K` is a multiple of 32 then each memory access will be going to the same shared memory bank, meaning those accesses need to be serialized and take 32 times as long. This can be fixed by padding `sX` with 4 extra bytes per row so that each thread in a warp ends up accessing a different shared memory bank.

Code after:
    const int rows_per_warp = (n + WARP_SIZE - 1) / WARP_SIZE;

#pragma unroll
    for (int i = 0; i < rows_per_warp; i++) {
        const int i0 = lane + i * WARP_SIZE;
        if (i0 < n) {
            sXt[col_idx * n + i0] = B_batch[i0 * k + col_idx];
        }
    }

    __syncthreads();

#pragma unroll
    for (int row = 0; row < n; ++row) {
        float sum = 0.0f;

        {
            int j = lane;
            if (j < row) {
                sum += sA[row * n + j] * sXt[col_idx * n + j];
            }
        }
        if (row >= WARP_SIZE) {
            int j = WARP_SIZE + lane;
            if (j < row) {
                sum += sA[row * n + j] * sXt[col_idx * n + j];
            }
        }

        sum = warp_reduce_sum(sum);

        if (lane == 0) {
            const float b_val      = sXt[col_idx * n + row];
            const float a_diag     = sA[row * n + row];
            // no safeguards for division by zero because that indicates corrupt
            // data anyway
            sXt[col_idx * n + row] = (b_val - sum) / a_diag;
        }
    }

    __syncthreads();

#pragma unroll
    for (int i = 0; i < rows_per_warp; i++) {
        const int i0 = lane + i * WARP_SIZE;
        if (i0 < n) {
            X_batch[i0 * k + col_idx] = sXt[col_idx * n + i0];
        }
    }
}
#ifdef __clang__