lib/accy/src/kernel/library/image.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const choir_abi = @import("choir_abi");
3
4 const artifact_product = @import("../../artifact/model/root.zig");
5 const shape = @import("../../choir/shape/root.zig");
6 const entry = @import("entry.zig");
7 const extent_mod = @import("extent.zig");
8 const geometry_mod = @import("geometry.zig");
9 const kernel = @import("../root.zig");
10
11 pub const blur_family_version: u32 = 1;
12 pub const resize_family_version: u32 = 1;
13 pub const blur_radius_max: u32 = 15;
14 pub const image_thread_caps = geometry_mod.ThreadCaps{
15 .budget = 256,
16 .x_max = 64,
17 .y_max = 16,
18 };
19
20 pub const Axis = enum {
21 horizontal,
22 vertical,
23
24 pub fn name(self: Axis) []const u8 {
25 return switch (self) {
26 .horizontal => "h",
27 .vertical => "v",
28 };
29 }
30 };
31
32 pub const BlurPass = struct {
33 radius: u32,
34 axis: Axis,
35 width: u64 = 1,
36 height: u64 = 1,
37 threads: entry.Threads2D = .{ .x = 16, .y = 16 },
38
39 pub fn taps(self: BlurPass) u32 {
40 return self.radius * 2 + 1;
41 }
42 };
43
44 pub const Resize = struct {
45 dst_width: u64 = 1,
46 dst_height: u64 = 1,
47 src_width: u64 = 1,
48 src_height: u64 = 1,
49 threads: entry.Threads2D = .{ .x = 16, .y = 16 },
50 };
51
52 pub fn imageExtentValid(extent: u64) bool {
53 return extent >= 1 and extent <= extent_mod.runtime_extent_max;
54 }
55
56 pub fn imageThreadsForExtents(width: u64, height: u64) entry.Threads2D {
57 return geometry_mod.threadsForGrid(.{ .rows = height, .cols = width }, image_thread_caps);
58 }
59
60 pub fn imageThreadCandidatesForExtents(width: u64, height: u64) geometry_mod.ThreadCandidates {
61 return geometry_mod.threadCandidatesForGrid(.{ .rows = height, .cols = width }, image_thread_caps);
62 }
63
64 pub fn blurPassInstanceValid(instance: BlurPass) bool {
65 if (instance.radius == 0 or instance.radius > blur_radius_max) return false;
66 if (!imageExtentValid(instance.width) or !imageExtentValid(instance.height)) return false;
67 if (instance.threads.x == 0 or instance.threads.y == 0) return false;
68 return true;
69 }
70
71 pub fn resizeInstanceValid(instance: Resize) bool {
72 if (!imageExtentValid(instance.dst_width) or !imageExtentValid(instance.dst_height)) return false;
73 if (!imageExtentValid(instance.src_width) or !imageExtentValid(instance.src_height)) return false;
74 if (instance.threads.x == 0 or instance.threads.y == 0) return false;
75 return true;
76 }
77
78 pub fn gaussianWeightsAlloc(allocator: std.mem.Allocator, radius: u32, sigma: f32) ![]f32 {
79 if (radius == 0 or radius > blur_radius_max) return error.UnsupportedBlurRadius;
80 if (!(sigma > 0)) return error.UnsupportedBlurSigma;
81 const taps = radius * 2 + 1;
82 const weights = try allocator.alloc(f32, taps);
83 errdefer allocator.free(weights);
84 _ = try gaussianWeights(weights, radius, sigma);
85 return weights;
86 }
87
88 pub fn gaussianWeights(out: []f32, radius: u32, sigma: f32) ![]f32 {
89 if (radius == 0 or radius > blur_radius_max) return error.UnsupportedBlurRadius;
90 if (!(sigma > 0)) return error.UnsupportedBlurSigma;
91 const taps = radius * 2 + 1;
92 if (out.len < taps) return error.BufferTooSmall;
93 const weights = out[0..taps];
94 var total: f32 = 0;
95 for (weights, 0..) |*weight, tap| {
96 const offset = @as(f32, @floatFromInt(@as(i64, @intCast(tap)) - @as(i64, radius)));
97 const value = @exp(-(offset * offset) / (2 * sigma * sigma));
98 weight.* = value;
99 total += value;
100 }
101 for (weights) |*weight| weight.* /= total;
102 return weights;
103 }
104
105 const Channels = struct {
106 r: kernel.Value,
107 g: kernel.Value,
108 b: kernel.Value,
109 a: kernel.Value,
110 };
111
112 fn unpackChannels(inner: anytype, pixel: kernel.Value) !Channels {
113 const mask = try inner.constantInt(.u32, 0xff);
114 return .{
115 .r = try inner.cast(try inner.and_(pixel, mask), .f32),
116 .g = try inner.cast(try inner.and_(try inner.ushr(pixel, try inner.constantInt(.u32, 8)), mask), .f32),
117 .b = try inner.cast(try inner.and_(try inner.ushr(pixel, try inner.constantInt(.u32, 16)), mask), .f32),
118 .a = try inner.cast(try inner.ushr(pixel, try inner.constantInt(.u32, 24)), .f32),
119 };
120 }
121
122 fn packChannel(inner: anytype, value: kernel.Value) !kernel.Value {
123 const zero = try inner.constantFloat(.f32, 0);
124 const limit = try inner.constantFloat(.f32, 255);
125 const half = try inner.constantFloat(.f32, 0.5);
126 const rounded = try inner.floor(try inner.add(value, half));
127 const clamped = try inner.min(try inner.max(rounded, zero), limit);
128 return inner.cast(clamped, .u32);
129 }
130
131 fn packChannels(inner: anytype, channels: Channels) !kernel.Value {
132 const r = try packChannel(inner, channels.r);
133 const g = try packChannel(inner, channels.g);
134 const b = try packChannel(inner, channels.b);
135 const a = try packChannel(inner, channels.a);
136 return inner.or_(
137 try inner.or_(r, try inner.shl(g, try inner.constantInt(.u32, 8))),
138 try inner.or_(try inner.shl(b, try inner.constantInt(.u32, 16)), try inner.shl(a, try inner.constantInt(.u32, 24))),
139 );
140 }
141
142 fn blur_pass_body_active(inner: anytype, ctx: anytype) !void {
143 const zero_i32 = try inner.constantInt(.i32, 0);
144 const axis_extent = switch (ctx.spec.axis) {
145 .horizontal => ctx.width,
146 .vertical => ctx.height,
147 };
148 const axis_base = switch (ctx.spec.axis) {
149 .horizontal => ctx.col,
150 .vertical => ctx.row,
151 };
152 const axis_last = try inner.sub(try inner.cast(axis_extent, .i32), try inner.constantInt(.i32, 1));
153 const base_i32 = try inner.cast(axis_base, .i32);
154
155 var acc = Channels{
156 .r = try inner.constantFloat(.f32, 0),
157 .g = try inner.constantFloat(.f32, 0),
158 .b = try inner.constantFloat(.f32, 0),
159 .a = try inner.constantFloat(.f32, 0),
160 };
161 const tap_count = ctx.spec.taps();
162 var tap: u32 = 0;
163 while (tap < tap_count) : (tap += 1) {
164 const offset = @as(i64, tap) - @as(i64, ctx.spec.radius);
165 const offset_value = try inner.constantInt(.i32, offset);
166 const sample_raw = try inner.add(base_i32, offset_value);
167 const sample_clamped = try inner.min(try inner.max(sample_raw, zero_i32), axis_last);
168 const sample_index = try inner.castIndex(sample_clamped);
169 const pixel_index = switch (ctx.spec.axis) {
170 .horizontal => try inner.add(try inner.mul(ctx.row, ctx.width), sample_index),
171 .vertical => try inner.add(try inner.mul(sample_index, ctx.width), ctx.col),
172 };
173 const pixel = (try ctx.args.param(.src).load(inner, pixel_index)).raw();
174 const weight = (try ctx.args.param(.weights).load(inner, try inner.constantIndex(tap))).raw();
175 const channels = try unpackChannels(inner, pixel);
176 acc = .{
177 .r = try inner.fma(channels.r, weight, acc.r),
178 .g = try inner.fma(channels.g, weight, acc.g),
179 .b = try inner.fma(channels.b, weight, acc.b),
180 .a = try inner.fma(channels.a, weight, acc.a),
181 };
182 }
183 const out_index = try inner.add(try inner.mul(ctx.row, ctx.width), ctx.col);
184 try ctx.args.param(.dst).store(inner, try packChannels(inner, acc), out_index);
185 }
186
187 fn blurPassBody(k: anytype, spec: BlurPass, args: anytype) !void {
188 if (!blurPassInstanceValid(spec)) return error.UnsupportedBlurPassInstance;
189 const col = try k.globalId(.x);
190 const row = try k.globalId(.y);
191 const width = try k.castIndex(args.param(.width).raw());
192 const height = try k.castIndex(args.param(.height).raw());
193 const col_active = try k.compare(.lt, col, width);
194 const row_active = try k.compare(.lt, row, height);
195 const active = try k.and_(col_active, row_active);
196 try k.guardDo(active, .{
197 .args = args,
198 .spec = spec,
199 .col = col,
200 .row = row,
201 .width = width,
202 .height = height,
203 }, blur_pass_body_active);
204 }
205
206 fn blurPassSchedule(instance: BlurPass) kernel.logical.schedule.ThreadBlocks {
207 return kernel.logical.schedule.threadBlocks(.{
208 .x = instance.threads.x,
209 .y = instance.threads.y,
210 });
211 }
212
213 fn blurPassFamily() type {
214 return kernel.logical.Family(.{
215 .name = "accy_kernel_image_blur_pass_rgba8",
216 .parameters = .{
217 .dst = kernel.dynamicBuffer(.u32),
218 .src = kernel.dynamicBuffer(.u32),
219 .weights = kernel.dynamicBuffer(.f32),
220 .width = kernel.scalar(.i32),
221 .height = kernel.scalar(.i32),
222 },
223 .Instance = BlurPass,
224 .schedule = blurPassSchedule,
225 .body = blurPassBody,
226 });
227 }
228
229 pub const BlurPassFamilyRgba8 = blurPassFamily();
230
231 fn resize_body_active(inner: anytype, ctx: anytype) !void {
232 const sample_x = try sampleAxis(inner, ctx.col, ctx.args.param(.scale_x).raw(), ctx.args.param(.src_last_x).raw());
233 const sample_y = try sampleAxis(inner, ctx.row, ctx.args.param(.scale_y).raw(), ctx.args.param(.src_last_y).raw());
234 const src_width_index = try inner.castIndex(ctx.args.param(.src_width).raw());
235
236 const row0_base = try inner.mul(sample_y.lo, src_width_index);
237 const row1_base = try inner.mul(sample_y.hi, src_width_index);
238 const p00 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row0_base, sample_x.lo))).raw());
239 const p01 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row0_base, sample_x.hi))).raw());
240 const p10 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row1_base, sample_x.lo))).raw());
241 const p11 = try unpackChannels(inner, (try ctx.args.param(.src).load(inner, try inner.add(row1_base, sample_x.hi))).raw());
242
243 const top = try lerpChannels(inner, p00, p01, sample_x.frac);
244 const bottom = try lerpChannels(inner, p10, p11, sample_x.frac);
245 const blended = try lerpChannels(inner, top, bottom, sample_y.frac);
246
247 const out_index = try inner.add(try inner.mul(ctx.row, ctx.dst_width), ctx.col);
248 try ctx.args.param(.dst).store(inner, try packChannels(inner, blended), out_index);
249 }
250
251 fn resizeBody(k: anytype, spec: Resize, args: anytype) !void {
252 if (!resizeInstanceValid(spec)) return error.UnsupportedResizeInstance;
253 const col = try k.globalId(.x);
254 const row = try k.globalId(.y);
255 const dst_width = try k.castIndex(args.param(.dst_width).raw());
256 const dst_height = try k.castIndex(args.param(.dst_height).raw());
257 const col_active = try k.compare(.lt, col, dst_width);
258 const row_active = try k.compare(.lt, row, dst_height);
259 const active = try k.and_(col_active, row_active);
260 try k.guardDo(active, .{
261 .args = args,
262 .col = col,
263 .row = row,
264 .dst_width = dst_width,
265 }, resize_body_active);
266 }
267
268 const AxisSample = struct {
269 lo: kernel.Value,
270 hi: kernel.Value,
271 frac: kernel.Value,
272 };
273
274 fn sampleAxis(inner: anytype, dst_index: kernel.Value, scale: kernel.Value, last: kernel.Value) !AxisSample {
275 const zero = try inner.constantFloat(.f32, 0);
276 const half = try inner.constantFloat(.f32, 0.5);
277 const one = try inner.constantFloat(.f32, 1);
278 const negative_half = try inner.constantFloat(.f32, -0.5);
279 const centered = try inner.fma(try inner.add(try inner.cast(try inner.cast(dst_index, .u32), .f32), half), scale, negative_half);
280 const clamped = try inner.min(try inner.max(centered, zero), last);
281 const lo_f = try inner.floor(clamped);
282 const frac = try inner.sub(clamped, lo_f);
283 const lo_u32 = try inner.cast(lo_f, .u32);
284 const lo = try inner.castIndex(lo_u32);
285 const hi_f = try inner.min(try inner.add(lo_f, one), last);
286 const hi = try inner.castIndex(try inner.cast(hi_f, .u32));
287 return .{ .lo = lo, .hi = hi, .frac = frac };
288 }
289
290 fn lerpChannels(inner: anytype, from: Channels, to: Channels, t: kernel.Value) !Channels {
291 return .{
292 .r = try lerp(inner, from.r, to.r, t),
293 .g = try lerp(inner, from.g, to.g, t),
294 .b = try lerp(inner, from.b, to.b, t),
295 .a = try lerp(inner, from.a, to.a, t),
296 };
297 }
298
299 fn lerp(inner: anytype, from: kernel.Value, to: kernel.Value, t: kernel.Value) !kernel.Value {
300 return inner.fma(try inner.sub(to, from), t, from);
301 }
302
303 fn resizeSchedule(instance: Resize) kernel.logical.schedule.ThreadBlocks {
304 return kernel.logical.schedule.threadBlocks(.{
305 .x = instance.threads.x,
306 .y = instance.threads.y,
307 });
308 }
309
310 fn resizeFamily() type {
311 return kernel.logical.Family(.{
312 .name = "accy_kernel_image_resize_bilinear_rgba8",
313 .parameters = .{
314 .dst = kernel.dynamicBuffer(.u32),
315 .src = kernel.dynamicBuffer(.u32),
316 .dst_width = kernel.scalar(.i32),
317 .dst_height = kernel.scalar(.i32),
318 .src_width = kernel.scalar(.i32),
319 .scale_x = kernel.scalar(.f32),
320 .scale_y = kernel.scalar(.f32),
321 .src_last_x = kernel.scalar(.f32),
322 .src_last_y = kernel.scalar(.f32),
323 },
324 .Instance = Resize,
325 .schedule = resizeSchedule,
326 .body = resizeBody,
327 });
328 }
329
330 pub const ResizeBilinearFamilyRgba8 = resizeFamily();
331
332 pub fn blurPassFamilyTarget(allocator: std.mem.Allocator, instance: BlurPass) ![]u8 {
333 return std.fmt.allocPrint(
334 allocator,
335 "image_blur_pass_family_r{d}{s}_{d}x{d}_rgba8",
336 .{ instance.radius, instance.axis.name(), instance.threads.x, instance.threads.y },
337 );
338 }
339
340 pub fn blurPassFamilyEntryName(allocator: std.mem.Allocator, instance: BlurPass) ![]u8 {
341 return std.fmt.allocPrint(
342 allocator,
343 "accy_image_blur_pass_r{d}{s}_rgba8",
344 .{ instance.radius, instance.axis.name() },
345 );
346 }
347
348 pub fn resizeFamilyTarget(allocator: std.mem.Allocator, instance: Resize) ![]u8 {
349 return std.fmt.allocPrint(
350 allocator,
351 "image_resize_bilinear_family_{d}x{d}_rgba8",
352 .{ instance.threads.x, instance.threads.y },
353 );
354 }
355
356 pub fn resizeFamilyEntryName(allocator: std.mem.Allocator, instance: Resize) ![]u8 {
357 _ = instance;
358 return allocator.dupe(u8, "accy_image_resize_bilinear_rgba8");
359 }
360
361 pub fn blurPassRuntimeArguments(width: u32, height: u32) ![2]choir_abi.ScalarArgument {
362 if (width == 0 or height == 0) return error.UnsupportedImageExtent;
363 if (width > std.math.maxInt(i32) or height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;
364 return .{
365 .{ .i32 = @intCast(width) },
366 .{ .i32 = @intCast(height) },
367 };
368 }
369
370 pub fn resizeScale(src_extent: u32, dst_extent: u32) f32 {
371 return @as(f32, @floatFromInt(src_extent)) / @as(f32, @floatFromInt(dst_extent));
372 }
373
374 pub fn resizeRuntimeArguments(dst_width: u32, dst_height: u32, src_width: u32, src_height: u32) ![7]choir_abi.ScalarArgument {
375 if (dst_width == 0 or dst_height == 0 or src_width == 0 or src_height == 0) return error.UnsupportedImageExtent;
376 if (dst_width > std.math.maxInt(i32) or dst_height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;
377 if (src_width > std.math.maxInt(i32) or src_height > std.math.maxInt(i32)) return error.UnsupportedImageExtent;
378 return .{
379 .{ .i32 = @intCast(dst_width) },
380 .{ .i32 = @intCast(dst_height) },
381 .{ .i32 = @intCast(src_width) },
382 .{ .f32 = resizeScale(src_width, dst_width) },
383 .{ .f32 = resizeScale(src_height, dst_height) },
384 .{ .f32 = @floatFromInt(src_width - 1) },
385 .{ .f32 = @floatFromInt(src_height - 1) },
386 };
387 }
388
389 pub fn imageLaunchGeometry(instance_threads: entry.Threads2D, width: u32, height: u32) choir_abi.LaunchGeometry {
390 return .{
391 .grid = .{
392 (width + instance_threads.x - 1) / instance_threads.x,
393 (height + instance_threads.y - 1) / instance_threads.y,
394 1,
395 },
396 .threadgroup = .{ instance_threads.x, instance_threads.y, 1 },
397 };
398 }
399
400 pub fn blurPassAxisParameter(axis: Axis) u64 {
401 return switch (axis) {
402 .horizontal => 0,
403 .vertical => 1,
404 };
405 }
406
407 pub fn blurPassShapeProfileDimensions(instance: BlurPass) [2]artifact_product.KernelCallShapeProfileDimension {
408 _ = instance;
409 const bounds = imageRuntimeExtentBounds();
410 return .{
411 .{
412 .name = "x",
413 .runtime_scalar_argument_index = 0,
414 .bounds = bounds,
415 },
416 .{
417 .name = "y",
418 .runtime_scalar_argument_index = 1,
419 .bounds = bounds,
420 },
421 };
422 }
423
424 pub fn resizeShapeProfileDimensions(instance: Resize) [3]artifact_product.KernelCallShapeProfileDimension {
425 _ = instance;
426 const bounds = imageRuntimeExtentBounds();
427 return .{
428 .{
429 .name = "dst_x",
430 .runtime_scalar_argument_index = 0,
431 .bounds = bounds,
432 },
433 .{
434 .name = "dst_y",
435 .runtime_scalar_argument_index = 1,
436 .bounds = bounds,
437 },
438 .{
439 .name = "src_x",
440 .runtime_scalar_argument_index = 2,
441 .bounds = bounds,
442 },
443 };
444 }
445
446 fn imageRuntimeExtentBounds() shape.Bounds {
447 return .{ .min = 1, .max = extent_mod.runtime_extent_max };
448 }
449
450 fn imageDerivedLaunch(threads: entry.Threads2D) !artifact_product.KernelCallLaunch {
451 if (threads.x == 0 or threads.y == 0) return error.KernelLibraryLaunchThreadgroupMustBeNonzero;
452 return .{ .derived = .{
453 .grid = .{
454 .{ .runtime_u32_ceil_div = .{ .argument_index = 0, .divisor = threads.x } },
455 .{ .runtime_u32_ceil_div = .{ .argument_index = 1, .divisor = threads.y } },
456 .{ .fixed = 1 },
457 },
458 .threadgroup = .{ threads.x, threads.y, 1 },
459 } };
460 }
461
462 pub fn blurPassShapeFamily(backing_allocator: std.mem.Allocator, instance: BlurPass) !shape.Family {
463 if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;
464 var builder = try shape.Builder.init(backing_allocator, "image_blur_pass");
465 errdefer builder.deinit();
466
467 const y = try builder.symbol("y");
468 const x = try builder.symbol("x");
469 const y_expr = try builder.symbolExpression(y);
470 const x_expr = try builder.symbolExpression(x);
471 const taps_expr = builder.constantExpression(@intCast(instance.taps()));
472
473 _ = try builder.tensor("src", &.{ y_expr, x_expr });
474 _ = try builder.tensor("weights", &.{taps_expr});
475 _ = try builder.tensor("out", &.{ y_expr, x_expr });
476 try builder.assumeBounds(y_expr, imageRuntimeExtentBounds());
477 try builder.assumeBounds(x_expr, imageRuntimeExtentBounds());
478
479 return builder.finish();
480 }
481
482 pub fn resizeShapeFamily(backing_allocator: std.mem.Allocator, instance: Resize) !shape.Family {
483 if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;
484 var builder = try shape.Builder.init(backing_allocator, "image_resize_bilinear");
485 errdefer builder.deinit();
486
487 const dst_y = try builder.symbol("dst_y");
488 const dst_x = try builder.symbol("dst_x");
489 const src_y = try builder.symbol("src_y");
490 const src_x = try builder.symbol("src_x");
491 const dst_y_expr = try builder.symbolExpression(dst_y);
492 const dst_x_expr = try builder.symbolExpression(dst_x);
493 const src_y_expr = try builder.symbolExpression(src_y);
494 const src_x_expr = try builder.symbolExpression(src_x);
495
496 _ = try builder.tensor("src", &.{ src_y_expr, src_x_expr });
497 _ = try builder.tensor("out", &.{ dst_y_expr, dst_x_expr });
498 try builder.assumeBounds(dst_y_expr, imageRuntimeExtentBounds());
499 try builder.assumeBounds(dst_x_expr, imageRuntimeExtentBounds());
500 try builder.assumeBounds(src_y_expr, imageRuntimeExtentBounds());
501 try builder.assumeBounds(src_x_expr, imageRuntimeExtentBounds());
502
503 return builder.finish();
504 }
505
506 pub fn blurPassFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: BlurPass) !u64 {
507 var family = try blurPassShapeFamily(backing_allocator, instance);
508 defer family.deinit();
509 return shape.fingerprint(family);
510 }
511
512 pub fn resizeFamilyFingerprint(backing_allocator: std.mem.Allocator, instance: Resize) !u64 {
513 var family = try resizeShapeFamily(backing_allocator, instance);
514 defer family.deinit();
515 return shape.fingerprint(family);
516 }
517
518 pub fn blurPassFamilySpecialization(backing_allocator: std.mem.Allocator, instance: BlurPass) !entry.OwnedSpecialization {
519 if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;
520 var owned = entry.OwnedSpecialization.init(backing_allocator);
521 errdefer owned.deinit();
522 const lifetime_allocator = owned.allocator();
523
524 const inputs = try lifetime_allocator.alloc(entry.Shape, 2);
525 inputs[0] = try entry.runtimeShape2D(lifetime_allocator, "y", instance.height, "x", instance.width);
526 inputs[1] = try entry.runtimeShape1D(lifetime_allocator, "tap", instance.taps());
527
528 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
529 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, "y", instance.height, "x", instance.width);
530
531 const reductions = try lifetime_allocator.alloc(entry.Reduction, 1);
532 reductions[0] = try entry.runtimeReduction(
533 lifetime_allocator,
534 "blur_tap",
535 .weighted_sum,
536 try entry.runtimeShape1D(lifetime_allocator, "tap", instance.taps()),
537 );
538
539 const static_parameters = try lifetime_allocator.alloc(entry.StaticParameter, 2);
540 static_parameters[0] = try entry.runtimeStaticParameter(lifetime_allocator, "radius", instance.radius);
541 static_parameters[1] = try entry.runtimeStaticParameter(lifetime_allocator, "axis", blurPassAxisParameter(instance.axis));
542
543 owned.value = .{
544 .dtype = .u32,
545 .accumulation_dtype = .f32,
546 .operation = .{ .image = .blur_pass },
547 .inputs = inputs,
548 .outputs = outputs,
549 .reductions = reductions,
550 .static_parameters = static_parameters,
551 .schedule = try entry.runtimeThreadBlocks2D(
552 lifetime_allocator,
553 "x",
554 instance.width,
555 "y",
556 instance.height,
557 instance.threads.x,
558 instance.threads.y,
559 ),
560 };
561 owned.value.launch = owned.value.schedule.?.launch();
562 var family = try blurPassShapeFamily(backing_allocator, instance);
563 errdefer family.deinit();
564 try owned.takeShapeFamily(&family);
565 return owned;
566 }
567
568 pub fn resizeFamilySpecialization(backing_allocator: std.mem.Allocator, instance: Resize) !entry.OwnedSpecialization {
569 if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;
570 var owned = entry.OwnedSpecialization.init(backing_allocator);
571 errdefer owned.deinit();
572 const lifetime_allocator = owned.allocator();
573
574 const inputs = try lifetime_allocator.alloc(entry.Shape, 1);
575 inputs[0] = try entry.runtimeShape2D(lifetime_allocator, "src_y", instance.src_height, "src_x", instance.src_width);
576
577 const outputs = try lifetime_allocator.alloc(entry.Shape, 1);
578 outputs[0] = try entry.runtimeShape2D(lifetime_allocator, "dst_y", instance.dst_height, "dst_x", instance.dst_width);
579
580 owned.value = .{
581 .dtype = .u32,
582 .accumulation_dtype = .f32,
583 .operation = .{ .image = .resize_bilinear },
584 .inputs = inputs,
585 .outputs = outputs,
586 .schedule = try entry.runtimeThreadBlocks2D(
587 lifetime_allocator,
588 "dst_x",
589 instance.dst_width,
590 "dst_y",
591 instance.dst_height,
592 instance.threads.x,
593 instance.threads.y,
594 ),
595 };
596 owned.value.launch = owned.value.schedule.?.launch();
597 var family = try resizeShapeFamily(backing_allocator, instance);
598 errdefer family.deinit();
599 try owned.takeShapeFamily(&family);
600 return owned;
601 }
602
603 pub fn blurPassInstanceFromSpecialization(specialization: entry.Specialization) ?BlurPass {
604 if (!specialization.scheduleMatchesLaunch()) return null;
605 if (!specialization.operationIs(.{ .image = .blur_pass })) return null;
606 if (specialization.dtype != .u32 or specialization.accumulation_dtype != .f32) return null;
607 if (specialization.inputs.len != 2 or specialization.outputs.len != 1) return null;
608 if (specialization.reductions.len != 1 or specialization.static_parameters.len != 2) return null;
609
610 const src = specialization.inputs[0];
611 const weights = specialization.inputs[1];
612 const output = specialization.outputs[0];
613 const reduction = specialization.reductions[0];
614 if (src.axes.len != 2 or weights.axes.len != 1 or output.axes.len != 2) return null;
615 if (reduction.shape.axes.len != 1) return null;
616 if (!std.mem.eql(u8, src.axes[0].name, output.axes[0].name)) return null;
617 if (!std.mem.eql(u8, src.axes[1].name, output.axes[1].name)) return null;
618 if (!std.mem.eql(u8, reduction.name, "blur_tap")) return null;
619 if (reduction.operator != .weighted_sum) return null;
620 if (!std.mem.eql(u8, reduction.shape.axes[0].name, weights.axes[0].name)) return null;
621
622 const radius_value = specialization.staticParameterValue("radius") orelse return null;
623 const axis_value = specialization.staticParameterValue("axis") orelse return null;
624 const radius = std.math.cast(u32, radius_value) orelse return null;
625 const axis: Axis = switch (axis_value) {
626 0 => .horizontal,
627 1 => .vertical,
628 else => return null,
629 };
630 const taps = 2 * @as(u64, radius) + 1;
631 if (weights.axes[0].extent != taps or reduction.shape.axes[0].extent != taps) return null;
632 if (!src.matchesExtents(&.{ output.axes[0].extent, output.axes[1].extent })) return null;
633
634 const launch = specialization.launch orelse return null;
635 const instance = BlurPass{
636 .radius = radius,
637 .axis = axis,
638 .width = output.axes[1].extent,
639 .height = output.axes[0].extent,
640 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
641 };
642 if (!blurPassInstanceValid(instance)) return null;
643 return instance;
644 }
645
646 pub fn resizeInstanceFromSpecialization(specialization: entry.Specialization) ?Resize {
647 if (!specialization.scheduleMatchesLaunch()) return null;
648 if (!specialization.operationIs(.{ .image = .resize_bilinear })) return null;
649 if (specialization.dtype != .u32 or specialization.accumulation_dtype != .f32) return null;
650 if (specialization.inputs.len != 1 or specialization.outputs.len != 1) return null;
651 if (specialization.reductions.len != 0 or specialization.static_parameters.len != 0) return null;
652
653 const src = specialization.inputs[0];
654 const output = specialization.outputs[0];
655 if (src.axes.len != 2 or output.axes.len != 2) return null;
656
657 const launch = specialization.launch orelse return null;
658 const instance = Resize{
659 .dst_width = output.axes[1].extent,
660 .dst_height = output.axes[0].extent,
661 .src_width = src.axes[1].extent,
662 .src_height = src.axes[0].extent,
663 .threads = .{ .x = launch.threadgroup[0], .y = launch.threadgroup[1] },
664 };
665 if (!resizeInstanceValid(instance)) return null;
666 return instance;
667 }
668
669 pub fn createBlurPassFamilyArtifact(
670 allocator: std.mem.Allocator,
671 handle: kernel.BackendHandle,
672 instance: BlurPass,
673 options: entry.ArtifactOptions,
674 ) !kernel.OwnedKernelCallArtifact {
675 if (!blurPassInstanceValid(instance)) return error.UnsupportedBlurPassInstance;
676 const target = try blurPassFamilyTarget(allocator, instance);
677 defer allocator.free(target);
678 const entry_name = try blurPassFamilyEntryName(allocator, instance);
679 defer allocator.free(entry_name);
680 const family_fingerprint = options.shape_family_fingerprint orelse try blurPassFamilyFingerprint(allocator, instance);
681 const shape_profile_dimensions = blurPassShapeProfileDimensions(instance);
682 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
683 .name = "image_blur_pass",
684 .fingerprint = family_fingerprint,
685 .dimensions = shape_profile_dimensions[0..],
686 };
687
688 var graph = try BlurPassFamilyRgba8.buildNamed(allocator, options.limits, entry_name, instance);
689 defer graph.deinit();
690 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
691 .target = target,
692 .version = blur_family_version,
693 .format = options.format,
694 .kernel_plan = options.kernel_plan,
695 .element_count_argument = options.element_count_argument,
696 .shape_family_fingerprint = family_fingerprint,
697 .shape_profile = shape_profile,
698 .launch = options.launch orelse try imageDerivedLaunch(instance.threads),
699 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 2 else options.runtime_scalar_argument_count,
700 .static_arguments = options.static_arguments,
701 });
702 }
703
704 pub fn createResizeFamilyArtifact(
705 allocator: std.mem.Allocator,
706 handle: kernel.BackendHandle,
707 instance: Resize,
708 options: entry.ArtifactOptions,
709 ) !kernel.OwnedKernelCallArtifact {
710 if (!resizeInstanceValid(instance)) return error.UnsupportedResizeInstance;
711 const target = try resizeFamilyTarget(allocator, instance);
712 defer allocator.free(target);
713 const entry_name = try resizeFamilyEntryName(allocator, instance);
714 defer allocator.free(entry_name);
715 const family_fingerprint = options.shape_family_fingerprint orelse try resizeFamilyFingerprint(allocator, instance);
716 const shape_profile_dimensions = resizeShapeProfileDimensions(instance);
717 const shape_profile = options.shape_profile orelse artifact_product.KernelCallShapeProfile{
718 .name = "image_resize_bilinear",
719 .fingerprint = family_fingerprint,
720 .dimensions = shape_profile_dimensions[0..],
721 };
722
723 var graph = try ResizeBilinearFamilyRgba8.buildNamed(allocator, options.limits, entry_name, instance);
724 defer graph.deinit();
725 return kernel.createKernelCallArtifact(allocator, handle, &graph, .{
726 .target = target,
727 .version = resize_family_version,
728 .format = options.format,
729 .kernel_plan = options.kernel_plan,
730 .element_count_argument = options.element_count_argument,
731 .shape_family_fingerprint = family_fingerprint,
732 .shape_profile = shape_profile,
733 .launch = options.launch orelse try imageDerivedLaunch(instance.threads),
734 .runtime_scalar_argument_count = if (options.runtime_scalar_argument_count == 0) 7 else options.runtime_scalar_argument_count,
735 .static_arguments = options.static_arguments,
736 });
737 }
738
739 pub fn referenceBlurPass(dst: []u32, src: []const u32, weights: []const f32, width: u32, height: u32, axis: Axis) void {
740 const radius: i64 = @intCast((weights.len - 1) / 2);
741 var row: u32 = 0;
742 while (row < height) : (row += 1) {
743 var col: u32 = 0;
744 while (col < width) : (col += 1) {
745 var acc = [4]f32{ 0, 0, 0, 0 };
746 for (weights, 0..) |weight, tap| {
747 const offset = @as(i64, @intCast(tap)) - radius;
748 const base: i64 = switch (axis) {
749 .horizontal => @as(i64, col),
750 .vertical => @as(i64, row),
751 };
752 const extent: i64 = switch (axis) {
753 .horizontal => @as(i64, width),
754 .vertical => @as(i64, height),
755 };
756 const sample = std.math.clamp(base + offset, 0, extent - 1);
757 const index: usize = switch (axis) {
758 .horizontal => @as(usize, row) * width + @as(usize, @intCast(sample)),
759 .vertical => @as(usize, @intCast(sample)) * width + col,
760 };
761 const pixel = src[index];
762 acc[0] = @mulAdd(f32, @floatFromInt(pixel & 0xff), weight, acc[0]);
763 acc[1] = @mulAdd(f32, @floatFromInt((pixel >> 8) & 0xff), weight, acc[1]);
764 acc[2] = @mulAdd(f32, @floatFromInt((pixel >> 16) & 0xff), weight, acc[2]);
765 acc[3] = @mulAdd(f32, @floatFromInt(pixel >> 24), weight, acc[3]);
766 }
767 dst[@as(usize, row) * width + col] = packReferenceChannels(acc);
768 }
769 }
770 }
771
772 pub fn referenceResizeBilinear(dst: []u32, src: []const u32, dst_width: u32, dst_height: u32, src_width: u32, src_height: u32) void {
773 var row: u32 = 0;
774 while (row < dst_height) : (row += 1) {
775 var col: u32 = 0;
776 while (col < dst_width) : (col += 1) {
777 const sx = referenceSampleAxis(col, dst_width, src_width);
778 const sy = referenceSampleAxis(row, dst_height, src_height);
779 const p00 = referenceUnpack(src[sy.lo * src_width + sx.lo]);
780 const p01 = referenceUnpack(src[sy.lo * src_width + sx.hi]);
781 const p10 = referenceUnpack(src[sy.hi * src_width + sx.lo]);
782 const p11 = referenceUnpack(src[sy.hi * src_width + sx.hi]);
783 var blended: [4]f32 = undefined;
784 for (0..4) |channel| {
785 const top = @mulAdd(f32, p01[channel] - p00[channel], sx.frac, p00[channel]);
786 const bottom = @mulAdd(f32, p11[channel] - p10[channel], sx.frac, p10[channel]);
787 blended[channel] = @mulAdd(f32, bottom - top, sy.frac, top);
788 }
789 dst[@as(usize, row) * dst_width + col] = packReferenceChannels(blended);
790 }
791 }
792 }
793
794 const ReferenceAxisSample = struct {
795 lo: usize,
796 hi: usize,
797 frac: f32,
798 };
799
800 fn referenceSampleAxis(dst_index: u32, dst_extent: u32, src_extent: u32) ReferenceAxisSample {
801 const scale = resizeScale(src_extent, dst_extent);
802 const centered = @mulAdd(f32, @as(f32, @floatFromInt(dst_index)) + 0.5, scale, -0.5);
803 const last = @as(f32, @floatFromInt(src_extent)) - 1;
804 const clamped = @max(@as(f32, 0), @min(centered, last));
805 const lo_f = @floor(clamped);
806 const hi_f = @min(lo_f + 1, last);
807 return .{
808 .lo = @intFromFloat(lo_f),
809 .hi = @intFromFloat(hi_f),
810 .frac = clamped - lo_f,
811 };
812 }
813
814 fn referenceUnpack(pixel: u32) [4]f32 {
815 return .{
816 @floatFromInt(pixel & 0xff),
817 @floatFromInt((pixel >> 8) & 0xff),
818 @floatFromInt((pixel >> 16) & 0xff),
819 @floatFromInt(pixel >> 24),
820 };
821 }
822
823 fn packReferenceChannels(channels: [4]f32) u32 {
824 var packed_pixel: u32 = 0;
825 for (channels, 0..) |value, channel| {
826 const rounded = @floor(value + 0.5);
827 const clamped = @max(@as(f32, 0), @min(rounded, 255));
828 packed_pixel |= @as(u32, @intFromFloat(clamped)) << @intCast(channel * 8);
829 }
830 return packed_pixel;
831 }
832
833 const testing = std.testing;
834
835 fn testPixel(seed: usize) u32 {
836 var value: u32 = @truncate(seed *% 2654435761);
837 value ^= value >> 13;
838 value *%= 0x5bd1e995;
839 value ^= value >> 15;
840 return value;
841 }
842
843 test "gaussian weights normalize and peak at the center" {
844 const weights = try gaussianWeightsAlloc(testing.allocator, 3, 1.4);
845 defer testing.allocator.free(weights);
846 try testing.expectEqual(@as(usize, 7), weights.len);
847 var total: f32 = 0;
848 for (weights) |weight| total += weight;
849 try testing.expectApproxEqAbs(@as(f32, 1), total, 0.0001);
850 for (weights) |weight| try testing.expect(weight <= weights[3]);
851 }
852
853 test "blur pass family matches the reference on the interpreter" {
854 const allocator = testing.allocator;
855 const width: u32 = 13;
856 const height: u32 = 7;
857 const pixel_count = @as(usize, width) * height;
858
859 const src = try allocator.alloc(u32, pixel_count);
860 defer allocator.free(src);
861 for (src, 0..) |*pixel, index| pixel.* = testPixel(index);
862
863 const weights = try gaussianWeightsAlloc(allocator, 2, 1.1);
864 defer allocator.free(weights);
865
866 inline for (.{ Axis.horizontal, Axis.vertical }) |axis| {
867 const instance = BlurPass{ .radius = 2, .axis = axis, .threads = .{ .x = 8, .y = 4 } };
868 const expected = try allocator.alloc(u32, pixel_count);
869 defer allocator.free(expected);
870 referenceBlurPass(expected, src, weights, width, height, axis);
871
872 const actual = try allocator.alloc(u32, pixel_count);
873 defer allocator.free(actual);
874 @memset(actual, 0);
875
876 var graph = try BlurPassFamilyRgba8.build(allocator, BlurPassFamilyRgba8.Limits.testing, instance);
877 defer graph.deinit();
878 const geometry = imageLaunchGeometry(instance.threads, width, height);
879 try graph.runCpuWithLaunch(allocator, &.{
880 kernel.argumentBuffer(u32, actual),
881 kernel.argumentBuffer(u32, @constCast(src)),
882 kernel.argumentBuffer(f32, @constCast(weights)),
883 kernel.argumentI32(@intCast(width)),
884 kernel.argumentI32(@intCast(height)),
885 }, .{
886 .grid = geometry.grid,
887 .block = geometry.threadgroup,
888 });
889
890 try testing.expectEqualSlices(u32, expected, actual);
891 }
892 }
893
894 test "resize bilinear family matches the reference on the interpreter" {
895 const allocator = testing.allocator;
896 const src_width: u32 = 12;
897 const src_height: u32 = 9;
898 const dst_width: u32 = 7;
899 const dst_height: u32 = 5;
900
901 const src = try allocator.alloc(u32, @as(usize, src_width) * src_height);
902 defer allocator.free(src);
903 for (src, 0..) |*pixel, index| pixel.* = testPixel(index +% 17);
904
905 const expected = try allocator.alloc(u32, @as(usize, dst_width) * dst_height);
906 defer allocator.free(expected);
907 referenceResizeBilinear(expected, src, dst_width, dst_height, src_width, src_height);
908
909 const actual = try allocator.alloc(u32, @as(usize, dst_width) * dst_height);
910 defer allocator.free(actual);
911 @memset(actual, 0);
912
913 const instance = Resize{ .threads = .{ .x = 8, .y = 4 } };
914 var graph = try ResizeBilinearFamilyRgba8.build(allocator, ResizeBilinearFamilyRgba8.Limits.testing, instance);
915 defer graph.deinit();
916 const geometry = imageLaunchGeometry(instance.threads, dst_width, dst_height);
917 try graph.runCpuWithLaunch(allocator, &.{
918 kernel.argumentBuffer(u32, actual),
919 kernel.argumentBuffer(u32, @constCast(src)),
920 kernel.argumentI32(@intCast(dst_width)),
921 kernel.argumentI32(@intCast(dst_height)),
922 kernel.argumentI32(@intCast(src_width)),
923 kernel.argumentF32(resizeScale(src_width, dst_width)),
924 kernel.argumentF32(resizeScale(src_height, dst_height)),
925 kernel.argumentF32(@floatFromInt(src_width - 1)),
926 kernel.argumentF32(@floatFromInt(src_height - 1)),
927 }, .{
928 .grid = geometry.grid,
929 .block = geometry.threadgroup,
930 });
931
932 try testing.expectEqualSlices(u32, expected, actual);
933 }
934
935 test "upscale resize keeps corner pixels exact" {
936 const allocator = testing.allocator;
937 const src = [_]u32{ 0xff000011, 0xff000022, 0xff000033, 0xff000044 };
938 const dst = try allocator.alloc(u32, 16);
939 defer allocator.free(dst);
940 referenceResizeBilinear(dst, src[0..], 4, 4, 2, 2);
941 try testing.expectEqual(src[0], dst[0]);
942 try testing.expectEqual(src[1], dst[3]);
943 try testing.expectEqual(src[2], dst[12]);
944 try testing.expectEqual(src[3], dst[15]);
945 }
946
947 test "blur instance validation bounds radius and threads" {
948 try testing.expect(blurPassInstanceValid(.{ .radius = 1, .axis = .horizontal }));
949 try testing.expect(!blurPassInstanceValid(.{ .radius = 0, .axis = .horizontal }));
950 try testing.expect(!blurPassInstanceValid(.{ .radius = blur_radius_max + 1, .axis = .vertical }));
951 try testing.expect(!blurPassInstanceValid(.{ .radius = 2, .axis = .vertical, .threads = .{ .x = 0, .y = 4 } }));
952 }