X-Git-Url: https://git.sesse.net/?p=narabu;a=blobdiff_plain;f=rans.shader;h=3e19c83eae83e4f31f6c0c47bd4f0bc4263a7cc5;hp=39efc52914bb579bb01208a870df2ac4301352cb;hb=0f84eb2a5fa77dbc20f7a073f83bdb674ebb7b75;hpb=fedc6308264085999789854152c9577ddfe1bc6e

diff --git a/rans.shader b/rans.shader
index 39efc52..3e19c83 100644
--- a/rans.shader
+++ b/rans.shader
@@ -1,5 +1,4 @@
 #version 440
-#extension GL_NV_gpu_shader5 : enable
 
 layout(local_size_x = 1) in;
 
@@ -7,7 +6,7 @@ const uint prob_bits = 12;
 const uint prob_scale = 1 << prob_bits;
 const uint RANS_BYTE_L = (1u << 23);
 const uint BLOCKS_PER_STREAM = 320;
-const uint STREAM_BUF_SIZE = 1024;  // 1 kB per stream ought to be enough for everyone :-)
+const uint STREAM_BUF_SIZE = 256;  // In uint32s. 1 kB per stream ought to be enough for everyone :-)
 const uint NUM_SYMS = 256;
 const uint ESCAPE_LIMIT = NUM_SYMS - 1;
 
@@ -26,12 +25,12 @@ const uint luma_mapping[8] = {
 
 layout(std430, binding = 10) buffer outputBuf
 {
-	uint8_t rans_output[];
+	uint rans_output[];
 };
 
 layout(std430, binding = 11) buffer outputBuf2
 {
-	uint rans_start_offset[];
+	uint rans_bytes_written[];
 };
 
 layout(std140, binding = 13) uniform DistBlock
@@ -41,11 +40,14 @@ layout(std140, binding = 13) uniform DistBlock
 };
 
 struct RansEncoder {
-	uint stream_num;   // const
-	uint lut_base;     // const
-	uint sign_bias;    // const
+	uint stream_num;         // const
+	uint lut_base;           // const
+	uint sign_bias;          // const
+	uint rans_start_offset;  // const
 	uint rans_offset;
 	uint rans;
+	uint bit_buffer;
+	uint bytes_in_buffer;
 };
 
 layout(r16ui) uniform restrict readonly uimage2D dc_ac7_tex;
@@ -56,7 +58,11 @@ layout(r8i) uniform restrict readonly iimage2D ac4_tex;
 
 void RansPutByte(uint x, inout RansEncoder enc)
 {
-	rans_output[--enc.rans_offset] = uint8_t(x);
+	enc.bit_buffer = (enc.bit_buffer << 8) | x;
+	if (++enc.bytes_in_buffer == 4) {
+		rans_output[--enc.rans_offset] = enc.bit_buffer;
+		enc.bytes_in_buffer = 0;
+	}
 }
 
 void RansEncInit(uint streamgroup_num, uint coeff_row, uint coeff_col, uint dist_num, out RansEncoder enc)
@@ -65,7 +71,10 @@ void RansEncInit(uint streamgroup_num, uint coeff_row, uint coeff_col, uint dist
 	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_start_offset = enc.rans_offset;
 	enc.rans = RANS_BYTE_L;
+	enc.bit_buffer = 0;
+	enc.bytes_in_buffer = 0;
 }
 
 void RansEncRenorm(inout RansEncoder enc, uint freq, uint prob_bits)
@@ -102,12 +111,21 @@ void RansEncPutSymbol(inout RansEncoder enc, uvec4 sym)
 	enc.rans += bias + (q >> rcp_shift) * cmpl_freq;
 }
 
-void RansEncFlush(inout RansEncoder enc)
+uint RansEncFlush(inout RansEncoder enc)
 {
 	RansPutByte(enc.rans >> 24, enc);
 	RansPutByte(enc.rans >> 16, enc);
 	RansPutByte(enc.rans >> 8, enc);
 	RansPutByte(enc.rans >> 0, enc);
+
+	uint num_bytes_written = (enc.rans_start_offset - enc.rans_offset) * 4 + enc.bytes_in_buffer;
+
+	// Make sure there's nothing left in the buffer.
+	RansPutByte(0, enc);
+	RansPutByte(0, enc);
+	RansPutByte(0, enc);
+
+	return num_bytes_written;
 }
 
 int sign_extend(uint coeff, uint bits)
@@ -137,8 +155,8 @@ void encode_coeff(int signed_k, inout RansEncoder enc)
 
 void encode_end(inout RansEncoder enc)
 {
-	RansEncFlush(enc);
-	rans_start_offset[enc.stream_num] = enc.rans_offset;
+	uint bytes_written = RansEncFlush(enc);
+	rans_bytes_written[enc.stream_num] = bytes_written;
 }
 
 void encode_9_7(uint streamgroup_num, uint coeff_row, layout(r16ui) restrict readonly uimage2D tex, uint col1, uint col2, uint dist1, uint dist2)