feat: AVX512 diff algorithm (#131)

* AVX diff

* Add avx CI run

* Change avx512_diff build flag to --enable-asm runtime flag

* Do not log unsupported windows OS for avx build

---------

Co-authored-by: Dmitriy Kovalenko <dmtr.kovalenko@outlook.com>

Author: serpent7776

.github/workflows/build.yml

@@ -86,6 +86,8 @@ jobs:
           path: binaries/odiff-macos-arm64
           retention-days: 14
 
+  
+
   test:
     name: Run Tests (all targets)
     runs-on: ${{ matrix.os }}
@@ -139,6 +141,43 @@ jobs:
         run: |
           echo "Skipping: no hosted runner available for ${{ matrix.target }}"
 
+  test-avx:
+    name: Run Tests (AVX x86_64)
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - target: x86_64-linux-gnu
+            os: ubuntu-latest
+          - target: x86_64-macos
+            os: macos-latest
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Setup Zig
+        uses: mlugg/setup-zig@v2
+        with:
+          version: 0.15.1
+
+      - name: Install nasm for x86 targets (Linux)
+        if: contains(matrix.target, 'x86_64') && runner.os == 'Linux'
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y nasm
+
+      - name: Install nasm for x86 targets (macOS)
+        if: contains(matrix.target, 'x86_64') && runner.os == 'macOS'
+        run: |
+          brew install nasm
+
+      - name: Run AVX tests
+        run: |
+          zig build test-all --summary all
+
   e2e-tests:
     name: End-to-end JavaScript tests
     needs: [build]

build.zig

@@ -10,7 +10,11 @@ pub fn build(b: *std.Build) !void {
     const is_cross_compiling = target.result.cpu.arch != native_target.result.cpu.arch or
         target.result.os.tag != native_target.result.os.tag;
 
-    const lib_mod, const exe = buildOdiff(b, target, optimize, dynamic);
+    const build_options = b.addOptions();
+    build_options.addOption([]const u8, "version", manifest.version);
+    const build_options_mod = build_options.createModule();
+
+    const lib_mod, const exe = buildOdiff(b, target, optimize, dynamic, build_options_mod);
     b.installArtifact(exe);
 
     const run_cmd = b.addRunArtifact(exe);
@@ -36,6 +40,7 @@ pub fn build(b: *std.Build) !void {
         "src/test_io_bmp.zig",
         "src/test_io_jpg.zig",
         "src/test_io_tiff.zig",
+        "src/test_avx.zig",
     };
 
     const integration_tests_pure_zig = [_][]const u8{
@@ -58,6 +63,7 @@ pub fn build(b: *std.Build) !void {
                     .optimize = optimize,
                 }),
             });
+            integration_test.root_module.addImport("build_options", build_options_mod);
             integration_test.linkLibC();
             integration_test.linkLibrary(root_lib);
             linkDeps(b, target, optimize, false, integration_test.root_module);
@@ -74,6 +80,7 @@ pub fn build(b: *std.Build) !void {
                     .optimize = optimize,
                 }),
             });
+            pure_test.root_module.addImport("build_options", build_options_mod);
 
             const run_pure_test = b.addRunArtifact(pure_test);
             integration_test_steps.append(run_pure_test) catch @panic("OOM");
