ronantakizawa/codereview-bench huggingface dataset:
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/cumsum.cu (CUDA)

Code before:

        // store result
        if (idx < ne00) {
            dst_row[idx] = final_val;
        }

        __syncthreads();

        if (tid == 0) {
            block_carry += block_total;
        }

        __syncthreads();
    }
#else
    NO_DEVICE_CODE;
#endif
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
    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) {

    GGML_UNUSED_VARS(nb00, nb0);

    const int tid = threadIdx.x;
    const int lane = tid & (WARP_SIZE - 1);
    const int warp = tid / WARP_SIZE;
    const int warps_per_block = blockDim.x / WARP_SIZE;

    extern __shared__ float smem[];
    float* s_vals = smem;
    float* s_warp_sums = smem + blockDim.x;
    float* s_carry = smem + blockDim.x + warps_per_block;
    float* s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;

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) {
```

Code after:

        // store result
        if (idx < ne00) {
            dst_row[idx] = final_val;
        }

        __syncthreads();

        if (tid == 0) {
            block_carry += block_total;
        }

        __syncthreads();
    }
#else
    NO_DEVICE_CODE;
#endif // GGML_CUDA_USE_CUB
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
        const T * src, T * dst,
        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
        const int64_t  s00, const int64_t  s01, const int64_t  s02, const int64_t  s03,
        const int64_t   s0, const int64_t   s1, const int64_t   s2, const int64_t   s3) {

    GGML_UNUSED_VARS(s00, s0);

    const int tid = threadIdx.x;
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    const int lane = tid % warp_size;
    const int warp = tid / warp_size;
    const int warps_per_block = blockDim.x / warp_size;

    extern __shared__ float smem[];
    float * s_vals = smem;
    float * s_warp_sums = smem + blockDim.x;
    float * s_carry = smem + blockDim.x + warps_per_block;
    float * s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/cumsum.cu (CUDA)

Code before:

        __syncthreads();

        if (tid == 0) {
            block_carry += block_total;
        }

        __syncthreads();
    }
#else
    NO_DEVICE_CODE;
#endif
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
    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) {

    GGML_UNUSED_VARS(nb00, nb0);

    const int tid = threadIdx.x;
    const int lane = tid & (WARP_SIZE - 1);
    const int warp = tid / WARP_SIZE;
    const int warps_per_block = blockDim.x / WARP_SIZE;

    extern __shared__ float smem[];
    float* s_vals = smem;
    float* s_warp_sums = smem + blockDim.x;
    float* s_carry = smem + blockDim.x + warps_per_block;
    float* s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

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


Review comment (suggestion):
Since `WARP_SIZE` is known at compile time the compiler will do this substitution automatically. In my opinion `tid % WARP_SIZE` is more readable.

Code after:

        __syncthreads();

        if (tid == 0) {
            block_carry += block_total;
        }

        __syncthreads();
    }
#else
    NO_DEVICE_CODE;
#endif // GGML_CUDA_USE_CUB
}

// Fallback kernel implementation (original)
template<typename T>
static __global__ void cumsum_kernel(
        const T * src, T * dst,
        const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03,
        const int64_t  s00, const int64_t  s01, const int64_t  s02, const int64_t  s03,
        const int64_t   s0, const int64_t   s1, const int64_t   s2, const int64_t   s3) {

    GGML_UNUSED_VARS(s00, s0);

    const int tid = threadIdx.x;
    constexpr int warp_size = ggml_cuda_get_physical_warp_size();
    const int lane = tid % warp_size;
    const int warp = tid / warp_size;
    const int warps_per_block = blockDim.x / warp_size;

    extern __shared__ float smem[];
    float * s_vals = smem;
    float * s_warp_sums = smem + blockDim.x;
    float * s_carry = smem + blockDim.x + warps_per_block;
    float * s_chunk_total = s_carry + 1;

    // Initialize carry
    if (tid == 0) {
        *s_carry = 0.0f;
    }
    __syncthreads();

    const int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

    const T * src_row = src + i1 * s01 + i2 * s02 + i3 * s03;
    T       * dst_row = dst + i1 * s1  + i2 * s2  + i3 * s3;
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/tri.cu (CUDA)

Code before:
#include "common.cuh"
#include "convert.cuh"
#include "tri.cuh"
#include "ggml.h"

template<typename T, bool prefix_keep, int add_to_split>
static __global__ void tri_kernel(
    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 int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    const int64_t split_point = i1 + add_to_split;

    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

    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];

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) {
```

Code after:
#include "common.cuh"
#include "convert.cuh"
#include "tri.cuh"
#include "ggml.h"

template<typename T, bool prefix_keep, int add_to_split>
static __global__ void tri_kernel(
        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 int64_t i3 = blockIdx.z;
    const int64_t i2 = blockIdx.y;
    const int64_t i1 = blockIdx.x;
    const int64_t split_point = i1 + add_to_split;

    GGML_UNUSED_VARS(nb00, nb0);

    if (i3 >= ne03 || i2 >= ne02 || i1 >= ne01) {
        return;
    }

    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);
        }
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/common.cuh (CUDA)

