git.sesse.net Git - narabu/blob - narabu-encoder.cpp

   1 #include <stdio.h>
   2 #include <stdint.h>
   3 #include <stdlib.h>
   4 #include <string.h>
   5 #include <assert.h>
   6 #include <math.h>
   7 #include <SDL2/SDL.h>
   8 #include <SDL2/SDL_error.h>
   9 #include <SDL2/SDL_video.h>
  10 #include <epoxy/gl.h>
  11
  12 #include <algorithm>
  13 #include <chrono>
  14 #include <memory>
  15 #include <numeric>
  16 #include <random>
  17 #include <vector>
  18 #include <unordered_map>
  19
  20 #include <movit/util.h>
  21
  22 #include "util.h"
  23
  24 #define WIDTH 1280
  25 #define HEIGHT 720
  26 #define WIDTH_BLOCKS (WIDTH/8)
  27 #define WIDTH_BLOCKS_CHROMA (WIDTH/16)
  28 #define HEIGHT_BLOCKS (HEIGHT/8)
  29 #define NUM_BLOCKS (WIDTH_BLOCKS * HEIGHT_BLOCKS)
  30 #define NUM_BLOCKS_CHROMA (WIDTH_BLOCKS_CHROMA * HEIGHT_BLOCKS)
  31
  32 #define NUM_SYMS 256
  33 #define ESCAPE_LIMIT (NUM_SYMS - 1)
  34 #define BLOCKS_PER_STREAM 320
  35
  36 static constexpr uint32_t prob_bits = 12;
  37 static constexpr uint32_t prob_scale = 1 << prob_bits;
  38
  39 unsigned char rgb[WIDTH * HEIGHT * 3];
  40 unsigned char pix_y[WIDTH * HEIGHT];
  41 unsigned char pix_cb[(WIDTH/2) * HEIGHT];
  42 unsigned char pix_cr[(WIDTH/2) * HEIGHT];
  43
  44 struct RansDistSSBO {
  45         unsigned dist[4 * 256];
  46         std::pair<unsigned, unsigned> ransdist[4 * 256];
  47 };
  48
  49 using namespace std;
  50 using namespace std::chrono;
  51
  52 void write_varint(int x, FILE *fp)
  53 {
  54         while (x >= 128) {
  55                 putc((x & 0x7f) | 0x80, fp);
  56                 x >>= 7;
  57         }
  58         putc(x, fp);
  59 }
  60
  61 void readpix(unsigned char *ptr, const char *filename)
  62 {
  63         FILE *fp = fopen(filename, "rb");
  64         if (fp == nullptr) {
  65                 perror(filename);
  66                 exit(1);
  67         }
  68
  69         fseek(fp, 0, SEEK_END);
  70         long len = ftell(fp);
  71         assert(len >= WIDTH * HEIGHT * 3);
  72         fseek(fp, len - WIDTH * HEIGHT * 3, SEEK_SET);
  73
  74         fread(ptr, 1, WIDTH * HEIGHT * 3, fp);
  75         fclose(fp);
  76 }
  77
  78 // Should be done on the GPU, of course, but irrelevant for the demonstration.
  79 void convert_ycbcr()
  80 {
  81         double coeff[3] = { 0.2126, 0.7152, 0.0722 };  // sum = 1.0
  82         double cb_fac = 1.0 / (coeff[0] + coeff[1] + 1.0f - coeff[2]);  // 0.539
  83         double cr_fac = 1.0 / (1.0f - coeff[0] + coeff[1] + coeff[2]);  // 0.635
  84
  85         unique_ptr<float[]> temp_cb(new float[WIDTH * HEIGHT]);
  86         unique_ptr<float[]> temp_cr(new float[WIDTH * HEIGHT]);
  87         for (unsigned yb = 0; yb < HEIGHT; ++yb) {
  88                 for (unsigned xb = 0; xb < WIDTH; ++xb) {
  89                         int r = rgb[((yb * WIDTH) + xb) * 3 + 0];
  90                         int g = rgb[((yb * WIDTH) + xb) * 3 + 1];
  91                         int b = rgb[((yb * WIDTH) + xb) * 3 + 2];
  92                         double y = std::min(std::max(coeff[0] * r + coeff[1] * g + coeff[2] * b, 0.0), 255.0);
  93                         double cb = (b - y) * cb_fac + 128.0;
  94                         double cr = (r - y) * cr_fac + 128.0;
  95                         pix_y[(yb * WIDTH) + xb] = lrint(y);
  96                         temp_cb[(yb * WIDTH) + xb] = cb;
  97                         temp_cr[(yb * WIDTH) + xb] = cr;
  98                 }
  99         }
 100
 101         // Simple 4:2:2 subsampling with left convention.
 102         for (unsigned yb = 0; yb < HEIGHT; ++yb) {
 103                 for (unsigned xb = 0; xb < WIDTH / 2; ++xb) {
 104                         int c0 = yb * WIDTH + std::max(int(xb) * 2 - 1, 0);
 105                         int c1 = yb * WIDTH + xb * 2;
 106                         int c2 = yb * WIDTH + xb * 2 + 1;
 107
 108                         double cb = 0.25 * temp_cb[c0] + 0.5 * temp_cb[c1] + 0.25 * temp_cb[c2];
 109                         double cr = 0.25 * temp_cr[c0] + 0.5 * temp_cr[c1] + 0.25 * temp_cr[c2];
 110                         cb = std::min(std::max(cb, 0.0), 255.0);
 111                         cr = std::min(std::max(cr, 0.0), 255.0);
 112                         pix_cb[(yb * WIDTH/2) + xb] = lrint(cb);
 113                         pix_cr[(yb * WIDTH/2) + xb] = lrint(cr);
 114                 }
 115         }
 116 }
 117
 118 int main(int argc, char **argv)
 119 {
 120         // Set up an OpenGL context using SDL.
 121         if (SDL_Init(SDL_INIT_VIDEO) == -1) {
 122                 fprintf(stderr, "SDL_Init failed: %s\n", SDL_GetError());
 123                 exit(1);
 124         }
 125         SDL_GL_SetAttribute(SDL_GL_DEPTH_SIZE, 0);
 126         SDL_GL_SetAttribute(SDL_GL_STENCIL_SIZE, 0);
 127         SDL_GL_SetAttribute(SDL_GL_DOUBLEBUFFER, 1);
 128         SDL_GL_SetAttribute(SDL_GL_CONTEXT_PROFILE_MASK, SDL_GL_CONTEXT_PROFILE_CORE);
 129         SDL_GL_SetAttribute(SDL_GL_CONTEXT_MAJOR_VERSION, 4);
 130         SDL_GL_SetAttribute(SDL_GL_CONTEXT_MINOR_VERSION, 5);
 131
 132         SDL_Window *window = SDL_CreateWindow("OpenGL window for unit test",
 133                 SDL_WINDOWPOS_UNDEFINED,
 134                 SDL_WINDOWPOS_UNDEFINED,
 135                 32, 32,
 136                 SDL_WINDOW_OPENGL);
 137         SDL_GLContext context = SDL_GL_CreateContext(window);
 138         assert(context != nullptr);
 139
 140         if (argc >= 2)
 141                 readpix(rgb, argv[1]);
 142         else
 143                 readpix(rgb, "color.pnm");
 144         convert_ycbcr();
 145
 146         // Compile the DCT shader.
 147         string shader_src = ::read_file("encoder.shader");
 148         GLuint shader_num = compile_shader(shader_src, GL_COMPUTE_SHADER);
 149         GLuint glsl_program_num = glCreateProgram();
 150         glAttachShader(glsl_program_num, shader_num);
 151         glLinkProgram(glsl_program_num);
 152
 153         GLint success;
 154         glGetProgramiv(glsl_program_num, GL_LINK_STATUS, &success);
 155         if (success == GL_FALSE) {
 156                 GLchar error_log[1024] = {0};
 157                 glGetProgramInfoLog(glsl_program_num, 1024, nullptr, error_log);
 158                 fprintf(stderr, "Error linking program: %s\n", error_log);
 159                 exit(1);
 160         }
 161
 162         // Compile the tally shader.
 163         shader_src = ::read_file("tally.shader");
 164         shader_num = compile_shader(shader_src, GL_COMPUTE_SHADER);
 165         GLuint glsl_tally_program_num = glCreateProgram();
 166         glAttachShader(glsl_tally_program_num, shader_num);
 167         glLinkProgram(glsl_tally_program_num);
 168
 169         glGetProgramiv(glsl_tally_program_num, GL_LINK_STATUS, &success);
 170         if (success == GL_FALSE) {
 171                 GLchar error_log[1024] = {0};
 172                 glGetProgramInfoLog(glsl_tally_program_num, 1024, nullptr, error_log);
 173                 fprintf(stderr, "Error linking program: %s\n", error_log);
 174                 exit(1);
 175         }
 176
 177         // Compile the rANS shader.
 178         shader_src = ::read_file("rans.shader");
 179         shader_num = compile_shader(shader_src, GL_COMPUTE_SHADER);
 180         GLuint glsl_rans_program_num = glCreateProgram();
 181         glAttachShader(glsl_rans_program_num, shader_num);
 182         glLinkProgram(glsl_rans_program_num);
 183
 184         glGetProgramiv(glsl_rans_program_num, GL_LINK_STATUS, &success);
 185         if (success == GL_FALSE) {
 186                 GLchar error_log[1024] = {0};
 187                 glGetProgramInfoLog(glsl_rans_program_num, 1024, nullptr, error_log);
 188                 fprintf(stderr, "Error linking program: %s\n", error_log);
 189                 exit(1);
 190         }
 191         check_error();
 192
 193         // An SSBO for the rANS distributions.
 194         GLuint ssbo;
 195         glGenBuffers(1, &ssbo);
 196         glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssbo);
 197         glNamedBufferStorage(ssbo, 256 * 16 * sizeof(uint32_t), nullptr, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 198         check_error();
 199
 200         // SSBOs for the rANS output (data and offsets).
 201         GLuint output_ssbo;
 202         glGenBuffers(1, &output_ssbo);
 203         glBindBuffer(GL_SHADER_STORAGE_BUFFER, output_ssbo);
 204         glNamedBufferStorage(output_ssbo, 45 * 64 * 1024, nullptr, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 205         check_error();
 206
 207         GLuint output_offset_ssbo;
 208         glGenBuffers(1, &output_offset_ssbo);
 209         glBindBuffer(GL_SHADER_STORAGE_BUFFER, output_offset_ssbo);
 210         glNamedBufferStorage(output_offset_ssbo, 45 * 64 * sizeof(uint32_t), nullptr, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 211         check_error();
 212
 213         // Bind SSBOs.
 214         glUseProgram(glsl_program_num);
 215         glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 9, ssbo);
 216
 217         glUseProgram(glsl_tally_program_num);
 218         glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 9, ssbo);
 219
 220         glUseProgram(glsl_rans_program_num);
 221         glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 9, ssbo);
 222         glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 10, output_ssbo);
 223         glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 11, output_offset_ssbo);
 224
 225         glUseProgram(glsl_program_num);
 226         check_error();
 227
 228         // Upload luma.
 229         GLuint y_tex;
 230         glGenTextures(1, &y_tex);
 231         glBindTexture(GL_TEXTURE_2D, y_tex);
 232         glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
 233         glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
 234         glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
 235         glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
 236         glTexImage2D(GL_TEXTURE_2D, 0, GL_R8UI, WIDTH, HEIGHT, 0, GL_RED_INTEGER, GL_UNSIGNED_BYTE, pix_y);
 237         check_error();
 238
 239         // Make destination textures.
 240         GLuint dc_ac7_tex, ac1_ac6_tex, ac2_ac5_tex;
 241         for (GLuint *tex : { &dc_ac7_tex, &ac1_ac6_tex, &ac2_ac5_tex }) {
 242                 glGenTextures(1, tex);
 243                 glBindTexture(GL_TEXTURE_2D, *tex);
 244                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
 245                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
 246                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
 247                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
 248                 glTexImage2D(GL_TEXTURE_2D, 0, GL_R16UI, WIDTH / 8, HEIGHT, 0, GL_RED_INTEGER, GL_UNSIGNED_SHORT, nullptr);
 249                 check_error();
 250         }
 251
 252         GLuint ac3_tex, ac4_tex;
 253         for (GLuint *tex : { &ac3_tex, &ac4_tex }) {
 254                 glGenTextures(1, tex);
 255                 glBindTexture(GL_TEXTURE_2D, *tex);
 256                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
 257                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
 258                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
 259                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
 260                 glTexImage2D(GL_TEXTURE_2D, 0, GL_R8I, WIDTH / 8, HEIGHT, 0, GL_RED_INTEGER, GL_BYTE, nullptr);
 261                 check_error();
 262         }
 263
 264         glBindImageTexture(0, dc_ac7_tex, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_R16UI);
 265         glBindImageTexture(1, ac1_ac6_tex, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_R16UI);
 266         glBindImageTexture(2, ac2_ac5_tex, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_R16UI);
 267         glBindImageTexture(3, ac3_tex, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_R8I);
 268         glBindImageTexture(4, ac4_tex, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_R8I);
 269         glBindImageTexture(5, y_tex, 0, GL_FALSE, 0, GL_READ_ONLY, GL_R8UI);
 270         check_error();
 271
 272         // Bind uniforms.
 273         glUseProgram(glsl_program_num);
 274         GLint dc_ac7_tex_uniform = glGetUniformLocation(glsl_program_num, "dc_ac7_tex");
 275         GLint ac1_ac6_tex_uniform = glGetUniformLocation(glsl_program_num, "ac1_ac6_tex");
 276         GLint ac2_ac5_tex_uniform = glGetUniformLocation(glsl_program_num, "ac2_ac5_tex");
 277         GLint ac3_tex_uniform = glGetUniformLocation(glsl_program_num, "ac3_tex");
 278         GLint ac4_tex_uniform = glGetUniformLocation(glsl_program_num, "ac4_tex");
 279         GLint image_tex_uniform = glGetUniformLocation(glsl_program_num, "image_tex");
 280         glUniform1i(dc_ac7_tex_uniform, 0);
 281         glUniform1i(ac1_ac6_tex_uniform, 1);
 282         glUniform1i(ac2_ac5_tex_uniform, 2);
 283         glUniform1i(ac3_tex_uniform, 3);
 284         glUniform1i(ac4_tex_uniform, 4);
 285         glUniform1i(image_tex_uniform, 5);
 286
 287         glUseProgram(glsl_rans_program_num);
 288         dc_ac7_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "dc_ac7_tex");
 289         ac1_ac6_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "ac1_ac6_tex");
 290         ac2_ac5_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "ac2_ac5_tex");
 291         ac3_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "ac3_tex");
 292         ac4_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "ac4_tex");
 293         image_tex_uniform = glGetUniformLocation(glsl_rans_program_num, "image_tex");
 294         glUniform1i(dc_ac7_tex_uniform, 0);
 295         glUniform1i(ac1_ac6_tex_uniform, 1);
 296         glUniform1i(ac2_ac5_tex_uniform, 2);
 297         glUniform1i(ac3_tex_uniform, 3);
 298         glUniform1i(ac4_tex_uniform, 4);
 299
 300         steady_clock::time_point start = steady_clock::now();
 301         unsigned num_iterations = 100;
 302         for (unsigned i = 0; i < num_iterations; ++i) {
 303                 glClearNamedBufferSubData(ssbo, GL_R8, 0, 256 * 16 * sizeof(uint32_t), GL_RED, GL_UNSIGNED_BYTE, nullptr);
 304                 glUseProgram(glsl_program_num);
 305                 glDispatchCompute(WIDTH_BLOCKS / 16, HEIGHT_BLOCKS, 1);
 306                 glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
 307
 308                 glUseProgram(glsl_tally_program_num);
 309                 glDispatchCompute(4, 1, 1);
 310                 glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);
 311
 312                 glUseProgram(glsl_rans_program_num);
 313                 glDispatchCompute(NUM_BLOCKS / BLOCKS_PER_STREAM, 8, 5);
 314         }
 315         check_error();
 316         glFinish();
 317         check_error();
 318         steady_clock::time_point now = steady_clock::now();
 319
 320 #if 0
 321         printf("%ld bytes + %ld escape bits (%ld) = %ld total bytes\n",
 322                 tot_bytes - extra_bits / 8,
 323                 extra_bits,
 324                 extra_bits / 8,
 325                 tot_bytes);
 326
 327         printf("\n");
 328 #endif
 329
 330         printf("Each iteration took %.3f ms.\n", 1e3 * duration<double>(now - start).count() / num_iterations);
 331
 332         FILE *codedfp = fopen("coded.dat", "wb");
 333         if (codedfp == nullptr) {
 334                 perror("coded.dat");
 335                 exit(1);
 336         }
 337
 338         // Write out the distributions.
 339         const RansDistSSBO *rans_dist = (const RansDistSSBO *)glMapNamedBufferRange(ssbo, 0, 256 * 16 * sizeof(uint32_t), GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 340         for (unsigned r = 0; r < 2; ++r) {  // Hack to write fake chroma tables.
 341                 // TODO: rather gamma-k or something
 342                 for (unsigned i = 0; i < 4; ++i) {
 343                         printf("writing table %d\n", i);
 344                         for (unsigned j = 0; j < NUM_SYMS; ++j) {
 345                                 printf("%d,%d: start=%d freq=%d\n", i, j, rans_dist->ransdist[i * 256 + j].first, rans_dist->ransdist[i * 256 + j].second);
 346                                 write_varint(rans_dist->ransdist[i * 256 + j].second, codedfp);
 347                         }
 348                 }
 349         }
 350
 351         // Write out the actual data.
 352         // TODO: Do the deduplication.
 353
 354         const uint32_t *offsets = (const uint32_t *)glMapNamedBufferRange(output_offset_ssbo, 0, 45 * 64 * sizeof(uint32_t), GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 355 #if 0
 356         for (int i = 0; i < 45*64; ++i) {
 357                 printf("%d,%d,%d: %u\n", i / 64, (i / 8) % 8, i % 8, 1024 * (i + 1) - offsets[i]);
 358         }
 359 #endif
 360
 361         const uint8_t *data = (const uint8_t *)glMapNamedBufferRange(output_ssbo, 0, 45 * 64 * 1024, GL_MAP_READ_BIT | GL_MAP_PERSISTENT_BIT);
 362
 363         for (unsigned y = 0; y < 8; ++y) {
 364                 for (unsigned x = 0; x < 8; ++x) {
 365                         for (unsigned int stream_idx = 0; stream_idx < 45; ++stream_idx) {
 366                                 const uint8_t *out_end = data + (stream_idx * 64 + y * 8 + x + 1) * 1024;
 367                                 const uint8_t *ptr = data + offsets[stream_idx * 64 + y * 8 + x];
 368                                 uint32_t num_rans_bytes = out_end - ptr;
 369 #if 0
 370                                 if (num_rans_bytes == last_block.size() &&
 371                                     memcmp(last_block.data(), ptr, last_block.size()) == 0) {
 372                                         write_varint(0, codedfp);
 373                                         clear();
 374                                         return 1;
 375                                 } else {
 376                                         last_block = string((const char *)ptr, num_rans_bytes);
 377                                 }
 378 #endif
 379
 380                                 write_varint(num_rans_bytes, codedfp);
 381                                 fwrite(ptr, 1, num_rans_bytes, codedfp);
 382                         }
 383                 }
 384         }
 385         fclose(codedfp);
 386 }