@@ -95,7 +102,7 @@ pub fn build(b: *std.Build) !void {
     const build_ci_step = b.step("ci", "Build the app for CI");
     for (build_targets) |target_query| {
         const t = b.resolveTargetQuery(target_query);
-        _, const odiff_exe = buildOdiff(b, t, optimize, dynamic);
+        _, const odiff_exe = buildOdiff(b, t, optimize, dynamic, build_options_mod);
         odiff_exe.root_module.strip = true;
         const odiff_output = b.addInstallArtifact(odiff_exe, .{
             .dest_dir = .{
@@ -113,6 +120,7 @@ fn buildOdiff(
     target: std.Build.ResolvedTarget,
     optimize: std.builtin.OptimizeMode,
     dynamic: bool,
+    build_options_mod: *std.Build.Module,
 ) struct { *std.Build.Module, *std.Build.Step.Compile } {
     const lib_mod = b.createModule(.{
         .root_source_file = b.path("src/root.zig"),
@@ -146,11 +154,24 @@ fn buildOdiff(
     });
 
     exe_mod.addImport("odiff_lib", lib_mod);
-
-    const options = b.addOptions();
-    options.addOption([]const u8, "version", manifest.version);
-    exe_mod.addImport("build_options", options.createModule());
-    lib_mod.addImport("build_options", options.createModule());
+    exe_mod.addImport("build_options", build_options_mod);
+    lib_mod.addImport("build_options", build_options_mod);
+
+    if (target.result.cpu.arch == .x86_64) {
+        const os_tag = target.result.os.tag;
+        const fmt: ?[]const u8 = switch (os_tag) {
+            .linux => "elf64",
+            .macos => "macho64",
+            else => null,
+        };
+
+        if (fmt) |nasm_fmt| {
+            const nasm = b.addSystemCommand(&.{ "nasm", "-f", nasm_fmt, "-o" });
+            const asm_obj = nasm.addOutputFileArg("vxdiff.o");
+            nasm.addFileArg(b.path("src/vxdiff.asm"));
+            lib_mod.addObjectFile(asm_obj);
+        }
+    }
 
     const exe = b.addExecutable(.{
         .name = "odiff",

src/cli.zig

@@ -16,6 +16,7 @@ pub const CliArgs = struct {
     antialiasing: bool = false,
     diff_lines: bool = false,
     reduce_ram_usage: bool = false,
+    enable_asm: bool = false,
     ignore_regions: std.array_list.Managed(diff.IgnoreRegion),
     allocator: std.mem.Allocator,
 
@@ -50,6 +51,7 @@ fn printUsage(program_name: []const u8) void {
     print("  --aa, --antialiasing        Ignore antialiased pixels in diff\n", .{});
     print("  --output-diff-lines         Output line numbers with differences\n", .{});
     print("  --reduce-ram-usage          Use less memory (slower)\n", .{});
+    print("  --enable-asm                Enable AVX-512 optimized asm path when supported (x86_64 only)\n", .{});
     print("  -i, --ignore <regions>      Ignore regions (format: x1:y1-x2:y2,x3:y3-x4:y4)\n", .{});
     print("  -h, --help                  Show this help message\n", .{});
     print("  --version                   Show version\n", .{});
@@ -158,6 +160,8 @@ pub fn parseArgs(allocator: std.mem.Allocator) !CliArgs {
             parsed_args.diff_lines = true;
         } else if (std.mem.eql(u8, arg, "--reduce-ram-usage")) {
             parsed_args.reduce_ram_usage = true;
+        } else if (std.mem.eql(u8, arg, "--enable-asm")) {
+            parsed_args.enable_asm = true;
         } else if (std.mem.eql(u8, arg, "-i") or std.mem.eql(u8, arg, "--ignore")) {
             i += 1;
             if (i >= args.len) {

src/diff.zig

@@ -8,7 +8,11 @@ const antialiasing = @import("antialiasing.zig");
 const Image = image_io.Image;
 const ArrayList = std.ArrayList;
 
-const HAS_AVX512 = std.Target.x86.featureSetHas(builtin.cpu.features, .avx512f);
+const HAS_AVX512f = std.Target.x86.featureSetHas(builtin.cpu.features, .avx512f);
+const HAS_AVX512bwvl =
+    HAS_AVX512f and
+    std.Target.x86.featureSetHas(builtin.cpu.features, .avx512bw) and
+    std.Target.x86.featureSetHas(builtin.cpu.features, .avx512vl);
 const HAS_NEON = std.Target.aarch64.featureSetHas(builtin.cpu.features, .neon);
 
 const RED_PIXEL: u32 = 0xFF0000FF;
@@ -73,6 +77,7 @@ pub const DiffOptions = struct {
     ignore_regions: ?[]const IgnoreRegion = null,
     capture_diff: bool = true,
     fail_on_layout_change: bool = true,
+    enable_asm: bool = false,
 };
 
 fn unrollIgnoreRegions(width: u32, regions: ?[]const IgnoreRegion, allocator: std.mem.Allocator) !?[]struct { u32, u32 } {
@@ -133,7 +138,14 @@ pub noinline fn compare(
 
     const layout_difference = base.width != comp.width or base.height != comp.height;
 
-    if (layout_difference) {
+    // AVX diff only supports default options
+    const threshold_ok = @abs(options.threshold - 0.1) < 0.0000001;
+    const no_ignore_regions = options.ignore_regions == null or options.ignore_regions.?.len == 0;
+    const avx_compatible = !options.antialiasing and no_ignore_regions and !options.capture_diff and !options.diff_lines and threshold_ok;
+
+    if (options.enable_asm and HAS_AVX512bwvl and avx_compatible) {
+        try compareAVX(base, comp, &diff_count);
+    } else if (layout_difference) {
         // slow path for different layout or weird widths
         try compareDifferentLayouts(base, comp, &diff_output, &diff_count, if (diff_lines != null) &diff_lines.? else null, ignore_regions, max_delta_i64, options);
     } else {
@@ -217,7 +229,7 @@ pub noinline fn compareSameLayouts(base: *const Image, comp: *const Image, diff_
     const base_data = base.data;
     const comp_data = comp.data;
 
-    const SIMD_SIZE = std.simd.suggestVectorLength(u32) orelse if (HAS_AVX512) 16 else if (HAS_NEON) 8 else 4;
+    const SIMD_SIZE = std.simd.suggestVectorLength(u32) orelse if (HAS_AVX512f) 16 else if (HAS_NEON) 8 else 4;
     const simd_end = (size / SIMD_SIZE) * SIMD_SIZE;
 
     var offset: usize = 0;
@@ -328,6 +340,29 @@ pub fn compareDifferentLayouts(base: *const Image, comp: *const Image, maybe_dif
     }
 }
 
+pub fn compareAVX(base: *const Image, comp: *const Image, diff_count: *u32) !void {
+    if (!HAS_AVX512bwvl) return error.Invalid;
+
+    const base_ptr: [*]const u8 = @ptrCast(@alignCast(base.data.ptr));
+    const comp_ptr: [*]const u8 = @ptrCast(@alignCast(comp.data.ptr));
+
+    const base_w: usize = base.width;
+    const base_h: usize = base.height;
+    const comp_w: usize = comp.width;
+    const comp_h: usize = comp.height;
+
+    diff_count.* = vxdiff(base_ptr, comp_ptr, base_w, comp_w, base_h, comp_h);
+}
+
+extern fn vxdiff(
+    base_rgba: [*]const u8,
+    comp_rgba: [*]const u8,
+    base_width: usize,
+    comp_width: usize,
+    base_height: usize,
+    comp_height: usize,
+) u32;
+
 pub fn diff(
     base: *const Image,
     comp: *const Image,

src/main.zig

@@ -78,6 +78,7 @@ pub fn main() !void {
         .diff_lines = args.diff_lines,
         .ignore_regions = args.ignore_regions.items,
         .capture_diff = args.diff_output != null,
+        .enable_asm = args.enable_asm,
     };
 
     const result = diff.diff(&base_img, &comp_img, diff_options, allocator) catch |err| {

src/test_avx.zig

@@ -0,0 +1,66 @@
+const std = @import("std");
+const testing = std.testing;
+const expect = testing.expect;
+const expectEqual = testing.expectEqual;
+const expectApproxEqRel = testing.expectApproxEqRel;
+
+const odiff = @import("root.zig");
+const image_io = odiff.image_io;
+const diff = odiff.diff;
+const color_delta = odiff.color_delta;
+
+fn loadTestImage(path: []const u8, allocator: std.mem.Allocator) !image_io.Image {
+    return image_io.loadImage(path, allocator) catch |err| {
+        std.debug.print("Failed to load image: {s}\nError: {}\n", .{ path, err });
+        return err;
+    };
+}
+
+test "layoutDifference: diff images with different layouts without capture" {
+    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
+    defer _ = gpa.deinit();
+    const allocator = gpa.allocator();
+
+    var img1 = try loadTestImage("test/png/white4x4.png", allocator);
+    defer img1.deinit();
+
+    var img2 = try loadTestImage("test/png/purple8x8.png", allocator);
+    defer img2.deinit();
+
+    const options = diff.DiffOptions{
+        .antialiasing = false,
+        .output_diff_mask = false,
+        .capture_diff = false,
+        .enable_asm = true,
+    };
+
+    var diff_output, const diff_count, const diff_percentage, var diff_lines = try diff.compare(&img1, &img2, options, allocator);
+    defer if (diff_output) |*img| img.deinit();
+    defer if (diff_lines) |*lines| lines.deinit();
+
+    try expectEqual(@as(u32, 16), diff_count); // diffPixels
+    try expectApproxEqRel(@as(f64, 100.0), diff_percentage, 0.001); // diffPercentage
+}
+
+test "PNG: finds difference between 2 images without capture" {
+    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
+    defer _ = gpa.deinit();
+    const allocator = gpa.allocator();
+
+    var img1 = try loadTestImage("test/png/orange.png", allocator);
+    defer img1.deinit();
+
+    var img2 = try loadTestImage("test/png/orange_changed.png", allocator);
+    defer img2.deinit();
+
+    const options = diff.DiffOptions{
+        .capture_diff = false,
+        .enable_asm = true,
+    };
+    var diff_output, const diff_count, const diff_percentage, var diff_lines = try diff.compare(&img1, &img2, options, allocator);
+    defer if (diff_output) |*img| img.deinit();
+    defer if (diff_lines) |*lines| lines.deinit();
+
+    try expectEqual(@as(u32, 1366), diff_count); // diffPixels
+    try expectApproxEqRel(@as(f64, 1.14), diff_percentage, 0.1); // diffPercentage
+}