Code before:
        int     cc;                 // compute capability
        int     nsm;                // number of streaming multiprocessors
        size_t  smpb;               // max. shared memory per block
        size_t  smpbo;              // max. shared memory per block (with opt-in)
        bool    integrated;         // Device is integrated as opposed to discrete
        bool    vmm;                // virtual memory support
        size_t  vmm_granularity;    // granularity of virtual memory
        size_t  total_vram;
        int     warp_size;          // Number of threads in a dispatch
    };

    cuda_device_info devices[GGML_CUDA_MAX_DEVICES] = {};

    std::array<float, GGML_CUDA_MAX_DEVICES> default_tensor_split = {};
};

const ggml_cuda_device_info & ggml_cuda_info();

void ggml_cuda_set_device(int device);
int ggml_cuda_get_device();

static __host__ int ggml_cuda_get_physical_warp_size_host() {
    const auto &info = ggml_cuda_info().devices[ggml_cuda_get_device()];
    return info.warp_size;
}

struct ggml_cuda_pool {
    virtual ~ggml_cuda_pool() = default;

    virtual void * alloc(size_t size, size_t * actual_size) = 0;
    virtual void free(void * ptr, size_t size) = 0;
};

template<typename T>
struct ggml_cuda_pool_alloc {
    ggml_cuda_pool * pool = nullptr;
    T * ptr = nullptr;
    size_t actual_size = 0;

    ggml_cuda_pool_alloc() = default;

    explicit ggml_cuda_pool_alloc(ggml_cuda_pool & pool) : pool(&pool) {
    }

    ggml_cuda_pool_alloc(ggml_cuda_pool & pool, size_t size) : pool(&pool) {
        alloc(size);
    }

