From f3eac72e679d3e0ec7ae1d4484736cd552c344dd Mon Sep 17 00:00:00 2001 From: "Steinar H. Gunderson" Date: Tue, 17 Oct 2017 22:05:45 +0200 Subject: [PATCH] Kill the division in the rANS GPU encoder. --- narabu-encoder.cpp | 22 +++++++++++++------ rans.shader | 54 ++++++++++++++++++++++++++++++---------------- tally.shader | 35 ++++++++++++++++++++++++++++-- 3 files changed, 85 insertions(+), 26 deletions(-) diff --git a/narabu-encoder.cpp b/narabu-encoder.cpp index d13f044..ca28528 100644 --- a/narabu-encoder.cpp +++ b/narabu-encoder.cpp @@ -43,7 +43,11 @@ unsigned char pix_cr[(WIDTH/2) * HEIGHT]; struct RansDistSSBO { unsigned dist[4 * 256]; - std::pair ransdist[4 * 256]; + unsigned ransfreq[4 * 256]; + struct { + uint32_t x_max, rcp_freq, bias, rcp_shift_and_cmpl_freq; + } ransdist[4 * 256]; + unsigned sign_biases[4]; }; using namespace std; @@ -194,7 +198,7 @@ int main(int argc, char **argv) GLuint ssbo; glGenBuffers(1, &ssbo); glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo); - glNamedBufferStorage(ssbo, 256 * 16 * sizeof(uint32_t), nullptr, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT); + glNamedBufferStorage(ssbo, sizeof(RansDistSSBO), nullptr, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT); check_error(); // SSBOs for the rANS output (data and offsets). @@ -300,7 +304,7 @@ int main(int argc, char **argv) steady_clock::time_point start = steady_clock::now(); unsigned num_iterations = 100; for (unsigned i = 0; i < num_iterations; ++i) { - glClearNamedBufferSubData(ssbo, GL_R8, 0, 256 * 16 * sizeof(uint32_t), GL_RED, GL_UNSIGNED_BYTE, nullptr); + glClearNamedBufferSubData(ssbo, GL_R8, 0, sizeof(RansDistSSBO), GL_RED, GL_UNSIGNED_BYTE, nullptr); glUseProgram(glsl_program_num); glDispatchCompute(WIDTH_BLOCKS / 16, HEIGHT_BLOCKS, 1); glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT); @@ -336,14 +340,20 @@ int main(int argc, char **argv) } // Write out the distributions. - const RansDistSSBO *rans_dist = (const RansDistSSBO *)glMapNamedBufferRange(ssbo, 0, 256 * 16 * sizeof(uint32_t), GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT); + const RansDistSSBO *rans_dist = (const RansDistSSBO *)glMapNamedBufferRange(ssbo, 0, sizeof(RansDistSSBO), GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT); for (unsigned r = 0; r < 2; ++r) { // Hack to write fake chroma tables. // TODO: rather gamma-k or something for (unsigned i = 0; i < 4; ++i) { printf("writing table %d\n", i); for (unsigned j = 0; j < NUM_SYMS; ++j) { - printf("%d,%d: start=%d freq=%d\n", i, j, rans_dist->ransdist[i * 256 + j].first, rans_dist->ransdist[i * 256 + j].second); - write_varint(rans_dist->ransdist[i * 256 + j].second, codedfp); + printf("%d,%d: freq=%d x_max=%d, rcp_freq=%08x, bias=%d, rcp_shift=%d, cmpl_freq=%d\n", + i, j, rans_dist->ransfreq[i * 256 + j], + rans_dist->ransdist[i * 256 + j].x_max, + rans_dist->ransdist[i * 256 + j].rcp_freq, + rans_dist->ransdist[i * 256 + j].bias, + rans_dist->ransdist[i * 256 + j].rcp_shift_and_cmpl_freq & 0xffff, + rans_dist->ransdist[i * 256 + j].rcp_shift_and_cmpl_freq >> 16); + write_varint(rans_dist->ransfreq[i * 256 + j], codedfp); } } } diff --git a/rans.shader b/rans.shader index 5b13d7e..b677415 100644 --- a/rans.shader +++ b/rans.shader @@ -27,7 +27,9 @@ const uint luma_mapping[8] = { layout(std430, binding = 9) buffer layoutName { uint dist[4 * 256]; - uvec2 ransdist[4 * 256]; + uint ransfreq[4 * 256]; + uvec4 ransdist[4 * 256]; + uint sign_biases[4]; }; layout(std430, binding = 10) buffer outputBuf @@ -43,6 +45,7 @@ layout(std430, binding = 11) buffer outputBuf2 struct RansEncoder { uint stream_num; // const uint lut_base; // const + uint sign_bias; // const uint rans_offset; uint rans; }; @@ -57,6 +60,7 @@ void RansEncInit(uint streamgroup_num, uint coeff_row, uint coeff_col, uint dist { enc.stream_num = streamgroup_num * 64 + coeff_row * 8 + coeff_col; enc.lut_base = dist_num * 256; + enc.sign_bias = sign_biases[dist_num]; enc.rans_offset = enc.stream_num * STREAM_BUF_SIZE + STREAM_BUF_SIZE; // Starts at the end. enc.rans = RANS_BYTE_L; } @@ -64,11 +68,9 @@ void RansEncInit(uint streamgroup_num, uint coeff_row, uint coeff_col, uint dist void RansEncRenorm(inout uint rans, inout uint rans_offset, uint freq, uint prob_bits) { uint x_max = ((RANS_BYTE_L >> prob_bits) << 8) * freq; // this turns into a shift. - if (rans >= x_max) { - do { - rans_output[--rans_offset] = uint8_t(rans & 0xff); - rans >>= 8; - } while (rans >= x_max); + while (rans >= x_max) { + rans_output[--rans_offset] = uint8_t(rans & 0xffu); + rans >>= 8; } } @@ -78,6 +80,27 @@ void RansEncPut(inout uint rans, inout uint rans_offset, uint start, uint freq, rans = ((rans / freq) << prob_bits) + (rans % freq) + start; } +void RansEncPutSymbol(inout uint rans, inout uint rans_offset, uvec4 sym) +{ + uint x_max = sym.x; + uint rcp_freq = sym.y; + uint bias = sym.z; + uint rcp_shift = (sym.w & 0xffffu); + uint cmpl_freq = (sym.w >> 16); + + // renormalize + if (rans >= x_max) { + do { + rans_output[--rans_offset] = uint8_t(rans & 0xffu); + rans >>= 8; + } while (rans >= x_max); + } + + uint q, unused; + umulExtended(rans, rcp_freq, q, unused); + rans += bias + (q >> rcp_shift) * cmpl_freq; +} + void RansEncFlush(uint rans, inout uint rans_offset) { rans_offset -= 4; @@ -92,7 +115,7 @@ int sign_extend(uint coeff, uint bits) return int(coeff << (32 - bits)) >> (32 - bits); } -void encode_coeff(int signed_k, uint sign_bias, inout RansEncoder enc) +void encode_coeff(int signed_k, inout RansEncoder enc) { uint k = abs(signed_k); @@ -104,11 +127,11 @@ void encode_coeff(int signed_k, uint sign_bias, inout RansEncoder enc) k = ESCAPE_LIMIT; } - uvec2 sym = ransdist[enc.lut_base + ((k - 1) & (NUM_SYMS - 1))]; - RansEncPut(enc.rans, enc.rans_offset, sym.x, sym.y, prob_bits + 1); + uvec4 sym = ransdist[enc.lut_base + ((k - 1) & (NUM_SYMS - 1))]; + RansEncPutSymbol(enc.rans, enc.rans_offset, sym); if (signed_k < 0) { - enc.rans += sign_bias; + enc.rans += enc.sign_bias; } } @@ -124,17 +147,14 @@ void encode_9_7(uint streamgroup_num, uint coeff_row, layout(r16ui) restrict rea RansEncInit(streamgroup_num, coeff_row, col1, dist1, enc1); RansEncInit(streamgroup_num, coeff_row, col2, dist2, enc2); - uint sign_bias1 = ransdist[enc1.lut_base + 255].x + ransdist[enc1.lut_base + 255].y; - uint sign_bias2 = ransdist[enc2.lut_base + 255].x + ransdist[enc2.lut_base + 255].y; - for (uint subblock_idx = 0; subblock_idx < BLOCKS_PER_STREAM; ++subblock_idx) { // TODO: Use SSBOs instead of a texture? uint x = (streamgroup_num * BLOCKS_PER_STREAM + subblock_idx) % 160; uint y = (streamgroup_num * BLOCKS_PER_STREAM + subblock_idx) / 160; uint f = imageLoad(tex, ivec2(x, y * 8 + coeff_row)).x; - encode_coeff(sign_extend(f & 0x1ffu, 9), sign_bias1, enc1); - encode_coeff(sign_extend(f >> 9, 7), sign_bias2, enc2); + encode_coeff(sign_extend(f & 0x1ffu, 9), enc1); + encode_coeff(sign_extend(f >> 9, 7), enc2); } encode_end(enc1); @@ -146,15 +166,13 @@ void encode_8(uint streamgroup_num, uint coeff_row, layout(r8i) restrict readonl RansEncoder enc; RansEncInit(streamgroup_num, coeff_row, col, dist, enc); - uint sign_bias = ransdist[enc.lut_base + 255].x + ransdist[enc.lut_base + 255].y; - for (uint subblock_idx = 0; subblock_idx < BLOCKS_PER_STREAM; ++subblock_idx) { // TODO: Use SSBOs instead of a texture? uint x = (streamgroup_num * BLOCKS_PER_STREAM + subblock_idx) % 160; uint y = (streamgroup_num * BLOCKS_PER_STREAM + subblock_idx) / 160; int f = imageLoad(tex, ivec2(x, y * 8 + coeff_row)).x; - encode_coeff(f, sign_bias, enc); + encode_coeff(f, enc); } encode_end(enc); diff --git a/tally.shader b/tally.shader index 53f73af..351d3fb 100644 --- a/tally.shader +++ b/tally.shader @@ -1,4 +1,5 @@ #version 440 +#extension GL_NV_gpu_shader5 : enable // http://cbloomrants.blogspot.no/2014/02/02-11-14-understanding-ans-10.html @@ -7,11 +8,14 @@ layout(local_size_x = 256) in; layout(std430, binding = 9) buffer layoutName { uint dist[4 * 256]; - uvec2 ransdist[4 * 256]; + uint ransfreq[4 * 256]; + uvec4 ransdist[4 * 256]; + uint sign_biases[4]; }; const uint prob_bits = 12; const uint prob_scale = 1 << prob_bits; +const uint RANS_BYTE_L = (1u << 23); // FIXME: should come through a uniform. const uint sums[4] = { 57600, 115200, 302400, 446400 }; @@ -165,5 +169,32 @@ void main() memoryBarrierShared(); barrier(); } - ransdist[base + i] = uvec2(new_dist[i] - new_val, new_val); + + uint start = new_dist[i] - new_val; + uint freq = new_val; + + uint x_max = ((RANS_BYTE_L >> (prob_bits + 1)) << 8) * freq; + uint cmpl_freq = ((1 << (prob_bits + 1)) - freq); + uint rcp_freq, rcp_shift, bias; + if (freq < 2) { + rcp_freq = ~0u; + rcp_shift = 0; + bias = start + (1 << (prob_bits + 1)) - 1; + } else { + uint shift = 0; + while (freq > (1u << shift)) { + shift++; + } + + rcp_freq = uint(((uint64_t(1) << (shift + 31)) + freq-1) / freq); + rcp_shift = shift - 1; + bias = start; + } + + ransfreq[base + i] = freq; + ransdist[base + i] = uvec4(x_max, rcp_freq, bias, (cmpl_freq << 16) | rcp_shift); + + if (i == 255) { + sign_biases[gl_WorkGroupID.x] = new_dist[i]; + } } -- 2.39.2