Method operator()

csrc/xpu_kernels.cpp:97–172

template <typename T, int TILE_SIZE, int NUM_PER_TH, int DATA_TYPE>
SYCL_EXTERNAL void kDequantizeBlockwise<T, TILE_SIZE, NUM_PER_TH, DATA_TYPE>::operator()(sycl::nd_item<1> item) const {
    const int64_t base_idx = static_cast<int64_t>(item.get_group(0)) * TILE_SIZE;
    int64_t local_idx = static_cast<int64_t>(item.get_local_id(0)) * NUM_PER_TH;
    float local_abs_max = -FLT_MAX;
    int64_t local_load_idx = 0;
    int64_t local_store_idx = 0;

    uint8_t qvals[NUM_PER_TH];
    T vals[NUM_PER_TH * ((DATA_TYPE > 0) ? 2 : 1)];

    if (DATA_TYPE > 0) {
        // Cast n to int64_t to avoid overflow for large n (same as CUDA)
        local_load_idx = sycl::min(static_cast<int64_t>(TILE_SIZE), (static_cast<int64_t>(n) + 1) / 2 - base_idx);
        local_store_idx = sycl::min(static_cast<int64_t>(TILE_SIZE * 2), static_cast<int64_t>(n) - base_idx * 2);
    } else {
        local_load_idx = sycl::min(static_cast<int64_t>(TILE_SIZE), static_cast<int64_t>(n) - base_idx);
        local_store_idx = local_load_idx;
    }

    // Avoid expensive division by the blocksize (as blocksize will always be a
    // power-of-2)
    local_abs_max = absmax[(base_idx + local_idx) >> (31 - std::countl_zero<unsigned int>(blocksize))];

    if (local_idx + NUM_PER_TH < local_load_idx) {
        reinterpret_cast<sycl::vec<uint8_t, NUM_PER_TH>(&)[NUM_PER_TH]>(qvals)[0] =
            reinterpret_cast<sycl::vec<uint8_t, NUM_PER_TH>*>(A)[(base_idx + local_idx) / NUM_PER_TH];
    } else {
#pragma unroll NUM_PER_TH
        for (int i = 0; i < NUM_PER_TH; i++) {
            if (local_idx + i < local_load_idx) {
                qvals[i] = A[base_idx + local_idx + i];
            } else {
                qvals[i] = (uint8_t)0;
            }
        }
    }

    switch (DATA_TYPE) {
    case General8bit:
#pragma unroll NUM_PER_TH
        for (int j = 0; j < NUM_PER_TH; j++)
            vals[j] = code[qvals[j]] * local_abs_max;
        break;
    case FP4:
#pragma unroll NUM_PER_TH
        for (int j = 0; j < NUM_PER_TH; j++) {
            vals[j * 2] = dDequantizeFP4(qvals[j] >> 4) * local_abs_max;
            vals[j * 2 + 1] = dDequantizeFP4(qvals[j] & 0x0F) * local_abs_max;
        }
        break;
    case NF4:
#pragma unroll NUM_PER_TH
        for (int j = 0; j < NUM_PER_TH; j++) {
            vals[j * 2] = dDequantizeNF4(qvals[j] >> 4) * local_abs_max;
            vals[j * 2 + 1] = dDequantizeNF4(qvals[j] & 0x0F) * local_abs_max;
        }
        break;
    }