    ~ggml_cuda_pool_alloc() {
        if (ptr != nullptr) {
            pool->free(ptr, actual_size);

Review comment (suggestion):
```suggestion
```

I think the more ergonomic pattern is to get `ggml_cuda_device_info` struct in some host function and to then access its properties. I would like to avoid adding functions like this.

Code after:
};

const ggml_cuda_device_info & ggml_cuda_info();

void ggml_cuda_set_device(int device);
int ggml_cuda_get_device();

struct ggml_cuda_pool {
    virtual ~ggml_cuda_pool() = default;

    virtual void * alloc(size_t size, size_t * actual_size) = 0;
    virtual void free(void * ptr, size_t size) = 0;
};

template<typename T>
struct ggml_cuda_pool_alloc {
    ggml_cuda_pool * pool = nullptr;
    T * ptr = nullptr;
    size_t actual_size = 0;

    ggml_cuda_pool_alloc() = default;

    explicit ggml_cuda_pool_alloc(ggml_cuda_pool & pool) : pool(&pool) {
    }

    ggml_cuda_pool_alloc(ggml_cuda_pool & pool, size_t size) : pool(&pool) {
        alloc(size);
    }

    ~ggml_cuda_pool_alloc() {
        if (ptr != nullptr) {
            pool->free(ptr, actual_size);
        }
    }

    // size is in number of elements
    T * alloc(size_t size) {
        GGML_ASSERT(pool != nullptr);
        GGML_ASSERT(ptr == nullptr);
        ptr = (T *) pool->alloc(size * sizeof(T), &this->actual_size);
        return ptr;
    }

    T * alloc(ggml_cuda_pool & pool, size_t size) {
        this->pool = &pool;
        return alloc(size);
    }

    T * get() {
        return ptr;
    }
Repository: ggerganov/llama.cpp
File: common/download.cpp (C++)

Code before:
    curl_easy_getinfo(curl.get(), CURLINFO_RESPONSE_CODE, &res_code);

    return { res_code, std::move(res_buffer) };
}

#elif defined(LLAMA_USE_HTTPLIB)

static bool is_output_a_tty() {
#if defined(_WIN32)
    return _isatty(_fileno(stdout));
#else
    return isatty(1);
#endif
}

static void print_progress(size_t current, size_t total) {
    if (!is_output_a_tty()) {
        return;
    }

    if (!total) {
        return;
    }

    static std::mutex mutex;
    static std::map<std::thread::id, int> lines;

    std::lock_guard<std::mutex> lock(mutex);
    std::thread::id id = std::this_thread::get_id();

    if (lines.find(id) == lines.end()) {
        lines[id] = lines.size();
        std::cout << "\n";
    }
    int lines_up = lines.size() - lines[id];

    size_t width = 50;
    size_t pct = (100 * current) / total;
    size_t pos = (width * current) / total;

    std::cout << "\033[s";

    if (lines_up > 0) {
        std::cout << "\033[" << lines_up << "A";
    }
    std::cout << "\033[2K\r["
              << std::string(pos, '=')
              << (pos < width ? ">" : "")
              << std::string(width - pos, ' ')
              << "] " << std::setw(3) << pct << "%  ("
              << current / (1024 * 1024) << " MB / "

Review comment (suggestion):
This looks like a local variable at first glance, but because it's `static`, that effectively make it a global variable.

IMO we should avoid global variable if possible. Probably a better way is to make a dedicated class to manage threads and implement the mutex as a member of the class.

Code after:
    curl_easy_getinfo(curl.get(), CURLINFO_RESPONSE_CODE, &res_code);

    return { res_code, std::move(res_buffer) };
}

#elif defined(LLAMA_USE_HTTPLIB)

class ProgressBar {
    static inline std::mutex mutex;
    static inline std::map<const ProgressBar *, int> lines;
    static inline int max_line = 0;

    static void cleanup(const ProgressBar * line) {
        lines.erase(line);
        if (lines.empty()) {
            max_line = 0;
        }
    }

    static bool is_output_a_tty() {
#if defined(_WIN32)
        return _isatty(_fileno(stdout));
#else
        return isatty(1);
#endif
    }

public:
    ProgressBar() = default;

    ~ProgressBar() {
        std::lock_guard<std::mutex> lock(mutex);
        cleanup(this);
    }

    void update(size_t current, size_t total) {
        if (!is_output_a_tty()) {
            return;
        }

        if (!total) {
            return;
        }

        std::lock_guard<std::mutex> lock(mutex);

        if (lines.find(this) == lines.end()) {
            lines[this] = max_line++;
            std::cout << "\n";
        }
        int lines_up = max_line - lines[this];
Repository: ggerganov/llama.cpp
File: tools/server/server-models.cpp (C++)

Code before:
        req.content_receiver = content_receiver;
    }

    // start the proxy thread
    SRV_DBG("start proxy thread %s %s\n", req.method.c_str(), req.path.c_str());
    this->thread = std::thread([cli, pipe, req]() {
        auto result = cli->send(std::move(req));
        if (result.error() != httplib::Error::Success) {
            auto err_str = httplib::to_string(result.error());
            SRV_ERR("http client error: %s\n", err_str.c_str());
            pipe->write({{}, 500, "", ""}); // header
            pipe->write({{}, 0, "proxy error: " + err_str, ""}); // body
        }
        pipe->close_write(); // signal EOF to reader
        SRV_DBG("%s", "client request thread ended\n");
    });
    this->thread.detach();

    // wait for the first chunk (headers)
    msg_t header;
    if (pipe->read(header, should_stop)) {
        SRV_DBG("%s", "received response headers\n");
        this->status  = header.status;
        this->headers = std::move(header.headers);
        if (!header.content_type.empty()) {
            this->content_type = std::move(header.content_type);
        }
    } else {
        SRV_DBG("%s", "no response headers received (request cancelled?)\n");
    }
}


Review comment (suggestion):
If we std::move here, I prefer restrict the scope of `header` to avoid someone (in the future) accidentally use `header` after this point

Code after:
        auto result = cli->send(std::move(req));
        if (result.error() != httplib::Error::Success) {
            auto err_str = httplib::to_string(result.error());
            SRV_ERR("http client error: %s\n", err_str.c_str());
            pipe->write({{}, 500, "", ""}); // header
            pipe->write({{}, 0, "proxy error: " + err_str, ""}); // body
        }
        pipe->close_write(); // signal EOF to reader
        SRV_DBG("%s", "client request thread ended\n");
    });
    this->thread.detach();

    // wait for the first chunk (headers)
    {
        msg_t header;
        if (pipe->read(header, should_stop)) {
            SRV_DBG("%s", "received response headers\n");
            this->status  = header.status;
            this->headers = std::move(header.headers);
            if (!header.content_type.empty()) {
                this->content_type = std::move(header.content_type);
            }
        } else {
            SRV_DBG("%s", "no response headers received (request cancelled?)\n");
        }
    }
}
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/solve_tri.cu (CUDA)

Code before:

    __shared__ float sA[MAX_N_FAST * MAX_N_FAST];

    const int offset = threadIdx.x + threadIdx.y * blockDim.x;

