Research: shared memory for big images (issue #66), benchmarks and findings

benchmarks/java/FijiExportBench.java

@@ -0,0 +1,77 @@
+import ij.ImagePlus;
+import ij.io.FileSaver;
+import ij.process.ShortProcessor;
+import java.io.File;
+import java.io.RandomAccessFile;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.util.Random;
+
+/**
+ * What Fiji's host code pays to hand an open image to the worker, measured with the ImageJ that Fiji ships (ij 1.54p):
+ *   tiff_disk  FileSaver.saveAsTiff to a file (what LabConstrictor_Tools.groovy does today)
+ *   tiff_shmfs the same into /dev/shm (CPU cost of the encoder, no disk)
+ *   shm_appose Appose 0.12.0 SharedMemory + ShortBuffer.put(short[]) from the ImagePlus's pixel array (JNA)
+ *   tiff_read  ij.io.Opener.openImage of the file just written (a result coming back as TIFF)
+ *   shm_to_short_array  block -> new short[] (a result coming back through shared memory)
+ *   shm_mmap   JDK-only MappedByteBuffer over a file in /dev/shm
+ * Args: side count scratchDir. Image = side x side uint16 (side 16384 = 512 MiB). One short[] per plane, like ImageJ.
+ */
+public class FijiExportBench {
+    public static void main(String[] a) throws Exception {
+        int side = Integer.parseInt(a[0]);
+        int repeats = Integer.parseInt(a[1]);
+        File scratch = new File(a[2]);
+        short[] pixels = new short[side * side];
+        Random r = new Random(1);
+        for (int i = 0; i < pixels.length; i += 7) pixels[i] = (short) r.nextInt(65536);
+        ImagePlus imp = new ImagePlus("bench", new ShortProcessor(side, side, pixels, null));
+        long nbytes = (long) pixels.length * 2;
+        for (int k = 0; k < repeats; k++) {
+            StringBuilder out = new StringBuilder("{\"MiB\":" + nbytes / (1024 * 1024));
+            File tif = new File(scratch, "fiji_bench.tif");
+            long t = System.nanoTime();
+            new FileSaver(imp).saveAsTiff(tif.getPath());
+            out.append(",\"tiff_disk_s\":").append(sec(t));
+            t = System.nanoTime();
+            ImagePlus back = new ij.io.Opener().openImage(tif.getPath());     // what the Fiji host does with an image result
+            out.append(",\"tiff_read_s\":").append(sec(t));
+            if (back == null || back.getWidth() != side) throw new IllegalStateException("read back failed");
+            t = System.nanoTime();
+            org.apposed.appose.SharedMemory rb = org.apposed.appose.SharedMemory.create((long) side * side * 2);
+            rb.buf(0, (long) side * side * 2).order(ByteOrder.nativeOrder()).asShortBuffer().put((short[]) pixels);   // stand-in for the worker's block
+            double rbMake = (System.nanoTime() - t) / 1e9;
+            t = System.nanoTime();
+            short[] copy = new short[side * side];
+            rb.buf(0, (long) side * side * 2).order(ByteOrder.nativeOrder()).asShortBuffer().get(copy);                 // block -> ImageJ short[]
+            out.append(",\"shm_to_short_array_s\":").append(sec(t));
+            rb.close();
+            back.close();
+            tif.delete();
+            File tif2 = new File("/dev/shm/fiji_bench.tif");
+            t = System.nanoTime();
+            new FileSaver(imp).saveAsTiff(tif2.getPath());
+            out.append(",\"tiff_shmfs_s\":").append(sec(t));
+            tif2.delete();
+            t = System.nanoTime();
+            org.apposed.appose.SharedMemory shm = org.apposed.appose.SharedMemory.create(nbytes);
+            shm.buf(0, nbytes).order(ByteOrder.nativeOrder()).asShortBuffer().put((short[]) imp.getProcessor().getPixels());
+            out.append(",\"shm_appose_s\":").append(sec(t));
+            shm.close();
+            File f = new File("/dev/shm/fiji_bench_map");
+            t = System.nanoTime();
+            RandomAccessFile raf = new RandomAccessFile(f, "rw");
+            raf.setLength(nbytes);
+            MappedByteBuffer buf = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, nbytes);
+            buf.order(ByteOrder.nativeOrder()).asShortBuffer().put((short[]) imp.getProcessor().getPixels());
+            out.append(",\"shm_mmap_s\":").append(sec(t));
+            raf.close();
+            f.delete();
+            System.out.println(out.append("}"));
+        }
+    }
+
+    static String sec(long t0) { return String.format("%.3f", (System.nanoTime() - t0) / 1e9); }
+}

benchmarks/java/ShmProof.java

