wehub-resource-sync
539 lines
19 KiB
Plaintext
539 lines
19 KiB
Plaintext
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
#include <iostream>
|
|
#include <cstdio>
|
|
#include <cuda_runtime.h>
|
|
#include "cachegen_kernels.cuh"
|
|
|
|
#define MAX_LP 64
|
|
#define MAX_THREAD_PER_BLOCK 128
|
|
#define MAX_SHARED_MEMORY_PER_THREAD (0xc000 / MAX_THREAD_PER_BLOCK)
|
|
#if MAX_SHARED_MEMORY_PER_THREAD - MAX_LP * 2 >= 256
|
|
#define MAX_TOKENS_PER_THREAD 256
|
|
#define OUTPUT_BUFFER_LENGTH_PER_THREAD 256
|
|
#else
|
|
#define OUTPUT_BUFFER_LENGTH_PER_THREAD \
|
|
(MAX_SHARED_MEMORY_PER_THREAD - MAX_LP * 2)
|
|
#define MAX_TOKENS_PER_THREAD (OUTPUT_BUFFER_LENGTH_PER_THREAD)
|
|
#endif
|
|
#define PRECISION 16
|
|
|
|
extern int get_block_size(int);
|
|
|
|
template <typename T>
|
|
__inline__ __device__ T big_to_small(T value) {
|
|
return value;
|
|
}
|
|
|
|
template <>
|
|
__inline__ __device__ uint32_t big_to_small<uint32_t>(uint32_t value) {
|
|
return ((value & 0xFF000000U) >> 24) | ((value & 0x00FF0000U) >> 8) |
|
|
((value & 0x0000FF00U) << 8) | ((value & 0x000000FFU) << 24);
|
|
}
|
|
|
|
template <>
|
|
__inline__ __device__ uint8_t big_to_small<uint8_t>(uint8_t value) {
|
|
return value;
|
|
}
|
|
|
|
template <int BUFFER_BITS, typename BUFFER_TYPE>
|
|
__inline__ __device__ void read_next_bit(uint32_t& value,
|
|
BUFFER_TYPE& byte_buffer,
|
|
int& bit_idx) {
|
|
value <<= 1;
|
|
value |= (byte_buffer >> (BUFFER_BITS - bit_idx)) & 1;
|
|
bit_idx += 1;
|
|
}
|
|
|
|
template <int BUFFER_BITS, typename BUFFER_TYPE>
|
|
__inline__ __device__ void check_and_update_byte_buffer(
|
|
BUFFER_TYPE& byte_buffer, int& bit_idx, int& byte_buffer_offset,
|
|
uint8_t* bytestream) {
|
|
if (bit_idx == BUFFER_BITS + 1) {
|
|
bit_idx = 1;
|
|
byte_buffer_offset++;
|
|
byte_buffer = big_to_small<BUFFER_TYPE>(
|
|
((BUFFER_TYPE*)bytestream)[byte_buffer_offset]);
|
|
}
|
|
}
|
|
|
|
template <int BLOCK_SIZE>
|
|
__inline__ __device__ uint16_t binsearch(const uint16_t* cdf_shared,
|
|
uint16_t target, uint8_t max_sym,
|
|
const int tid) {
|
|
uint16_t left = 0;
|
|
uint16_t right = max_sym + 1; // len(cdf) == max_sym + 2
|
|
|
|
while (left + 1 < right) { // ?
|
|
const auto m = static_cast<uint16_t>((left + right) / 2);
|
|
const auto offset = m * BLOCK_SIZE + tid;
|
|
const auto v = cdf_shared[offset];
|
|
if (v < target) {
|
|
left = m;
|
|
} else if (v > target) {
|
|
right = m;
|
|
} else {
|
|
return m;
|
|
}
|
|
}
|
|
return left;
|
|
}
|
|
|
|
// BLOCK_SIZE SHOULD ALWAYS BE THE SAME AS blockDim.x
|
|
template <int BLOCK_SIZE, typename CDF_ACC_T, typename BS_ACC_T,
|
|
typename LEN_ACC_T, typename OUT_ACC_T>
|
|
__global__ void decode_with_accessor_kernel(CDF_ACC_T cdf, BS_ACC_T bytestreams,
|
|
LEN_ACC_T lengths, OUT_ACC_T output,
|
|
int32_t lp, int32_t ntokens) {
|
|
// The shared memory will be split to 3 parts:
|
|
// 1. The CDF tensor, with shape [MAX_LP, BLOCK_SIZE)] (only used [LP,
|
|
// BLOCK_SIZE] part)
|
|
// 2. The bytestream buffer, with shape [BLOCK_SIZE,
|
|
// OUTPUT_BUFFER_LENGTH_PER_THREAD] uint8s
|
|
// 3. The lengths buffer, with shape [BLOCK_SIZE] int32s
|
|
__shared__ __align__(4) uint16_t cdf_shared[MAX_LP][BLOCK_SIZE];
|
|
__shared__ __align__(4)
|
|
uint8_t bytestreams_shared[BLOCK_SIZE][OUTPUT_BUFFER_LENGTH_PER_THREAD];
|
|
int32_t* lengths_shared = (int32_t*)&cdf_shared[0][0];
|
|
|
|
const int layer_id = blockIdx.x;
|
|
const int global_channel_offset = blockIdx.y * BLOCK_SIZE;
|
|
const int local_channel_id = threadIdx.x;
|
|
const int global_channel_id = global_channel_offset + local_channel_id;
|
|
const int max_symbol = lp - 2;
|
|
|
|
// copy lengths[layer_id,
|
|
// global_channel_offset:global_channel_offset+BLOCK_SIZE] to shared memory
|
|
for (int i = threadIdx.x; i < BLOCK_SIZE; i += BLOCK_SIZE) {
|
|
lengths_shared[i] = lengths[layer_id][global_channel_offset + i];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// copy bytestreams[layer_id,
|
|
// global_channel_offset:global_channel_offset+BLOCK_SIZE, :] to shared
|
|
// memory, do this channel by channel
|
|
for (int i = 0; i < BLOCK_SIZE; i++) {
|
|
const int channel_id = global_channel_offset + i;
|
|
const int length = lengths_shared[i]; // shared memory broadcast
|
|
// TODO: optimized this by a packed-32bits read instead of 8bits read
|
|
for (int j = threadIdx.x; j < OUTPUT_BUFFER_LENGTH_PER_THREAD;
|
|
j += BLOCK_SIZE) {
|
|
const uint8_t value =
|
|
j < length ? bytestreams[layer_id][channel_id][j] : 0;
|
|
bytestreams_shared[i][j] = value;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// copy CDF[layer_id, global_channel_offset:global_channel_offset+BLOCK_SIZE,
|
|
// :] to shared memory
|
|
const int cdf_size = lp * BLOCK_SIZE;
|
|
for (int i = threadIdx.x; i < cdf_size; i += BLOCK_SIZE) {
|
|
const int cid = i / lp;
|
|
const int lid = i % lp;
|
|
cdf_shared[lid][cid] = cdf[layer_id][global_channel_offset + cid][lid];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// decode the bytestreams
|
|
uint32_t low = 0;
|
|
uint32_t high = 0xFFFFFFFFU;
|
|
uint32_t value = 0;
|
|
const uint32_t c_count = 0x10000U;
|
|
const int precision = 16;
|
|
|
|
uint8_t byte_buffer = 0;
|
|
int bit_idx = 1; // next bit to read: (byte_buffer >> (8 - bit_idx)) & 1
|
|
int byte_buffer_offset =
|
|
sizeof(value) / sizeof(byte_buffer); // where to read the next byte
|
|
|
|
// Get the initial value and byte buffer
|
|
value = big_to_small<uint32_t>(
|
|
((uint32_t*)bytestreams_shared[local_channel_id])[0]);
|
|
// byte_buffer = ((uint32_t
|
|
// *)bytestreams_shared[local_channel_id])[byte_buffer_offset];
|
|
byte_buffer = bytestreams_shared[local_channel_id][byte_buffer_offset];
|
|
|
|
for (int i = 0; i < ntokens; ++i) {
|
|
const uint64_t span =
|
|
static_cast<uint64_t>(high) - static_cast<uint64_t>(low) + 1;
|
|
// always < 0x10000 ???
|
|
const uint16_t count =
|
|
((static_cast<uint64_t>(value) - static_cast<uint64_t>(low) + 1) *
|
|
c_count -
|
|
1) /
|
|
span;
|
|
|
|
// TODO: implement the binsearch here!
|
|
auto sym_i = binsearch<BLOCK_SIZE>(&cdf_shared[0][0], count, max_symbol,
|
|
local_channel_id);
|
|
|
|
output[layer_id][i][global_channel_id] = sym_i;
|
|
|
|
if (i == ntokens - 1) {
|
|
break;
|
|
}
|
|
|
|
const uint32_t c_low = cdf_shared[sym_i][local_channel_id];
|
|
const uint32_t c_high = sym_i == max_symbol
|
|
? 0x10000U
|
|
: cdf_shared[sym_i + 1][local_channel_id];
|
|
|
|
high = (low - 1) + ((span * static_cast<uint64_t>(c_high)) >> precision);
|
|
low = (low) + ((span * static_cast<uint64_t>(c_low)) >> precision);
|
|
|
|
while (true) {
|
|
if (low >= 0x80000000U || high < 0x80000000U) {
|
|
low <<= 1;
|
|
high <<= 1;
|
|
high |= 1;
|
|
read_next_bit<8>(value, byte_buffer, bit_idx);
|
|
check_and_update_byte_buffer<8, uint8_t>(
|
|
byte_buffer, bit_idx, byte_buffer_offset,
|
|
bytestreams_shared[local_channel_id]);
|
|
} else if (low >= 0x40000000U && high < 0xC0000000U) {
|
|
low <<= 1;
|
|
low &= 0x7FFFFFFFU; // make MSB 0
|
|
high <<= 1;
|
|
high |= 0x80000001U; // add 1 at the end, retain MSB = 1
|
|
value -= 0x40000000U;
|
|
read_next_bit<8>(value, byte_buffer, bit_idx);
|
|
check_and_update_byte_buffer<8, uint8_t>(
|
|
byte_buffer, bit_idx, byte_buffer_offset,
|
|
bytestreams_shared[local_channel_id]);
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// BLOCK_SIZE SHOULD ALWAYS BE THE SAME AS blockDim.x
|
|
template <int BLOCK_SIZE, typename CDF_ACC_T, typename BS_ACC_T,
|
|
typename LEN_ACC_T, typename OUT_ACC_T>
|
|
__global__ void decode_prefix_with_accessor_kernel(CDF_ACC_T cdf,
|
|
BS_ACC_T bytestreams,
|
|
LEN_ACC_T lengths_prefix,
|
|
OUT_ACC_T output, int32_t lp,
|
|
int32_t ntokens) {
|
|
// The shared memory will be split to 3 parts:
|
|
// 1. The CDF tensor, with shape [MAX_LP, BLOCK_SIZE)] (only used [LP,
|
|
// BLOCK_SIZE] part)
|
|
// 2. The bytestream buffer, with shape [BLOCK_SIZE,
|
|
// OUTPUT_BUFFER_LENGTH_PER_THREAD] uint8s
|
|
// 3. The lengths buffer, with shape [BLOCK_SIZE] int32s
|
|
__shared__ __align__(4) uint16_t cdf_shared[MAX_LP][BLOCK_SIZE];
|
|
__shared__ __align__(4)
|
|
uint8_t bytestreams_shared[BLOCK_SIZE][OUTPUT_BUFFER_LENGTH_PER_THREAD];
|
|
int32_t* sum_lengths_shared = (int32_t*)&cdf_shared[0][0];
|
|
|
|
const int layer_id = blockIdx.x;
|
|
const int global_channel_offset = blockIdx.y * BLOCK_SIZE;
|
|
const int local_channel_id = threadIdx.x;
|
|
const int global_channel_id = global_channel_offset + local_channel_id;
|
|
const int max_symbol = lp - 2;
|
|
const int nchannels = gridDim.y * BLOCK_SIZE;
|
|
|
|
// copy lengths[layer_id,
|
|
// global_channel_offset:global_channel_offset+BLOCK_SIZE] to shared memory
|
|
for (int i = threadIdx.x; i < BLOCK_SIZE + 1; i += BLOCK_SIZE) {
|
|
int gid = layer_id * nchannels + global_channel_offset + i - 1;
|
|
sum_lengths_shared[i] =
|
|
gid >= 0 ? lengths_prefix[gid / nchannels][gid % nchannels] : 0;
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// copy bytestreams[layer_id,
|
|
// global_channel_offset:global_channel_offset+BLOCK_SIZE, :] to shared
|
|
// memory, do this channel by channel
|
|
for (int i = 0; i < BLOCK_SIZE; i++) {
|
|
[[maybe_unused]] const int channel_id = global_channel_offset + i;
|
|
const int start_offset = sum_lengths_shared[i];
|
|
const int end_offset = sum_lengths_shared[i + 1];
|
|
const int length = end_offset - start_offset;
|
|
// TODO: optimized this by a packed-32bits read instead of 8bits read
|
|
for (int j = threadIdx.x; j < OUTPUT_BUFFER_LENGTH_PER_THREAD;
|
|
j += BLOCK_SIZE) {
|
|
const uint8_t value = j < length ? bytestreams[start_offset + j] : 0;
|
|
bytestreams_shared[i][j] = value;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// copy CDF[layer_id, global_channel_offset:global_channel_offset+BLOCK_SIZE,
|
|
// :] to shared memory
|
|
const int cdf_size = lp * BLOCK_SIZE;
|
|
for (int i = threadIdx.x; i < cdf_size; i += BLOCK_SIZE) {
|
|
const int cid = i / lp;
|
|
const int lid = i % lp;
|
|
cdf_shared[lid][cid] = cdf[layer_id][global_channel_offset + cid][lid];
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// decode the bytestreams
|
|
uint32_t low = 0;
|
|
uint32_t high = 0xFFFFFFFFU;
|
|
uint32_t value = 0;
|
|
const uint32_t c_count = 0x10000U;
|
|
const int precision = 16;
|
|
|
|
uint8_t byte_buffer = 0;
|
|
int bit_idx = 1; // next bit to read: (byte_buffer >> (8 - bit_idx)) & 1
|
|
int byte_buffer_offset =
|
|
sizeof(value) / sizeof(byte_buffer); // where to read the next byte
|
|
|
|
// Get the initial value and byte buffer
|
|
value = big_to_small<uint32_t>(
|
|
((uint32_t*)bytestreams_shared[local_channel_id])[0]);
|
|
// byte_buffer = ((uint32_t
|
|
// *)bytestreams_shared[local_channel_id])[byte_buffer_offset];
|
|
byte_buffer = bytestreams_shared[local_channel_id][byte_buffer_offset];
|
|
|
|
for (int i = 0; i < ntokens; ++i) {
|
|
const uint64_t span =
|
|
static_cast<uint64_t>(high) - static_cast<uint64_t>(low) + 1;
|
|
// always < 0x10000 ???
|
|
const uint16_t count =
|
|
((static_cast<uint64_t>(value) - static_cast<uint64_t>(low) + 1) *
|
|
c_count -
|
|
1) /
|
|
span;
|
|
|
|
// TODO: implement the binsearch here!
|
|
auto sym_i = binsearch<BLOCK_SIZE>(&cdf_shared[0][0], count, max_symbol,
|
|
local_channel_id);
|
|
|
|
output[layer_id][i][global_channel_id] = sym_i;
|
|
|
|
if (i == ntokens - 1) {
|
|
break;
|
|
}
|
|
|
|
const uint32_t c_low = cdf_shared[sym_i][local_channel_id];
|
|
const uint32_t c_high = sym_i == max_symbol
|
|
? 0x10000U
|
|
: cdf_shared[sym_i + 1][local_channel_id];
|
|
|
|
high = (low - 1) + ((span * static_cast<uint64_t>(c_high)) >> precision);
|
|
low = (low) + ((span * static_cast<uint64_t>(c_low)) >> precision);
|
|
|
|
while (true) {
|
|
if (low >= 0x80000000U || high < 0x80000000U) {
|
|
low <<= 1;
|
|
high <<= 1;
|
|
high |= 1;
|
|
read_next_bit<8>(value, byte_buffer, bit_idx);
|
|
check_and_update_byte_buffer<8, uint8_t>(
|
|
byte_buffer, bit_idx, byte_buffer_offset,
|
|
bytestreams_shared[local_channel_id]);
|
|
} else if (low >= 0x40000000U && high < 0xC0000000U) {
|
|
low <<= 1;
|
|
low &= 0x7FFFFFFFU; // make MSB 0
|
|
high <<= 1;
|
|
high |= 0x80000001U; // add 1 at the end, retain MSB = 1
|
|
value -= 0x40000000U;
|
|
read_next_bit<8>(value, byte_buffer, bit_idx);
|
|
check_and_update_byte_buffer<8, uint8_t>(
|
|
byte_buffer, bit_idx, byte_buffer_offset,
|
|
bytestreams_shared[local_channel_id]);
|
|
} else {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief CUDA kernel to decode a compressed bytestream using the given CDF.
|
|
*
|
|
* @param cdf the int16 CDF tensor, with shape [nlayers, nchannels, LP], should
|
|
* be on GPU
|
|
* @param bytestreams The uint8 bytestreams tensor, with shape [nlayers,
|
|
* nchannels, OUTPUT_BUFFER_LENGTH_PER_THREAD], should be on GPU
|
|
* @param lengths The int32 lengths tensor, with shape [nlayers, nchannels],
|
|
* should be on GPU
|
|
* @param output The uint8 output tensor, with shape nlayers, ntokens,
|
|
* nchannels], should be on GPU.
|
|
*/
|
|
void decode_cuda_new(const at::Tensor& cdf, const at::Tensor& bytestreams,
|
|
const at::Tensor& lengths, at::Tensor& output) {
|
|
TORCH_CHECK(cdf.is_cuda(), "CDF should be on GPU");
|
|
TORCH_CHECK(bytestreams.is_cuda(), "Bytestreams should be on GPU");
|
|
TORCH_CHECK(lengths.is_cuda(), "Lengths should be on GPU");
|
|
TORCH_CHECK(output.is_cuda(), "Output should be on GPU");
|
|
|
|
const auto cdf_shape = cdf.sizes();
|
|
const auto bs_shape = bytestreams.sizes();
|
|
const auto lengths_shape = lengths.sizes();
|
|
const auto output_shape = output.sizes();
|
|
TORCH_CHECK(cdf_shape[0] == bs_shape[0],
|
|
"CDF and bytestreams should have the same number of layers");
|
|
TORCH_CHECK(cdf_shape[1] == bs_shape[1],
|
|
"CDF and bytestreams should have the same number of channels");
|
|
TORCH_CHECK(cdf_shape[0] == lengths_shape[0],
|
|
"CDF and lengths should have the same number of layers");
|
|
TORCH_CHECK(cdf_shape[1] == lengths_shape[1],
|
|
"CDF and lengths should have the same number of channels");
|
|
TORCH_CHECK(cdf_shape[0] == output_shape[0],
|
|
"CDF and output should have the same number of layers");
|
|
TORCH_CHECK(cdf_shape[1] == output_shape[2],
|
|
"CDF and output should have the same number of channels");
|
|
|
|
const int nlayers = cdf_shape[0];
|
|
const int nchannels = cdf_shape[1];
|
|
const int ntokens = output_shape[1];
|
|
const int lp = cdf_shape[2];
|
|
const int block_size = get_block_size(nchannels);
|
|
TORCH_CHECK(ntokens <= MAX_TOKENS_PER_THREAD,
|
|
"Number of tokens should be less than or equal to",
|
|
MAX_TOKENS_PER_THREAD);
|
|
TORCH_CHECK(nchannels % block_size == 0,
|
|
"Number of channels should be divisible by block size");
|
|
TORCH_CHECK(lp <= MAX_LP, "CDF should have at most", MAX_LP, "Lps");
|
|
|
|
dim3 block_dim(block_size, 1, 1);
|
|
dim3 grid_dim(nlayers, nchannels / block_size, 1);
|
|
|
|
auto cdf_accessor =
|
|
cdf.packed_accessor32<int16_t, 3, torch::RestrictPtrTraits>();
|
|
auto bytestreams_accessor =
|
|
bytestreams.packed_accessor32<uint8_t, 3, torch::RestrictPtrTraits>();
|
|
auto lengths_accessor =
|
|
lengths.packed_accessor32<int32_t, 2, torch::RestrictPtrTraits>();
|
|
auto output_accessor =
|
|
output.packed_accessor32<uint8_t, 3, torch::RestrictPtrTraits>();
|
|
|
|
#ifndef LAUNCH_DECODE_KERNEL
|
|
#define LAUNCH_DECODE_KERNEL(block_size) \
|
|
decode_with_accessor_kernel<block_size><<<grid_dim, block_dim>>>( \
|
|
cdf_accessor, bytestreams_accessor, lengths_accessor, output_accessor, \
|
|
lp, ntokens)
|
|
#endif
|
|
|
|
switch (block_size) {
|
|
case 1:
|
|
LAUNCH_DECODE_KERNEL(1);
|
|
break;
|
|
case 2:
|
|
LAUNCH_DECODE_KERNEL(2);
|
|
break;
|
|
case 4:
|
|
LAUNCH_DECODE_KERNEL(4);
|
|
break;
|
|
case 8:
|
|
LAUNCH_DECODE_KERNEL(8);
|
|
break;
|
|
case 16:
|
|
LAUNCH_DECODE_KERNEL(16);
|
|
break;
|
|
case 32:
|
|
LAUNCH_DECODE_KERNEL(32);
|
|
break;
|
|
case 64:
|
|
LAUNCH_DECODE_KERNEL(64);
|
|
break;
|
|
case 128:
|
|
LAUNCH_DECODE_KERNEL(128);
|
|
break;
|
|
default:
|
|
throw std::runtime_error("Unsupported block size");
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief CUDA kernel to decode a compressed bytestream using the given CDF.
|
|
*
|
|
* @param cdf the int16 CDF tensor, with shape [nlayers, nchannels, LP], should
|
|
* be on GPU
|
|
* @param bytestreams The 1-D uint8 bytestreams tensor containing [nlayers,
|
|
* nchannels] bytestreams, should be on GPU
|
|
* @param lengths_prefsum The int64 tensor containing the prefix sum of the
|
|
* lengths, with shape [nlayers, nchannels], should be on GPU
|
|
* @param output The uint8 output tensor, with shape nlayers, ntokens,
|
|
* nchannels], should be on GPU.
|
|
*/
|
|
void decode_cuda_prefsum(const at::Tensor& cdf, const at::Tensor& bytestreams,
|
|
const at::Tensor& lengths_prefsum,
|
|
at::Tensor& output) {
|
|
TORCH_CHECK(cdf.is_cuda(), "CDF should be on GPU");
|
|
TORCH_CHECK(bytestreams.is_cuda(), "Bytestreams should be on GPU");
|
|
TORCH_CHECK(lengths_prefsum.is_cuda(), "Lengths should be on GPU");
|
|
TORCH_CHECK(output.is_cuda(), "Output should be on GPU");
|
|
|
|
const auto cdf_shape = cdf.sizes();
|
|
const auto lengths_shape = lengths_prefsum.sizes();
|
|
const auto output_shape = output.sizes();
|
|
TORCH_CHECK(cdf_shape[0] == lengths_shape[0],
|
|
"CDF and lengths should have the same number of layers");
|
|
TORCH_CHECK(cdf_shape[1] == lengths_shape[1],
|
|
"CDF and lengths should have the same number of channels");
|
|
TORCH_CHECK(cdf_shape[0] == output_shape[0],
|
|
"CDF and output should have the same number of layers");
|
|
TORCH_CHECK(cdf_shape[1] == output_shape[2],
|
|
"CDF and output should have the same number of channels");
|
|
|
|
const int nlayers = cdf_shape[0];
|
|
const int nchannels = cdf_shape[1];
|
|
const int ntokens = output_shape[1];
|
|
const int lp = cdf_shape[2];
|
|
const int block_size = get_block_size(nchannels);
|
|
TORCH_CHECK(ntokens <= MAX_TOKENS_PER_THREAD,
|
|
"Number of tokens should be less than or equal to",
|
|
MAX_TOKENS_PER_THREAD);
|
|
TORCH_CHECK(nchannels % block_size == 0,
|
|
"Number of channels should be divisible by block size");
|
|
TORCH_CHECK(lp <= MAX_LP, "CDF should have at most", MAX_LP, "Lps");
|
|
|
|
dim3 block_dim(block_size, 1, 1);
|
|
dim3 grid_dim(nlayers, nchannels / block_size, 1);
|
|
|
|
auto cdf_accessor =
|
|
cdf.packed_accessor32<int16_t, 3, torch::RestrictPtrTraits>();
|
|
auto bytestreams_accessor =
|
|
bytestreams.packed_accessor32<uint8_t, 1, torch::RestrictPtrTraits>();
|
|
auto lengths_accessor =
|
|
lengths_prefsum.packed_accessor32<int64_t, 2, torch::RestrictPtrTraits>();
|
|
auto output_accessor =
|
|
output.packed_accessor32<uint8_t, 3, torch::RestrictPtrTraits>();
|
|
|
|
#ifndef LAUNCH_DECODE_PREFIX_KERNEL
|
|
#define LAUNCH_DECODE_PREFIX_KERNEL(block_size) \
|
|
decode_prefix_with_accessor_kernel<block_size><<<grid_dim, block_dim>>>( \
|
|
cdf_accessor, bytestreams_accessor, lengths_accessor, output_accessor, \
|
|
lp, ntokens)
|
|
#endif
|
|
|
|
switch (block_size) {
|
|
case 1:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(1);
|
|
break;
|
|
case 2:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(2);
|
|
break;
|
|
case 4:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(4);
|
|
break;
|
|
case 8:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(8);
|
|
break;
|
|
case 16:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(16);
|
|
break;
|
|
case 32:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(32);
|
|
break;
|
|
case 64:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(64);
|
|
break;
|
|
case 128:
|
|
LAUNCH_DECODE_PREFIX_KERNEL(128);
|
|
break;
|
|
default:
|
|
throw std::runtime_error("Unsupported block size");
|
|
}
|
|
}
|