#pragma unroll
    for (int i = 0; i < n * n; i += k * WARP_SIZE) {
        int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    __syncthreads();

    float x_low  = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f;
    float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f;

    const int half = WARP_SIZE;
    const int nrows_low = (n < half) ? n : half;

#pragma unroll
    for (int row = 0; row < nrows_low; ++row) {
        float sum = 0.0f;
        if (lane < row) {
            sum = fmaf(sA[row * n + lane], x_low, sum);
        }
        sum = warp_reduce_sum(sum);

        if (lane == row) {
            float diag = sA[row * n + row];
            float idiv = 1.0f / diag;
            x_low = fmaf(sum, -idiv, x_low * idiv);
        }
    }

#pragma unroll
    for (int row = half; row < n; ++row) {
        float sum = fmaf(sA[row * n + lane], x_low, 0.0f);
        int j = half + lane;
        if (j < row) {
            sum = fmaf(sA[row * n + j], x_high, sum);
        }
        sum = warp_reduce_sum(sum);

        int updater = row - half;
        if (lane == updater) {
            float diag = sA[row * n + row];
            float idiv = 1.0f / diag;
            x_high = fmaf(sum, -idiv, x_high * idiv);
        }

Review comment (suggestion):
```suggestion
            sum += sA[row * n + lane] * x_low;
```

The compiler will do the substitution automatically, please replace the `fmaf` instructions with syntax like this as it will be more readable.

Code after:

    __shared__ float sA[MAX_N_FAST * MAX_N_FAST];

    const int offset = threadIdx.x + threadIdx.y * blockDim.x;

#pragma unroll
    for (int i = 0; i < n * n; i += k * WARP_SIZE) {
        const int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    __syncthreads();

    float x_low  = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f;
    float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f;

    const int half = WARP_SIZE;
    const int nrows_low = (n < half) ? n : half;

#pragma unroll
    for (int row = 0; row < nrows_low; ++row) {
        float sum = 0.0f;
        if (lane < row) {
            sum += sA[row * n + lane] * x_low;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row) {
            x_low = (x_low - sum) / sA[row * n + row];
        }
    }

#pragma unroll
    for (int row = half; row < n; ++row) {
        float sum = sA[row * n + lane] * x_low;
        const int j = half + lane;
        if (j < row) {
            sum += sA[row * n + j] * x_high;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row - half) {
            x_high = (x_high - sum) / sA[row * n + row];
        }
    }

#pragma unroll
    for (int rr = 0; rr < 2; ++rr) {
        const int row = rr * WARP_SIZE + lane;
Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/solve_tri.cu (CUDA)

Code before:
        const int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    __syncthreads();

    float x_low  = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f;
    float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f;

    const int half = WARP_SIZE;
    const int nrows_low = (n < half) ? n : half;

#pragma unroll
    for (int row = 0; row < nrows_low; ++row) {
        float sum = 0.0f;
        if (lane < row) {
            sum += sA[row * n + lane] * x_low;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row) {
            const float diag = sA[row * n + row];
            const float idiv = 1.0f / diag;
            x_low = fmaf(sum, -idiv, x_low * idiv);
        }
    }

#pragma unroll
    for (int row = half; row < n; ++row) {
        float sum = sA[row * n + lane] * x_low;
        const int j = half + lane;
        if (j < row) {
            sum += sA[row * n + j] * x_high;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row - half) {
            x_high = (x_high - sum) / sA[row * n + row];
        }
    }

#pragma unroll
    for (int rr = 0; rr < 2; ++rr) {
        const int row = rr * WARP_SIZE + lane;
        if (row < n) {
            const float val = (row < half) ? x_low : x_high;
            X_batch[row * k + col_idx] = val;
        }
    }

Review comment (suggestion):
```suggestion
            x_low = (x_low - sum) / sA[row * n + row]
```

Please change this spot as well.

Code after:
        const int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    __syncthreads();

    float x_low  = (lane < n) ? B_batch[lane * k + col_idx] : 0.0f;
    float x_high = (WARP_SIZE + lane < n) ? B_batch[(WARP_SIZE + lane) * k + col_idx] : 0.0f;

    const int half = WARP_SIZE;
    const int nrows_low = (n < half) ? n : half;

#pragma unroll
    for (int row = 0; row < nrows_low; ++row) {
        float sum = 0.0f;
        if (lane < row) {
            sum += sA[row * n + lane] * x_low;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row) {
            x_low = (x_low - sum) / sA[row * n + row];
        }
    }

#pragma unroll
    for (int row = half; row < n; ++row) {
        float sum = sA[row * n + lane] * x_low;
        const int j = half + lane;
        if (j < row) {
            sum += sA[row * n + j] * x_high;
        }
        sum = warp_reduce_sum(sum);

        if (lane == row - half) {
            x_high = (x_high - sum) / sA[row * n + row];
        }
    }

#pragma unroll
    for (int rr = 0; rr < 2; ++rr) {
        const int row = rr * WARP_SIZE + lane;
        if (row < n) {
            const float val = (row < half) ? x_low : x_high;
            X_batch[row * k + col_idx] = val;
        }
    }
}
#ifdef __clang__