@@ -0,0 +1,65 @@
+import java.io.RandomAccessFile;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.nio.MappedByteBuffer;
+import java.nio.channels.FileChannel;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.Random;
+
+/**
+ * Java side of the shared-memory proofs (see docs/PERF_SHARED_MEMORY.md, host Fiji / QuPath).
+ *
+ *   mode "appose": the Appose 0.12.0 SharedMemory (JNA: shm_open + mmap on Linux/macOS, CreateFileMapping on Windows), the
+ *                  class Fiji already ships; compiled with --release 8.
+ *   mode "mmap"  : JDK only (Java 8+): FileChannel.map over a file in /dev/shm (RAM-backed on Linux).
+ *   mode "ffm"   : Java 21 Foreign Function and Memory API (preview in 21): shm_open + mmap without JNA. Compile with
+ *                  --enable-preview --release 21 and run with --enable-preview.
+ *
+ * Fills a short[] image (what ImageJ holds: one short[] per plane) into the block, prints one JSON line the Python side reads:
+ * {"mode","name","shape":[h,w],"dtype":"uint16","max":N,"fill_s":x}. The block stays alive until the Python side writes "done" to stdin.
+ */
+public class ShmProof {
+    static final int SIDE = Integer.getInteger("side", 4096);
+
+    public static void main(String[] args) throws Exception {
+        String mode = args[0];
+        int h = SIDE, w = SIDE;
+        short[] pixels = new short[h * w];
+        Random r = new Random(1);
+        int max = 0; long sum = 0;
+        for (int i = 0; i < pixels.length; i++) {
+            int v = r.nextInt(65536);               // above 255 and above 32767: exposes byte order and signedness mistakes
+            pixels[i] = (short) v;
+            max = Math.max(max, v); sum += v;
+        }
+        long nbytes = (long) pixels.length * 2;
+        String name;
+        long t0 = System.nanoTime();
+        AutoCloseable owner;
+        if (mode.equals("appose")) {
+            org.apposed.appose.SharedMemory shm = org.apposed.appose.SharedMemory.create(nbytes);
+            ByteBuffer buf = shm.buf(0, nbytes).order(ByteOrder.nativeOrder());
+            buf.asShortBuffer().put(pixels);
+            name = shm.name();
+            owner = shm;
+        } else if (mode.equals("mmap")) {
+            name = "lc_java_proof_" + java.util.UUID.randomUUID().toString().substring(0, 8);
+            Path file = Paths.get("/dev/shm", name);
+            RandomAccessFile raf = new RandomAccessFile(file.toFile(), "rw");
+            raf.setLength(nbytes);
+            MappedByteBuffer buf = raf.getChannel().map(FileChannel.MapMode.READ_WRITE, 0, nbytes);
+            buf.order(ByteOrder.nativeOrder()).asShortBuffer().put(pixels);
+            owner = () -> { raf.close(); Files.deleteIfExists(file); };
+        } else {
+            throw new IllegalArgumentException(mode);
+        }
+        double fill = (System.nanoTime() - t0) / 1e9;
+        System.out.println("{\"mode\":\"" + mode + "\",\"name\":\"" + name + "\",\"shape\":[" + h + "," + w
+            + "],\"dtype\":\"uint16\",\"max\":" + max + ",\"sum\":" + sum + ",\"first\":" + (pixels[0] & 0xFFFF) + ",\"fill_s\":" + fill + "}");
+        System.out.flush();
+        System.in.read();                           // the host keeps the block until the task is over
+        owner.close();
+    }
+}

benchmarks/java/ShmProofFfm.java

@@ -0,0 +1,48 @@
+import java.lang.foreign.*;
+import java.lang.invoke.MethodHandle;
+import java.nio.ByteOrder;
+import java.util.Random;
+
+/** Java 21 FFM (preview in 21, final in 22): shm_open + ftruncate + mmap straight from libc, no JNA, no JNI. POSIX only. */
+public class ShmProofFfm {
+    public static void main(String[] args) throws Throwable {
+        int side = Integer.getInteger("side", 4096);
+        long nbytes = (long) side * side * 2;
+        Linker linker = Linker.nativeLinker();
+        SymbolLookup libc = linker.defaultLookup();
+        // shm_open is variadic (mode is optional): declare the variadic part explicitly
+        MethodHandle shmOpen = linker.downcallHandle(libc.find("shm_open").orElseThrow(),
+            FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS, ValueLayout.JAVA_INT, ValueLayout.JAVA_INT),
+            Linker.Option.firstVariadicArg(2));
+        MethodHandle ftruncate = linker.downcallHandle(libc.find("ftruncate").orElseThrow(),
+            FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, ValueLayout.JAVA_LONG));
+        MethodHandle mmap = linker.downcallHandle(libc.find("mmap").orElseThrow(),
+            FunctionDescriptor.of(ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT,
+                ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, ValueLayout.JAVA_LONG));
+        MethodHandle shmUnlink = linker.downcallHandle(libc.find("shm_unlink").orElseThrow(),
+            FunctionDescriptor.of(ValueLayout.JAVA_INT, ValueLayout.ADDRESS));
+        String name = "/lc_ffm_" + ProcessHandle.current().pid();
+        try (Arena arena = Arena.ofConfined()) {
+            int fd = (int) shmOpen.invokeExact(arena.allocateUtf8String(name), 02 | 0100, 0600);   // O_RDWR | O_CREAT
+            if (fd < 0) throw new IllegalStateException("shm_open failed");
+            if ((int) ftruncate.invokeExact(fd, nbytes) != 0) throw new IllegalStateException("ftruncate failed");
+            MemorySegment base = (MemorySegment) mmap.invokeExact(MemorySegment.NULL, nbytes, 3, 1, fd, 0L); // PROT_RW, MAP_SHARED
+            MemorySegment block = base.reinterpret(nbytes);
+            Random r = new Random(1);
+            int max = 0; long sum = 0; int first = -1;
+            long t0 = System.nanoTime();
+            ValueLayout.OfShort le = ValueLayout.JAVA_SHORT.withOrder(ByteOrder.nativeOrder());
+            for (long i = 0; i < (long) side * side; i++) {
+                int v = r.nextInt(65536);
+                max = Math.max(max, v); sum += v; if (first < 0) first = v;
+                block.setAtIndex(le, i, (short) v);
+            }
+            double fill = (System.nanoTime() - t0) / 1e9;
+            System.out.println("{\"mode\":\"ffm\",\"name\":\"" + name.substring(1) + "\",\"shape\":[" + side + "," + side
+                + "],\"dtype\":\"uint16\",\"max\":" + max + ",\"sum\":" + sum + ",\"first\":" + first + ",\"fill_s\":" + fill + "}");
+            System.out.flush();
+            System.in.read();
+            int rc = (int) shmUnlink.invokeExact(arena.allocateUtf8String(name));
+        }
+    }
+}

