restructured to allow qtables

This commit is contained in:
Muaz Ahmad 2023-11-28 13:44:46 +05:00
parent d670f8f880
commit af6d436079
3 changed files with 36 additions and 7 deletions

View file

@ -58,10 +58,10 @@ fn get_opts() !util.Options {
pub fn main() !void {
const opts = try get_opts();
var buffs = try util.Buffers.init(std.heap.page_allocator, opts.width, opts.height, 1000);
var buffs = try util.Buffers.init(std.heap.page_allocator, opts.width, opts.height, 1000, opts.quality);
defer buffs.deinit();
var thread_manager = try threads.ThreadManager.init(std.heap.page_allocator, opts.n_quant_jobs);
var thread_manager = try threads.ThreadManager.init(std.heap.page_allocator, opts.n_quant_jobs, buffs.Q_Lum, buffs.Q_Chrom);
defer thread_manager.deinit();
var f = std.io.getStdIn();

View file

@ -45,7 +45,7 @@ pub const ThreadManager = struct {
const Self = @This();
pub fn init(alloc: std.mem.Allocator, n_quant_jobs: usize) !Self {
pub fn init(alloc: std.mem.Allocator, n_quant_jobs: usize, Q_Lum: *util.QTable, Q_Chrom: *util.QTable) !Self {
var arena = std.heap.ArenaAllocator.init(alloc);
var arena_alloc = arena.allocator();
var job_pool = util.JobPool.init(alloc);
@ -59,7 +59,7 @@ pub const ThreadManager = struct {
thread_mgr.signals.* = Signals.init();
for (0..n_quant_jobs) |_| {
try thread_mgr.threads.append(try std.Thread.spawn(.{}, quantize_loop, .{ thread_mgr.queue_wrp.queue, thread_mgr.signals }));
try thread_mgr.threads.append(try std.Thread.spawn(.{}, quantize_loop, .{ thread_mgr.queue_wrp.queue, thread_mgr.signals, Q_Lum, Q_Chrom }));
}
return thread_mgr;
@ -71,10 +71,10 @@ pub const ThreadManager = struct {
}
};
fn quantize_loop(queue: *util.JobQueue, signals: *Signals) void {
fn quantize_loop(queue: *util.JobQueue, signals: *Signals, Q_Lum: *util.QTable, Q_Chrom: *util.QTable) void {
while (queue.HasJobs() or !signals.quit.load(std.builtin.AtomicOrder.Acquire)) : (std.time.sleep(1)) {
const job = queue.pop() orelse continue;
transform.quantize(job.source, job.target);
transform.quantize(job.source, job.target, if (job.is_lum) Q_Lum else Q_Chrom);
_ = @atomicRmw(u32, &signals.processed.value, .Add, 1, .SeqCst);
}
}

View file

@ -9,11 +9,13 @@ pub const Options = struct {
pub const Block = [64]u8;
pub const BlockQuantized = [64]i16;
pub const QTable = @Vector(64, f16);
pub const ThreadList = std.ArrayList(std.Thread);
pub const Job = struct {
source: *Block,
target: *BlockQuantized,
is_lum: bool,
};
pub const JobQueue = struct {
const List = std.TailQueue(Job);
@ -65,6 +67,9 @@ pub const Buffers = struct {
U: [][]Block,
V: [][]Block,
Q_Lum: *QTable,
Q_Chrom: *QTable,
Y_quant: [][][4]BlockQuantized,
U_quant: [][]BlockQuantized,
V_quant: [][]BlockQuantized,
@ -76,7 +81,7 @@ pub const Buffers = struct {
const Self = @This();
pub fn init(root_alloc: std.mem.Allocator, w: usize, h: usize, output_buff_len: usize) !Self {
pub fn init(root_alloc: std.mem.Allocator, w: usize, h: usize, output_buff_len: usize, q: f16) !Self {
var arena = std.heap.ArenaAllocator.init(root_alloc);
var alloc = arena.allocator();
const block_w = w / 16;
@ -88,6 +93,9 @@ pub const Buffers = struct {
.U = try alloc.alloc([]Block, block_h),
.V = try alloc.alloc([]Block, block_h),
.Q_Lum = try alloc.create(QTable),
.Q_Chrom = try alloc.create(QTable),
.Y_quant = try alloc.alloc([][4]BlockQuantized, block_h),
.U_quant = try alloc.alloc([]BlockQuantized, block_h),
.V_quant = try alloc.alloc([]BlockQuantized, block_h),
@ -106,6 +114,8 @@ pub const Buffers = struct {
buffs.U_quant[i] = try alloc.alloc(BlockQuantized, block_w);
buffs.V_quant[i] = try alloc.alloc(BlockQuantized, block_w);
}
buffs.Q_Lum.* = gen_qtable(255, q, 12);
buffs.Q_Chrom.* = gen_qtable(255, q, 6);
return buffs;
}
@ -113,3 +123,22 @@ pub const Buffers = struct {
self.arena.deinit();
}
};
pub fn gen_qtable(q_max: usize, q: f16, band_range: usize) @Vector(64, f16) {
var ret: @Vector(64, f16) = [_]f16{0.0} ** 64;
const q_min: usize = @intFromFloat(255 - 100 * q);
var step = (q_max - q_min) / band_range;
var idx: usize = 0;
for (0..16) |band_i| {
const band_len = band_i + 1;
for (0..band_len) |j| {
if (band_i < band_range) {
ret[idx + j] = @floatFromInt(q_min + step * band_i);
} else {
ret[idx + j] = @floatFromInt(q_max);
}
}
idx += band_len;
}
return ret;
}