benchmarks/java_shm_proof.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python3
+"""Run a Java process that creates a shared block with a short[] image, then let the REAL LabConstrictor worker read it.
+
+    python benchmarks/java_shm_proof.py --mode appose --classpath DIR_WITH_COMPILED_CLASSES:jna.jar:jna-platform.jar
+    python benchmarks/java_shm_proof.py --mode mmap   --classpath DIR
+    python benchmarks/java_shm_proof.py --mode ffm    --classpath DIR   (ShmProofFfm, Java 21 --enable-preview)
+
+Passes when the worker's `image_checksum` equals the sum and first pixel Java computed: the pixels, their byte order and their unsigned values
+survived Java heap -> block -> worker copy. Exit code 0 = pass.
+"""
+
+import argparse
+import json
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE.parent))
+from labconstrictor_tools.client import WorkerProcess
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--mode", choices=["appose", "mmap", "ffm"], required=True)
+    ap.add_argument("--classpath", required=True)
+    ap.add_argument("--side", type=int, default=4096)
+    args = ap.parse_args()
+    cmd = ["java", "-Dside=%d" % args.side, "-Xmx3g"]
+    main_class = "ShmProof"
+    if args.mode == "ffm":
+        cmd += ["--enable-preview"]
+        main_class = "ShmProofFfm"
+    cmd += ["-cp", args.classpath, main_class] + ([args.mode] if args.mode != "ffm" else [])
+    java = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True)
+    try:
+        line = java.stdout.readline()
+        info = json.loads(line)
+        descriptor = {
+            "appose_type": "ndarray",
+            "shm": {"name": info["name"]},
+            "shape": info["shape"],
+            "dtype": info["dtype"],
+        }
+        with WorkerProcess(module="shm_bench_app", pythonpath=[str(HERE)]) as worker:
+            task = worker.task("image_checksum", {"image": descriptor, "_job_dir": tempfile.mkdtemp()}).wait(
+                60
+            )
+        values = next(r["values"] for r in task.outputs["results"] if r["type"] == "values")
+        ok = (
+            values["sum"] == info["sum"]
+            and values["first"] == info["first"]
+            and values["shape"] == info["shape"]
+        )
+        print(
+            "%s: java sum %d first %d, worker sum %d first %d, java fill %.3f s -> %s"
+            % (
+                info["mode"],
+                info["sum"],
+                info["first"],
+                values["sum"],
+                values["first"],
+                info["fill_s"],
+                "PASS" if ok else "FAIL",
+            )
+        )
+        return 0 if ok else 1
+    finally:
+        java.stdin.write("done\n")
+        java.stdin.close()
+        java.wait(10)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

benchmarks/results/sandbox_1024_invert_disk_then_oom.txt

@@ -0,0 +1,6 @@
+== 1024 MiB ((23170, 23170), uint16)
+  invert    disk  run1  export 2.75  task 11.58 (load 5.74 tool 5.33 write 0.50 other 0.008)  back 0.00  total 15.29  worker faults 2888 rss 2079 MiB  psi 31ms scan 0
+  invert    disk  run2  export 1.20  task 4.63 (load 1.41 tool 2.76 write 0.45 other 0.010)  back 0.00  total 6.46  worker faults 2028 rss 2079 MiB  psi 8ms scan 0
+  invert    disk  run3  export 0.53  task 2.79 (load 0.39 tool 1.89 write 0.51 other 0.008)  back 0.00  total 3.94  worker faults 2027 rss 2079 MiB  psi 47ms scan 71
+/usr/lib/python3.11/multiprocessing/resource_tracker.py:254: UserWarning: resource_tracker: There appear to be 1 leaked shared_memory objects to clean up at shutdown
+  warnings.warn('resource_tracker: There appear to be %d '