Run image atomics natively on Metal where the runtime allows (#2226)

build.gradle.kts

@@ -148,6 +148,17 @@ fun Test.configureAngelicaJava8() {
     }
 }
 
+if (isMacOs && providers.gradleProperty("clientJvmVendor").isPresent) {
+    for (javaVersion in listOf(17, 21, 25)) {
+        tasks.named<JavaExec>("runClient$javaVersion") {
+            javaLauncher = javaToolchains.launcherFor {
+                languageVersion = JavaLanguageVersion.of(javaVersion)
+                vendor = JvmVendorSpec.matching(providers.gradleProperty("clientJvmVendor").get())
+            }
+        }
+    }
+}
+
 tasks.test {
     useJUnitPlatform { excludeTags = setOf("gl-core") }
     configureAngelicaJava8()

glsm/src/main/java/com/gtnewhorizons/angelica/config/SystemProperties.java

@@ -62,6 +62,7 @@ public final class SystemProperties {
     public static final long FLYBY_TIME_OF_DAY = Math.floorMod(Long.getLong("angelica.flyby.timeOfDay", 6000L), 24000L);
     public static final String FLYBY_COMMANDS = System.getProperty("angelica.flyby.commands", "");
     public static final String FLYBY_ORIGIN = System.getProperty("angelica.flyby.origin", "");
+    public static final Integer FLYBY_DIMENSION = Integer.getInteger("angelica.flyby.dimension");
     public static final FlybyPacing FLYBY_PACING = parseEnum("angelica.flyby.pacing", FlybyPacing.UNCAPPED, FlybyPacing.UNCAPPED, FlybyPacing.class);
     public static final FlybyWeather FLYBY_WEATHER = parseEnum("angelica.flyby.weather", FlybyWeather.CLEAR, FlybyWeather.CLEAR, FlybyWeather.class);
     public static final FlybyCamera FLYBY_CAMERA = parseEnum("angelica.flyby.camera", FlybyCamera.FIRST_PERSON, FlybyCamera.FIRST_PERSON, FlybyCamera.class);

glsm/src/main/java/com/gtnewhorizons/angelica/glsm/backend/RenderBackend.java

@@ -280,6 +280,8 @@ public abstract class RenderBackend {
 
     public boolean supportsComputeImageAtomics() { return true; }
 
+    public boolean supportsComputeImageAtomics(int width, int height, int depth) { return supportsComputeImageAtomics(); }
+
     public abstract void enable(int cap);
     public abstract void enablei(int cap, int index);
     public abstract void disable(int cap);

scripts/flyby-capture.sh

@@ -10,6 +10,8 @@
 #   length - Units: blocks (straight), blocks per leg (circuit), degrees (pan), ticks (static).
 #   extra  - passed through to gradle; -Dangelica.* reaches the client (e.g. -Dangelica.flyby.timeOfDay=1500).
 #            -Dangelica.flyby.camera=FIRST_PERSON|THIRD_PERSON_BACK|THIRD_PERSON_FRONT (default FIRST_PERSON).
+#            -Dangelica.flyby.dimension=<id> moves the player there first (1 = End), after world saving is off.
+#            -PclientJvmVendor=adoptium (macOS) picks the JVM vendor for runClient17/21/25.
 #
 #   FLYBY_SPEED     - blocks per tick. (0.5 ~= creative flight).
 #   FLYBY_PACING    - UNCAPPED | CONFIGURED. Passed through when set.

scripts/flyby-scenes/end-crystals.txt

@@ -0,0 +1,4 @@
+# End crystals next to the arrival platform. Use with -Dangelica.flyby.dimension=1, route pan.
+summon EnderCrystal ~ ~2 ~12 {ForgeData:{flyby:1b}}
+summon EnderCrystal ~8 ~2 ~8 {ForgeData:{flyby:1b}}
+summon EnderCrystal ~-8 ~4 ~10 {ForgeData:{flyby:1b}}

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/SDLGPURenderBackend.java

@@ -4448,6 +4448,10 @@ public class SDLGPURenderBackend extends RenderBackend {
         return device.supportsSpirv() || !device.supportsMsl();
     }
 
+    @Override public boolean supportsComputeImageAtomics(int width, int height, int depth) {
+        return supportsComputeImageAtomics() || device.supportsImageAtomics(width, height, depth);
+    }
+
     private void prepareVertexWrites(ContextState st) {
         final ShaderManager.ProgramObject prog = st.boundProgramObj;
         if (prog != null && prog.vertexWriteReplay != null) vertexWriteReplayer.ensureImageUsage(prog.vertexWriteReplay);

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/compute/ImageAtomicsProbe.java

@@ -0,0 +1,152 @@
+package com.gtnewhorizons.angelica.sdlgpu.compute;
+
+import com.gtnewhorizons.angelica.glsm.shader.SpirvCompiler;
+import com.gtnewhorizons.angelica.sdlgpu.shader.ShaderManager;
+import com.gtnewhorizons.angelica.sdlgpu.shader.cross.CrossCompileCache;
+import com.gtnewhorizons.angelica.sdlgpu.shader.msl.MslCrossCompile;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.embeddedt.embeddium.impl.render.shader.ShaderLoader;
+import org.lwjgl.opengl.GL43;
+import org.lwjgl.sdl.SDLError;
+import org.lwjgl.sdl.SDL_GPUComputePipelineCreateInfo;
+import org.lwjgl.sdl.SDL_GPUStorageTextureReadWriteBinding;
+import org.lwjgl.sdl.SDL_GPUTextureCreateInfo;
+import org.lwjgl.sdl.SDL_GPUTextureRegion;
+import org.lwjgl.sdl.SDL_GPUTextureTransferInfo;
+import org.lwjgl.sdl.SDL_GPUTransferBufferCreateInfo;
+import org.lwjgl.system.JNI;
+import org.lwjgl.system.MemoryStack;
+import org.lwjgl.system.MemoryUtil;
+import org.lwjgl.system.macosx.ObjCRuntime;
+import org.lwjgl.util.shaderc.Shaderc;
+
+import static org.lwjgl.sdl.SDLGPU.*;
+
+public final class ImageAtomicsProbe {
+    private static final Logger LOG = LogManager.getLogger("Angelica-SDLGPU");
+
+    private static final String SOURCE = "angelica:sdlgpu/atomic_probe.csh";
+    public static final int MTL_LANGUAGE_VERSION_3_1 = (3 << 16) + 1;
+    private static final int GROUPS = 64;
+
+    private final long device;
+    private long pipeline;
+
+    public ImageAtomicsProbe(long device) {
+        this.device = device;
+        final SpirvCompiler.Result compiled = SpirvCompiler.compile(ShaderLoader.getShaderSource(SOURCE), Shaderc.shaderc_compute_shader, "atomic_probe", SpirvCompiler.Options.vulkanForced460Core(), false);
+        if (compiled.spirv() == null) {
+            LOG.warn("Metal image atomics: probe pipeline failed: {}", compiled.error());
+            return;
+        }
+        CrossCompileCache.Output msl = null;
+        try (MemoryStack stack = MemoryStack.stackPush()) {
+            ShaderManager.remapSpirvForComputeSDLGPU(compiled.spirv());
+            msl = MslCrossCompile.compileAtVersion(compiled.spirv(), GL43.GL_COMPUTE_SHADER, MslCrossCompile.SPVC_MSL_3_1);
+            pipeline = SDL_CreateGPUComputePipeline(device, SDL_GPUComputePipelineCreateInfo.calloc(stack)
+                .code(msl.code())
+                .entrypoint(stack.ASCII(msl.entrypoint()))
+                .format(SDL_GPU_SHADERFORMAT_MSL)
+                .num_readwrite_storage_textures(1)
+                .threadcount_x(InvocationDispatch.WORKGROUP_SIZE).threadcount_y(1).threadcount_z(1));
+            if (pipeline == 0) LOG.warn("Metal image atomics: probe pipeline failed: {}", SDLError.SDL_GetError());
+        } catch (RuntimeException e) {
+            LOG.warn("Metal image atomics: probe pipeline failed", e);
+        } finally {
+            if (msl != null) MemoryUtil.memFree(msl.code());
+            MemoryUtil.memFree(compiled.spirv());
+        }
+    }
+
+    public static int defaultMslVersion() {
+        try {
+            final long msgSend = ObjCRuntime.getLibrary().getFunctionAddress("objc_msgSend");
+            final long cls = ObjCRuntime.objc_getClass("MTLCompileOptions");
+            if (msgSend == 0 || cls == 0) return 0;
+            final long options = JNI.invokePPP(cls, ObjCRuntime.sel_getUid("new"), msgSend);
+            if (options == 0) return 0;
+            final long version = JNI.invokePPP(options, ObjCRuntime.sel_getUid("languageVersion"), msgSend);
+            JNI.invokePPV(options, ObjCRuntime.sel_getUid("release"), msgSend);
+            return (int) version;
+        } catch (RuntimeException | LinkageError e) {
+            LOG.warn("Could not read the default Metal language version", e);
+            return 0;
+        }
+    }
+
+    public boolean run(int width, int height) {
+        if (pipeline == 0) return false;
+        long texture = 0, upload = 0, download = 0, cb = 0;
+        try (MemoryStack stack = MemoryStack.stackPush()) {
+            texture = SDL_CreateGPUTexture(device, SDL_GPUTextureCreateInfo.calloc(stack)
+                .type(SDL_GPU_TEXTURETYPE_2D)
+                .format(SDL_GPU_TEXTUREFORMAT_R32_INT)
+                .usage(SDL_GPU_TEXTUREUSAGE_SAMPLER | SDL_GPU_TEXTUREUSAGE_COMPUTE_STORAGE_WRITE)
+                .width(width).height(height).layer_count_or_depth(1).num_levels(1)
+                .sample_count(SDL_GPU_SAMPLECOUNT_1));
+            if (texture == 0) return failSdl(width, height, "texture");
+
+            upload = SDL_CreateGPUTransferBuffer(device, SDL_GPUTransferBufferCreateInfo.calloc(stack).usage(SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD).size(Integer.BYTES));
+            download = SDL_CreateGPUTransferBuffer(device, SDL_GPUTransferBufferCreateInfo.calloc(stack).usage(SDL_GPU_TRANSFERBUFFERUSAGE_DOWNLOAD).size(Integer.BYTES));
+            if (upload == 0 || download == 0) return failSdl(width, height, "transfer buffer");
+            final long zero = nSDL_MapGPUTransferBuffer(device, upload, false);
+            if (zero == 0) return failSdl(width, height, "map");
+            MemoryUtil.memPutInt(zero, 0);
+            SDL_UnmapGPUTransferBuffer(device, upload);
+
+            cb = SDL_AcquireGPUCommandBuffer(device);
+            if (cb == 0) return failSdl(width, height, "command buffer");
+            final SDL_GPUTextureRegion region = SDL_GPUTextureRegion.calloc(stack).texture(texture).w(1).h(1).d(1);
+            long copy = SDL_BeginGPUCopyPass(cb);
+            SDL_UploadToGPUTexture(copy, SDL_GPUTextureTransferInfo.calloc(stack).transfer_buffer(upload), region, false);
+            SDL_EndGPUCopyPass(copy);
+
+            final SDL_GPUStorageTextureReadWriteBinding.Buffer binding = SDL_GPUStorageTextureReadWriteBinding.calloc(1, stack);
+            binding.get(0).texture(texture);
+            final long pass = SDL_BeginGPUComputePass(cb, binding, null);
+            if (pass == 0) return failSdl(width, height, "compute pass");
+            SDL_BindGPUComputePipeline(pass, pipeline);
+            SDL_DispatchGPUCompute(pass, GROUPS, 1, 1);
+            SDL_EndGPUComputePass(pass);
+
+            copy = SDL_BeginGPUCopyPass(cb);
+            SDL_DownloadFromGPUTexture(copy, region, SDL_GPUTextureTransferInfo.calloc(stack).transfer_buffer(download));
+            SDL_EndGPUCopyPass(copy);
+            final boolean submitted = SDL_SubmitGPUCommandBuffer(cb);
+            cb = 0;
+            if (!submitted) return failSdl(width, height, "submit");
+            SDL_WaitForGPUIdle(device);
+
+            final long result = nSDL_MapGPUTransferBuffer(device, download, false);
+            if (result == 0) return failSdl(width, height, "map");
+            final int adds = MemoryUtil.memGetInt(result);
+            SDL_UnmapGPUTransferBuffer(device, download);
+
+            final int invocations = GROUPS * InvocationDispatch.WORKGROUP_SIZE;
+            if (adds != invocations) LOG.info("Metal image atomics on {}x{}: {} invocations added {}", width, height, invocations, adds);
+            return adds == invocations;
+        } catch (RuntimeException e) {
+            return fail(width, height, e.toString());
+        } finally {
+            if (cb != 0) SDL_CancelGPUCommandBuffer(cb);
+            if (download != 0) SDL_ReleaseGPUTransferBuffer(device, download);
+            if (upload != 0) SDL_ReleaseGPUTransferBuffer(device, upload);
+            if (texture != 0) SDL_ReleaseGPUTexture(device, texture);
+        }
+    }
+
+    public void release() {
+        if (pipeline != 0) SDL_ReleaseGPUComputePipeline(device, pipeline);
+        pipeline = 0;
+    }
+
+    private static boolean failSdl(int width, int height, String what) {
+        return fail(width, height, what + " failed: " + SDLError.SDL_GetError());
+    }
+
+    private static boolean fail(int width, int height, String reason) {
+        LOG.warn("Metal image atomics on {}x{}: {}", width, height, reason);
+        return false;
+    }
+}

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/device/Device.java

@@ -5,7 +5,9 @@ import com.gtnewhorizons.angelica.glsm.backend.BackendOptions;
 import com.gtnewhorizons.angelica.glsm.backend.MoltenVK;
 import com.gtnewhorizons.angelica.glsm.backend.VSyncMode;
 import com.gtnewhorizons.angelica.glsm.backend.RenderBackend;
+import com.gtnewhorizons.angelica.sdlgpu.compute.ImageAtomicsProbe;
 import com.gtnewhorizons.angelica.sdlgpu.util.DebugMessageRelay;
+import it.unimi.dsi.fastutil.longs.Long2BooleanOpenHashMap;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.lwjgl.sdl.SDLError;
@@ -49,6 +51,9 @@ public final class Device {
     private String driverName;
     private boolean fencePollEnabled;
     private boolean fenceQueryInverted;
+    private int defaultMslVersion;
+    private ImageAtomicsProbe imageAtomicsProbe;
+    private final Long2BooleanOpenHashMap imageAtomicsBySize = new Long2BooleanOpenHashMap();
     private final FenceReleaser fenceReleaser = new FenceReleaser(fence -> device == 0 || FenceWait.isSignaled(this, fence), fence -> { if (device != 0) SDL_ReleaseGPUFence(device, fence); });
     private String deviceName;
     private String driverVersion;
@@ -153,10 +158,36 @@ public final class Device {
         }
 
         probeFences(ver);
+        readMetalLanguageVersion();
 
         return true;
     }
 
+    private void readMetalLanguageVersion() {
+        if (!isMetal(driverName)) return;
+        defaultMslVersion = ImageAtomicsProbe.defaultMslVersion();
+        if (metalTextureAtomics()) {
+            LOG.info("Metal image atomics: MSL 3.1 available, checked per image size");
+        } else {
+            LOG.info("Metal image atomics: unavailable, this Java runtime compiles Metal shaders as MSL {}.{} and they need 3.1. A runtime built with the macOS 14 SDK or newer enables them.", defaultMslVersion >>> 16, defaultMslVersion & 0xFFFF);
+        }
+    }
+
+    public boolean metalTextureAtomics() {
+        return defaultMslVersion >= ImageAtomicsProbe.MTL_LANGUAGE_VERSION_3_1;
+    }
+
+    public boolean supportsImageAtomics(int width, int height, int depth) {
+        if (!metalTextureAtomics() || width <= 0 || height <= 0 || depth > 1) return false;
+        final long key = (long) width << 32 | height;
+        if (imageAtomicsBySize.containsKey(key)) return imageAtomicsBySize.get(key);
+        if (imageAtomicsProbe == null) imageAtomicsProbe = new ImageAtomicsProbe(device);
+        final boolean ok = imageAtomicsProbe.run(width, height);
+        if (ok) LOG.info("Metal image atomics on {}x{}: ok", width, height);
+        imageAtomicsBySize.put(key, ok);
+        return ok;
+    }
+
     private void probeFences(int ver) {
         fencePollEnabled = FenceWait.pollEnabled(driverName, ver);
         final boolean expectInverted = isMetal(driverName) && ver == SDLVersion.SDL_VERSIONNUM(3, 4, 14);
@@ -256,11 +287,15 @@ public final class Device {
             claimed = false;
         }
         fenceReleaser.releaseAll();
+        if (imageAtomicsProbe != null) imageAtomicsProbe.release();
+        imageAtomicsProbe = null;
         SDL_DestroyGPUDevice(device);
         device = 0;
         driverName = null;
         fencePollEnabled = false;
         fenceQueryInverted = false;
+        defaultMslVersion = 0;
+        imageAtomicsBySize.clear();
         deviceName = null;
         driverVersion = null;
         driverInfo = null;

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/shader/ShaderManager.java

@@ -1162,7 +1162,7 @@ public final class ShaderManager {
             return new CrossCompiled(SDL_GPU_SHADERFORMAT_SPIRV, spirv, "main", false);
         }
         if (device.supportsMsl()) {
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, glShaderType);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, glShaderType, device.metalTextureAtomics());
             return new CrossCompiled(SDL_GPU_SHADERFORMAT_MSL, out.code(), out.entrypoint(), true);
         }
         if (device.supportsDxbc()) {

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/shader/cross/CrossCompileUtil.java

@@ -150,18 +150,30 @@ public final class CrossCompileUtil {
         return false;
     }
 
+    public static boolean usesImageAtomics(ByteBuffer spirv) {
+        final IntBuffer w = spirv.asIntBuffer();
+        for (int i = SPIRV_HEADER_WORDS; i < w.limit(); i = nextInstruction(w, i)) {
+            if ((w.get(i) & 0xFFFF) == Spv.SpvOpImageTexelPointer) return true;
+        }
+        return false;
+    }
+
+    private static final int SPIRV_HEADER_WORDS = 5; // magic, version, generator, bound, schema
+
+    private static int nextInstruction(IntBuffer w, int i) {
+        final int len = w.get(i) >>> 16;
+        return len == 0 ? w.limit() : i + len;
+    }
+
     public static int[] parseLocalSize(ByteBuffer spirv) {
         final IntBuffer w = spirv.asIntBuffer();
-        final int wordCount = w.limit();
-        int i = 5; // 5-word SPIR-V header (magic, version, generator, bound, schema)
         int lx = 1, ly = 1, lz = 1;
         boolean foundLiteral = false;
         boolean sawExecutionModeId = false;
-        while (i < wordCount) {
+        for (int i = SPIRV_HEADER_WORDS; i < w.limit(); i = nextInstruction(w, i)) {
             final int word0 = w.get(i);
             final int opcode = word0 & 0xFFFF;
-            final int instrLen = (word0 >>> 16) & 0xFFFF;
-            if (instrLen == 0) break;
+            final int instrLen = word0 >>> 16;
             if (opcode == 16 && instrLen >= 6) { // SpvOpExecutionMode
                 final int mode = w.get(i + 2);
                 if (mode == 17) { // SpvExecutionModeLocalSize
@@ -173,7 +185,6 @@ public final class CrossCompileUtil {
             } else if (opcode == 331) { // SpvOpExecutionModeId
                 sawExecutionModeId = true;
             }
-            i += instrLen;
         }
         if (!foundLiteral && sawExecutionModeId) {
             LOG.warn("SPIR-V uses OpExecutionModeId (spec-constant LocalSizeId); falling back to {{1,1,1}} - not supported by parseLocalSize");

sdl-gpu/src/main/java/com/gtnewhorizons/angelica/sdlgpu/shader/msl/MslCrossCompile.java

@@ -26,7 +26,11 @@ public final class MslCrossCompile {
     private static final Logger LOG = LogManager.getLogger("Angelica-SDLGPU");
     private static final String BACKEND = "MSL";
 
-    private static final CrossCompileCache CACHE = new CrossCompileCache("msl", CrossCompileUtil::spvcId, MslCrossCompile::crossCompileUncached);
+    private static final int SPVC_MSL_2_2 = mslVersion(2, 2, 0);
+    public static final int SPVC_MSL_3_1 = mslVersion(3, 1, 0);
+
+    private static final CrossCompileCache CACHE = new CrossCompileCache("msl", CrossCompileUtil::spvcId, (spirv, glShaderType) -> compileAtVersion(spirv, glShaderType, SPVC_MSL_2_2));
+    private static final CrossCompileCache CACHE_3_1 = new CrossCompileCache("msl31", CrossCompileUtil::spvcId, (spirv, glShaderType) -> compileAtVersion(spirv, glShaderType, SPVC_MSL_3_1));
 
     @FunctionalInterface
     private interface BufferSlotFn { public int slot(int set, int binding); }
@@ -36,8 +40,9 @@ public final class MslCrossCompile {
 
     private MslCrossCompile() {}
 
-    public static Output compile(ByteBuffer spirv, int glShaderType) {
-        return CACHE.compile(spirv, glShaderType);
+    public static Output compile(ByteBuffer spirv, int glShaderType, boolean textureAtomics) {
+        final boolean atomics = textureAtomics && glShaderType == GL43.GL_COMPUTE_SHADER && CrossCompileUtil.usesImageAtomics(spirv);
+        return (atomics ? CACHE_3_1 : CACHE).compile(spirv, glShaderType);
     }
 
     private static int executionModelFor(int glShaderType) {
@@ -47,7 +52,7 @@ public final class MslCrossCompile {
         throw new UnsupportedOperationException("Unsupported shader type for MSL cross-compile: 0x" + Integer.toHexString(glShaderType));
     }
 
-    private static Output crossCompileUncached(ByteBuffer spirv, int glShaderType) {
+    public static Output compileAtVersion(ByteBuffer spirv, int glShaderType, int version) {
         final int executionModel = executionModelFor(glShaderType);
         final int dumpId = SystemProperties.dumpShaders() ? CrossCompileUtil.SHADER_DUMP_COUNTER.getAndIncrement() : -1;
         if (dumpId >= 0) CrossCompileUtil.dumpSpirv(spirv, dumpId, glShaderType);
@@ -77,7 +82,7 @@ public final class MslCrossCompile {
                     throw CrossCompileUtil.spvcError(ctx, BACKEND, "create_compiler_options");
                 }
                 final long opts = pOpts.get(0);
-                Spvc.spvc_compiler_options_set_uint(opts, Spvc.SPVC_COMPILER_OPTION_MSL_VERSION, mslVersion(2, 2, 0));
+                Spvc.spvc_compiler_options_set_uint(opts, Spvc.SPVC_COMPILER_OPTION_MSL_VERSION, version);
                 Spvc.spvc_compiler_options_set_uint(opts, Spvc.SPVC_COMPILER_OPTION_MSL_PLATFORM, Spvc.SPVC_MSL_PLATFORM_MACOS);
                 if (Spvc.spvc_compiler_install_compiler_options(compiler, opts) != Spvc.SPVC_SUCCESS) {
                     throw CrossCompileUtil.spvcError(ctx, BACKEND, "install_compiler_options");

sdl-gpu/src/main/resources/assets/angelica/shaders/sdlgpu/atomic_probe.csh

@@ -0,0 +1,7 @@
+#version 460 core
+layout(local_size_x = 64) in;
+layout(binding = 0, r32i) uniform iimage2D counter;
+
+void main() {
+    imageAtomicAdd(counter, ivec2(0), 1);
+}

sdl-gpu/src/test/java/com/gtnewhorizons/angelica/sdlgpu/compute/ImageAtomicsGpuTest.java

@@ -0,0 +1,35 @@
+package com.gtnewhorizons.angelica.sdlgpu.compute;
+
+import com.gtnewhorizons.angelica.sdlgpu.SdlTestRig;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+import static org.lwjgl.sdl.SDLGPU.SDL_GetGPUDeviceDriver;
+
+class ImageAtomicsGpuTest {
+
+    private static SdlTestRig rig;
+    private static ImageAtomicsProbe probe;
+
+    @BeforeAll
+    static void setUp() throws Exception {
+        rig = SdlTestRig.acquireRealDevice();
+        assumeTrue("metal".equalsIgnoreCase(SDL_GetGPUDeviceDriver(rig.sdlHandle)), "native Metal only");
+        assumeTrue(ImageAtomicsProbe.defaultMslVersion() >= ImageAtomicsProbe.MTL_LANGUAGE_VERSION_3_1, "this process compiles Metal source below MSL 3.1");
+        probe = new ImageAtomicsProbe(rig.sdlHandle);
+    }
+
+    @AfterAll
+    static void tearDown() {
+        if (probe != null) probe.release();
+        SdlTestRig.releaseRealDevice();
+    }
+
+    @Test
+    void atomicsCountEveryInvocationOnAnImageTheSizeOfEuphoriasEndCrystalMap() {
+        assertTrue(probe.run(40, 10));
+    }
+}

sdl-gpu/src/test/java/com/gtnewhorizons/angelica/sdlgpu/device/DeviceImageAtomicsGpuTest.java

@@ -0,0 +1,62 @@
+package com.gtnewhorizons.angelica.sdlgpu.device;
+
+import com.gtnewhorizons.angelica.glsm.testutil.Reflect;
+import com.gtnewhorizons.angelica.sdlgpu.SdlTestRig;
+import it.unimi.dsi.fastutil.longs.Long2BooleanOpenHashMap;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.junit.jupiter.api.Assumptions.assumeTrue;
+
+class DeviceImageAtomicsGpuTest {
+
+    private Device device;
+
+    @BeforeAll
+    static void initSdl() {
+        SdlTestRig.initSdlVideo();
+    }
+
+    @AfterEach
+    void tearDown() {
+        if (device != null) device.destroyDevice();
+    }
+
+    private Long2BooleanOpenHashMap verdicts() {
+        return Reflect.get(device, "imageAtomicsBySize");
+    }
+
+    @Test
+    void verdictsAreCachedPerSizeAndDroppedWithTheDevice() {
+        device = new Device();
+        assumeTrue(device.createDevice(), "no SDL GPU device on this machine");
+        assumeTrue(device.metalTextureAtomics(), "needs native Metal compiling at MSL 3.1 or newer");
+
+        assertFalse(device.supportsImageAtomics(40, 10, 2));
+        assertFalse(device.supportsImageAtomics(0, 10, 1));
+        assertFalse(device.supportsImageAtomics(40, -1, 1));
+        assertEquals(0, verdicts().size());
+
+        assertNull(Reflect.get(device, "imageAtomicsProbe"));
+
+        assertTrue(device.supportsImageAtomics(40, 10, 1));
+        assertFalse(device.supportsImageAtomics(64, 64, 1));
+        assertEquals(2, verdicts().size());
+        final Object probe = Reflect.get(device, "imageAtomicsProbe");
+        assertTrue(device.supportsImageAtomics(40, 10, 1));
+        assertFalse(device.supportsImageAtomics(64, 64, 1));
+        assertEquals(2, verdicts().size());
+        assertEquals(probe, Reflect.get(device, "imageAtomicsProbe"));
+
+        device.destroyDevice();
+        assertFalse(device.metalTextureAtomics());
+        assertFalse(device.supportsImageAtomics(40, 10, 1));
+        assertEquals(0, verdicts().size());
+        assertNull(Reflect.get(device, "imageAtomicsProbe"));
+    }
+}

sdl-gpu/src/test/java/com/gtnewhorizons/angelica/sdlgpu/shader/msl/FragmentStorageTextureSlotTest.java

@@ -57,7 +57,7 @@ class FragmentStorageTextureSlotTest {
         final ByteBuffer spirv = r.spirv();
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_FRAGMENT_SHADER);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER, false);
             try {
                 final String msl = StandardCharsets.UTF_8.decode(out.code().duplicate()).toString();
                 final boolean[] seen = new boolean[samplerCount];
@@ -136,7 +136,7 @@ class FragmentStorageTextureSlotTest {
         final ByteBuffer spirv = Reflect.get(obj, "spirv");
         final ShaderManager.StageReflection refl = Reflect.get(obj, "reflection");
 
-        final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER);
+        final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER, false);
         try {
             final String msl = StandardCharsets.UTF_8.decode(out.code().duplicate()).toString();
             final int samplerCount = refl.counts().numSamplers();

sdl-gpu/src/test/java/com/gtnewhorizons/angelica/sdlgpu/shader/msl/MslCrossCompileTest.java

@@ -8,6 +8,7 @@ import com.gtnewhorizons.angelica.glsm.testutil.Reflect;
 import com.gtnewhorizons.angelica.sdlgpu.shader.ShaderManager;
 import com.gtnewhorizons.angelica.sdlgpu.shader.SpirvTestShaders;
 import com.gtnewhorizons.angelica.sdlgpu.shader.cross.CrossCompileCache;
+import com.gtnewhorizons.angelica.sdlgpu.shader.cross.CrossCompileUtil;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
@@ -52,7 +53,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(SpirvTestShaders.VERTEX_GLSL, Shaderc.shaderc_vertex_shader);
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_VERTEX_SHADER);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.length() > 0, "MSL output non-empty");
@@ -75,7 +76,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(SpirvTestShaders.FRAGMENT_GLSL, Shaderc.shaderc_fragment_shader);
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_FRAGMENT_SHADER);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.contains("fragment "), "MSL should contain fragment stage qualifier:\n" + msl);
@@ -114,7 +115,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(SpirvTestShaders.LDS_MULTI_MIP_GLSL, Shaderc.shaderc_compute_shader);
         try {
             ShaderManager.remapSpirvForComputeSDLGPU(spirv);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.contains("kernel "), "Compute MSL should contain kernel stage qualifier:\n" + msl);
@@ -146,7 +147,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(glsl, Shaderc.shaderc_compute_shader);
         try {
             ShaderManager.remapSpirvForComputeSDLGPU(spirv);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.contains("kernel "), "compute MSL should contain kernel stage qualifier:\n" + msl);
@@ -167,7 +168,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(CHUNK_CULL_GLSL, Shaderc.shaderc_compute_shader);
         try {
             ShaderManager.remapSpirvForComputeSDLGPU(spirv);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.contains("kernel "), "Compute MSL should contain kernel stage qualifier:\n" + msl);
@@ -202,7 +203,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(VERTEX_WRITES_IMAGE_GLSL, Shaderc.shaderc_vertex_shader);
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_VERTEX_SHADER);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertFalse(msl.contains("vertex void main0"),
@@ -236,7 +237,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(pre.rewrittenSource(), Shaderc.shaderc_fragment_shader);
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_FRAGMENT_SHADER);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertTrue(msl.contains("angelica_renamed_try"), "the local must keep a renamed identifier:\n" + msl);
@@ -255,7 +256,7 @@ class MslCrossCompileTest {
             final ByteBuffer spirv = compile(SpirvTestShaders.VERTEX_GLSL, Shaderc.shaderc_vertex_shader);
             try {
                 ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_VERTEX_SHADER);
-                final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER);
+                final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER, false);
                 MemoryUtil.memFree(out.code());
             } finally {
                 MemoryUtil.memFree(spirv);
@@ -381,7 +382,7 @@ class MslCrossCompileTest {
         final ByteBuffer spirv = compile(glsl, Shaderc.shaderc_compute_shader);
         try {
             ShaderManager.remapSpirvForComputeSDLGPU(spirv);
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, false);
             try {
                 final String msl = decode(out.code());
                 assertUniqueBufferSlots(msl);
@@ -436,7 +437,7 @@ class MslCrossCompileTest {
             final ByteBuffer spirv = compile(glsl, Shaderc.shaderc_compute_shader);
             try {
                 ShaderManager.remapSpirvForComputeSDLGPU(spirv);
-                final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER);
+                final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, false);
                 try {
                     assertTrue(decode(out.code()).contains("kernel "), name + " should produce compute MSL");
                 } finally {
@@ -456,10 +457,10 @@ class MslCrossCompileTest {
         try {
             ShaderManager.remapSpirvForSDLGPU(spirv, GL20.GL_VERTEX_SHADER);
             final byte[] heap = ShaderCacheIO.toHeap(spirv);
-            MemoryUtil.memFree(MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER).code());
+            MemoryUtil.memFree(MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER, false).code());
             Reflect.<Map<?, ?>>get(cache(), "cache").clear();
             ShaderDiskCache.putBlob(cache().diskKey(heap, GL20.GL_VERTEX_SHADER), "poison_entry", "poison".getBytes(StandardCharsets.UTF_8));
-            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_VERTEX_SHADER, false);
             try {
                 assertEquals("poison_entry", out.entrypoint());
                 assertEquals("poison", new String(ShaderCacheIO.toHeap(out.code()), StandardCharsets.UTF_8));
@@ -471,6 +472,95 @@ class MslCrossCompileTest {
         }
     }
 
+    private static final String IMAGE_ATOMIC_GLSL = """
+        #version 460 core
+        layout(local_size_x = 64) in;
+        layout(binding = 0, r32i) uniform iimage2D counter;
+        void main() {
+            imageAtomicAdd(counter, ivec2(0, 0), 1);
+        }
+        """;
+
+    @Test
+    void imageAtomicsUseTheAuxBufferWithoutTextureAtomics() {
+        assertTrue(computeMsl(IMAGE_ATOMIC_GLSL, false).contains("counter_atomic"));
+    }
+
+    @Test
+    void imageAtomicsCompileToTextureAtomicsWhenAvailable() {
+        final String msl = computeMsl(IMAGE_ATOMIC_GLSL, true);
+        assertFalse(msl.contains("counter_atomic"), msl);
+        assertTrue(msl.contains("counter.atomic_fetch_add("), msl);
+    }
+
+    @Test
+    void textureAtomicsDoNotChangeShadersWithoutImageAtomics() {
+        final String base = computeMsl(CHUNK_CULL_GLSL, false);
+        Reflect.<Map<?, ?>>get(cache(), "cache").clear();
+        assertEquals(base, computeMsl(CHUNK_CULL_GLSL, true));
+    }
+
+    private static final String FRAGMENT_IMAGE_ATOMIC_GLSL = """
+        #version 460 core
+        layout(binding = 0, r32i) uniform iimage2D counter;
+        layout(location = 0) out vec4 fragColor;
+        void main() {
+            fragColor = vec4(imageAtomicAdd(counter, ivec2(0, 0), 1));
+        }
+        """;
+
+    @Test
+    void textureAtomicsLeaveRasterStagesAtMsl22() {
+        final ByteBuffer spirv = compile(FRAGMENT_IMAGE_ATOMIC_GLSL, Shaderc.shaderc_fragment_shader);
+        try {
+            assertTrue(CrossCompileUtil.usesImageAtomics(spirv));
+            final String base = fragmentMsl(spirv, false);
+            Reflect.<Map<?, ?>>get(cache(), "cache").clear();
+            final String atomics = fragmentMsl(spirv, true);
+            assertEquals(base, atomics);
+            assertFalse(atomics.contains("counter.atomic_fetch_add("), atomics);
+        } finally {
+            MemoryUtil.memFree(spirv);
+        }
+    }
+
+    private static String fragmentMsl(ByteBuffer spirv, boolean textureAtomics) {
+        final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL20.GL_FRAGMENT_SHADER, textureAtomics);
+        try {
+            return decode(out.code());
+        } finally {
+            MemoryUtil.memFree(out.code());
+        }
+    }
+
+    private static String computeMsl(String glsl, boolean textureAtomics) {
+        final ByteBuffer spirv = compile(glsl, Shaderc.shaderc_compute_shader);
+        try {
+            ShaderManager.remapSpirvForComputeSDLGPU(spirv);
+            final CrossCompileCache.Output out = MslCrossCompile.compile(spirv, GL43.GL_COMPUTE_SHADER, textureAtomics);
+            try {
+                return decode(out.code());
+            } finally {
+                MemoryUtil.memFree(out.code());
+            }
+        } finally {
+            MemoryUtil.memFree(spirv);
+        }
+    }
+
+    @Test
+    void onlyShadersWithImageAtomicsAreDetected() {
+        final ByteBuffer atomic = compile(IMAGE_ATOMIC_GLSL, Shaderc.shaderc_compute_shader);
+        final ByteBuffer plain = compile(CHUNK_CULL_GLSL, Shaderc.shaderc_compute_shader);
+        try {
+            assertTrue(CrossCompileUtil.usesImageAtomics(atomic));
+            assertFalse(CrossCompileUtil.usesImageAtomics(plain));
+        } finally {
+            MemoryUtil.memFree(atomic);
+            MemoryUtil.memFree(plain);
+        }
+    }
+
     private static String decode(ByteBuffer buf) {
         final ByteBuffer slice = buf.duplicate();
         int len = slice.remaining();

src/main/java/com/gtnewhorizons/angelica/debug/flyby/FlybyRunner.java

@@ -32,6 +32,7 @@ import net.minecraft.world.World;
 import net.minecraft.world.WorldServer;
 import net.minecraft.world.chunk.Chunk;
 import net.minecraft.world.chunk.IChunkProvider;
+import net.minecraftforge.common.DimensionManager;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.lwjgl.opengl.Display;
@@ -75,6 +76,9 @@ public final class FlybyRunner {
     private boolean jfr;
     private long recordingId;
     private final AtomicReference<FlybyRequest> pendingRequest = new AtomicReference<>();
+    private final AtomicReference<String> pendingDimensionPlayer = new AtomicReference<>();
+    private boolean dimensionRequested;
+    private volatile boolean dimensionAbandoned;
     private FlybyRequest activeRequest;
     private String[] sceneCommands = NO_COMMANDS;
     private volatile boolean sceneClearRequested;
@@ -197,6 +201,8 @@ public final class FlybyRunner {
         this.waitForTracy = false;
         this.waitForFocus = false;
         this.exitWhenDone = false;
+        this.dimensionRequested = false;
+        this.dimensionAbandoned = false;
         this.jfr = false;
         this.armed = false;
         this.fixedOrigin = null;
@@ -263,6 +269,7 @@ public final class FlybyRunner {
                     if (this.waitTicks++ % 100 == 0) LOGGER.info("Flyby waiting for window focus");
                     return;
                 }
+                if (this.awaitingDimension(mc, player)) return;
                 this.begin(mc, player);
             }
             case PREPARING -> {
@@ -331,6 +338,22 @@ public final class FlybyRunner {
         }
     }
 
+    private boolean awaitingDimension(Minecraft mc, EntityClientPlayerMP player) {
+        final Integer dimension = SystemProperties.FLYBY_DIMENSION;
+        if (!this.startedFromProperties || dimension == null || player.dimension == dimension || this.dimensionAbandoned) return false;
+        if (mc.getIntegratedServer() == null || !DimensionManager.isDimensionRegistered(dimension)) {
+            LOGGER.warn("Flyby: {}, staying in dimension {}", mc.getIntegratedServer() == null ? "not singleplayer" : "dimension " + dimension + " is not registered", player.dimension);
+            this.dimensionAbandoned = true;
+            return false;
+        }
+        if (!this.dimensionRequested) {
+            this.dimensionRequested = true;
+            LOGGER.info("Flyby: moving from dimension {} to {}", player.dimension, dimension);
+            this.pendingDimensionPlayer.set(player.getCommandSenderName());
+        }
+        return true;
+    }
+
     private void overridePauseOnLostFocus(Minecraft mc) {
         if (this.pauseOnLostFocusOverridden) return;
         this.pauseOnLostFocusSaved = mc.gameSettings.pauseOnLostFocus;
@@ -600,11 +623,23 @@ public final class FlybyRunner {
     public void onServerTick(TickEvent.ServerTickEvent event) {
         if (event.phase != TickEvent.Phase.END) return;
         final boolean clear = this.sceneClearRequested;
-        if (!clear && this.pendingRequest.get() == null) return;
+        if (!clear && this.pendingRequest.get() == null && this.pendingDimensionPlayer.get() == null) return;
 
         final MinecraftServer server = MinecraftServer.getServer();
         if (server == null) return;
 
+        final String traveller = this.pendingDimensionPlayer.getAndSet(null);
+        if (traveller != null) {
+            this.discardWorldChanges(server);
+            final EntityPlayerMP player = server.getConfigurationManager().func_152612_a(traveller);
+            if (player == null) {
+                LOGGER.warn("Flyby: no server player named '{}', staying in the current dimension", traveller);
+                this.dimensionAbandoned = true;
+            } else {
+                server.getConfigurationManager().transferPlayerToDimension(player, SystemProperties.FLYBY_DIMENSION);
+            }
+        }
+
         if (clear) {
             this.sceneClearRequested = false;
             this.clearScene(server);

src/main/java/net/coderbot/iris/pipeline/transform/RwImageStoreExtractor.java

@@ -1,5 +1,6 @@
 package net.coderbot.iris.pipeline.transform;
 
+import com.gtnewhorizons.angelica.glsm.backend.BackendManager;
 import com.gtnewhorizons.angelica.glsm.GlslTransformUtils;
 import net.coderbot.iris.gl.image.ImageInformation;
 import com.gtnewhorizons.angelica.glsm.texture.InternalTextureFormat;
@@ -23,6 +24,7 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.HashSet;
 import java.util.Set;
+import java.util.TreeSet;
 
 public final class RwImageStoreExtractor {
 
@@ -45,7 +47,7 @@ public final class RwImageStoreExtractor {
         }
     }
 
-    public record Result(String strippedSource, String computeSource, Set<String> writtenImages, RwExtractMode mode, boolean keepsImageAtomics) {}
+    public record Result(String strippedSource, String computeSource, Set<String> writtenImages, RwExtractMode mode, Set<String> atomicImages) {}
 
     public record ImageDecl(String name, String glslType, String layoutFormat) {}
 
@@ -69,8 +71,30 @@ public final class RwImageStoreExtractor {
         return activeCustomImages;
     }
 
+    private static volatile Set<String> atomicCustomImages = Set.of();
+
     public static void setActiveCustomImages(Map<String, ImageInformation> images) {
         activeCustomImages = (images == null) ? Map.of() : Map.copyOf(images);
+        final TreeSet<String> atomic = new TreeSet<>();
+        if (!activeCustomImages.isEmpty() && !BackendManager.RENDER_BACKEND.supportsComputeImageAtomics()) {
+            for (ImageInformation image : activeCustomImages.values()) {
+                if (image.isRelative()) continue;
+                if (BackendManager.RENDER_BACKEND.supportsComputeImageAtomics(image.width(), image.height(), image.depth())) atomic.add(image.name());
+            }
+        }
+        if (!atomic.equals(atomicCustomImages)) {
+            atomicCustomImages = atomic;
+            TransformPatcher.clearCache();
+            ShaderTransformer.clearCache();
+        }
+    }
+
+    public static boolean supportsAtomics(String image) {
+        return atomicCustomImages.contains(image) || BackendManager.RENDER_BACKEND.supportsComputeImageAtomics();
+    }
+
+    static String atomicCustomImagesKey() {
+        return String.join(",", atomicCustomImages);
     }
 
     private static final Set<String> RW_CALLS = Set.of(
@@ -119,7 +143,7 @@ public final class RwImageStoreExtractor {
         applyComputeMutations(tc, writtenImages, declared, mode == RwExtractMode.COMPOSITE_VSH ? quadInputs : chunkAttrsUsed, mode, nonWriteonlyImages);
         final String compute = buildComputeOutput(root, mode, writtenImages, nonWriteonlyImages, declared, attrs, quadInputs, definesFtransform, null);
 
-        return new Result(stripped, compute, writtenImages, mode, keepsImageAtomics(root, writtenImages));
+        return new Result(stripped, compute, writtenImages, mode, atomicImages(root, writtenImages));
     }
 
     public static Result tryExtractVertexReplay(String source) {
@@ -142,7 +166,7 @@ public final class RwImageStoreExtractor {
 
         final String stripped = buildRasterOutputViaSourceSlicing(source, decls, exprs, writtenImages, declared, RwExtractMode.VERTEX_REPLAY);
         final List<VertexInput> inputs = collectVertexInputs(decls);
-        if (inputs == null) return new Result(stripped, null, writtenImages, RwExtractMode.VERTEX_REPLAY, false);
+        if (inputs == null) return new Result(stripped, null, writtenImages, RwExtractMode.VERTEX_REPLAY, Set.of());
 
         final List<String> readImages = readOnlyImages(exprs, declared, writtenImages);
         final Set<String> inputNames = new LinkedHashSet<>();
@@ -153,7 +177,7 @@ public final class RwImageStoreExtractor {
         final ReplayLayout layout = new ReplayLayout(inputs, readImages, replayHeader(inputs, writtenImages, readImages, collectSamplerNames(decls)));
         final String compute = buildComputeOutput(root, RwExtractMode.VERTEX_REPLAY, writtenImages, nonWriteonlyImages, declared, List.of(), Set.of(), false, layout);
 
-        return new Result(stripped, compute, writtenImages, RwExtractMode.VERTEX_REPLAY, keepsImageAtomics(root, writtenImages));
+        return new Result(stripped, compute, writtenImages, RwExtractMode.VERTEX_REPLAY, atomicImages(root, writtenImages));
     }
 
     private record ReplayLayout(List<VertexInput> inputs, List<String> readImages, String header) {}
@@ -498,12 +522,15 @@ public final class RwImageStoreExtractor {
         });
     }
 
-    private static boolean keepsImageAtomics(GLSLParser.Translation_unitContext computeRoot, Set<String> writtenImages) {
+    private static Set<String> atomicImages(GLSLParser.Translation_unitContext computeRoot, Set<String> writtenImages) {
+        final Set<String> out = new LinkedHashSet<>();
         for (var expr : GlslAstHelpers.collectAll(computeRoot, GLSLParser.Postfix_expressionContext.class)) {
             final String fname = GlslAstHelpers.extractCallName(expr);
-            if (fname != null && ATOMIC_CALLS.contains(fname) && writtenImages.contains(GlslAstHelpers.firstArgIdentifier(expr))) return true;
+            if (fname == null || !ATOMIC_CALLS.contains(fname)) continue;
+            final String image = GlslAstHelpers.firstArgIdentifier(expr);
+            if (writtenImages.contains(image)) out.add(image);
         }
-        return false;
+        return out;
     }
 
     private static boolean isCallTopLevelStatement(GLSLParser.Postfix_expressionContext expr) {

src/main/java/net/coderbot/iris/pipeline/transform/ShaderTransformer.java

@@ -35,6 +35,7 @@ import java.util.Objects;
 import java.util.Set;
 import java.util.function.BooleanSupplier;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -237,7 +238,7 @@ public class ShaderTransformer {
         final ShaderDiskCache.Key k = ShaderDiskCache.key("iris-transform").str(patch.name());
         parameters.appendDiskKey(k);
         k.i(RenderSystem.getMaxGlslVersion()).b(RenderSystem.supportsSSBO()).b(RenderSystem.supportsImageLoadStore())
-            .b(RenderSystem.isGLES()).b(BackendManager.RENDER_BACKEND.isSDLGPU()).b(BackendManager.RENDER_BACKEND.supportsComputeImageAtomics())
+            .b(RenderSystem.isGLES()).b(BackendManager.RENDER_BACKEND.isSDLGPU()).b(BackendManager.RENDER_BACKEND.supportsComputeImageAtomics()).str(RwImageStoreExtractor.atomicCustomImagesKey())
             .i(maxSupportedHoistVersion);
         if (enabledRequirements == null) {
             k.i(-1);
@@ -560,16 +561,16 @@ public class ShaderTransformer {
                 artifactsOut.put(shaderType, new StageArtifact(treeHolder[0], formattedShader.length() - printedBody.length()));
             }
         }
-        if (BackendManager.RENDER_BACKEND.isSDLGPU()) extractRwImageStores(result, patchType, artifactsOut, BackendManager.RENDER_BACKEND.supportsComputeImageAtomics());
+        if (BackendManager.RENDER_BACKEND.isSDLGPU()) extractRwImageStores(result, patchType, artifactsOut, RwImageStoreExtractor::supportsAtomics);
         watch.stop();
         Iris.logger.info("[Load #{}] Transformed shader for {} in {}", Iris.getShaderPackLoadId(), patchType.name(), watch);
         return result;
     }
 
-    static void extractRwImageStores(EnumMap<PatchShaderType, String> result, Patch patchType, EnumMap<PatchShaderType, StageArtifact> artifactsOut, boolean computeImageAtomics) {
+    static void extractRwImageStores(EnumMap<PatchShaderType, String> result, Patch patchType, EnumMap<PatchShaderType, StageArtifact> artifactsOut, Predicate<String> imageAtomics) {
         if (result.containsKey(PatchShaderType.COMPUTE)) return;
         if (patchType == Patch.ATTRIBUTES) {
-            extractVertexReplay(result, artifactsOut, computeImageAtomics);
+            extractVertexReplay(result, artifactsOut, imageAtomics);
             return;
         }
 
@@ -589,7 +590,7 @@ public class ShaderTransformer {
         if (artifactsOut != null) artifactsOut.remove(stage);
         if (patchType != Patch.COMPOSITE && patchType != Patch.CELERITAS_TERRAIN) {
             Iris.logger.info("[RwImageStoreExtractor] Dropped image writes from {} (written={})", patchType.name(), extracted.writtenImages());
-        } else if (extracted.keepsImageAtomics() && !computeImageAtomics) {
+        } else if (needsUnsupportedAtomics(extracted, imageAtomics)) {
             Iris.logger.info("[RwImageStoreExtractor] Dropped image writes from {}, they need image atomics this backend cannot run (written={})", patchType.name(), extracted.writtenImages());
         } else {
             result.put(PatchShaderType.COMPUTE, extracted.computeSource());
@@ -597,7 +598,11 @@ public class ShaderTransformer {
         }
     }
 
-    private static void extractVertexReplay(EnumMap<PatchShaderType, String> result, EnumMap<PatchShaderType, StageArtifact> artifactsOut, boolean computeImageAtomics) {
+    private static boolean needsUnsupportedAtomics(RwImageStoreExtractor.Result extracted, Predicate<String> imageAtomics) {
+        return !extracted.atomicImages().stream().allMatch(imageAtomics);
+    }
+
+    private static void extractVertexReplay(EnumMap<PatchShaderType, String> result, EnumMap<PatchShaderType, StageArtifact> artifactsOut, Predicate<String> imageAtomics) {
         final String fsh = result.get(PatchShaderType.FRAGMENT);
         final RwImageStoreExtractor.Result fshResult = (fsh != null) ? RwImageStoreExtractor.tryExtract(fsh, PatchShaderType.FRAGMENT, Patch.ATTRIBUTES.name()) : null;
         if (fshResult != null) {
@@ -615,7 +620,7 @@ public class ShaderTransformer {
             Iris.logger.info("[RwImageStoreExtractor] Dropped vertex image writes from ATTRIBUTES, an input cannot be replayed (written={})", vshResult.writtenImages());
             return;
         }
-        if (vshResult.keepsImageAtomics() && !computeImageAtomics) {
+        if (needsUnsupportedAtomics(vshResult, imageAtomics)) {
             Iris.logger.info("[RwImageStoreExtractor] Dropped vertex image writes from ATTRIBUTES, they need image atomics this backend cannot run (written={})", vshResult.writtenImages());
             return;
         }

src/test/java/net/coderbot/iris/pipeline/transform/CompositeVertexWriteExtractionTest.java

@@ -60,7 +60,7 @@ class CompositeVertexWriteExtractionTest {
         final EnumMap<PatchShaderType, String> result = new EnumMap<>(transformed);
         result.remove(PatchShaderType.COMPUTE);
 
-        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, null, true);
+        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, null, image -> true);
 
         final String compute = result.get(PatchShaderType.COMPUTE);
         assertEquals(RwImageStoreExtractor.RwExtractMode.COMPOSITE_VSH, RwImageStoreExtractor.parseSentinel(compute));
@@ -77,7 +77,7 @@ class CompositeVertexWriteExtractionTest {
         result.remove(PatchShaderType.COMPUTE);
 
         // The CompSwap's result is used, so it stays a real atomic, which native Metal cannot bind
-        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, null, false);
+        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, null, image -> false);
 
         assertFalse(result.containsKey(PatchShaderType.COMPUTE));
         assertFalse(result.get(PatchShaderType.VERTEX).contains("imageAtomic"), result.get(PatchShaderType.VERTEX));

src/test/java/net/coderbot/iris/pipeline/transform/CustomImageAtomicsTest.java

@@ -0,0 +1,86 @@
+package net.coderbot.iris.pipeline.transform;
+
+import com.gtnewhorizons.angelica.glsm.backend.BackendManager;
+import com.gtnewhorizons.angelica.glsm.backend.RenderBackend;
+import com.gtnewhorizons.angelica.glsm.testutil.Reflect;
+import com.gtnewhorizons.angelica.glsm.texture.InternalTextureFormat;
+import com.gtnewhorizons.angelica.glsm.texture.PixelFormat;
+import com.gtnewhorizons.angelica.glsm.texture.PixelType;
+import com.gtnewhorizons.angelica.glsm.texture.TextureType;
+import net.coderbot.iris.gl.image.ImageInformation;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.AdditionalAnswers.delegatesTo;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.doReturn;
+import static org.mockito.Mockito.mock;
+
+class CustomImageAtomicsTest {
+
+    private RenderBackend original;
+    private RenderBackend backend;
+
+    @BeforeEach
+    void installBackend() {
+        original = BackendManager.RENDER_BACKEND;
+        backend = mock(RenderBackend.class, delegatesTo(original));
+        doAnswer(call -> call.<Integer>getArgument(0) == 40 && call.<Integer>getArgument(1) == 10 && call.<Integer>getArgument(2) == 1).when(backend).supportsComputeImageAtomics(anyInt(), anyInt(), anyInt());
+        Reflect.setStaticFinal(BackendManager.class, "RENDER_BACKEND", RenderBackend.class, backend);
+    }
+
+    @AfterEach
+    void restoreBackend() {
+        Reflect.setStaticFinal(BackendManager.class, "RENDER_BACKEND", RenderBackend.class, original);
+        RwImageStoreExtractor.setActiveCustomImages(null);
+    }
+
+    private static ImageInformation image(String name, int width, int height, boolean relative) {
+        return new ImageInformation(name, name + "_sampler", TextureType.TEXTURE_2D, PixelFormat.RED_INTEGER, InternalTextureFormat.R32I, PixelType.INT, width, height, 1, false, relative, relative ? 1.0f : 0.0f, relative ? 1.0f : 0.0f);
+    }
+
+    private static Map<String, ImageInformation> images() {
+        return Map.of("small_img", image("small_img", 40, 10, false), "large_img", image("large_img", 64, 64, false), "screen_img", image("screen_img", 40, 10, true));
+    }
+
+    @Test
+    void onlyImagesTheBackendAcceptsBySizeKeepAtomics() {
+        doReturn(false).when(backend).supportsComputeImageAtomics();
+        RwImageStoreExtractor.setActiveCustomImages(images());
+
+        assertTrue(RwImageStoreExtractor.supportsAtomics("small_img"));
+        assertFalse(RwImageStoreExtractor.supportsAtomics("large_img"));
+        assertFalse(RwImageStoreExtractor.supportsAtomics("screen_img"));
+        assertFalse(RwImageStoreExtractor.supportsAtomics("colortex0"));
+        assertEquals("small_img", RwImageStoreExtractor.atomicCustomImagesKey());
+    }
+
+    @Test
+    void backendWithImageAtomicsEverywhereNeedsNoPerImageList() {
+        doReturn(true).when(backend).supportsComputeImageAtomics();
+        RwImageStoreExtractor.setActiveCustomImages(images());
+
+        assertTrue(RwImageStoreExtractor.supportsAtomics("small_img"));
+        assertTrue(RwImageStoreExtractor.supportsAtomics("large_img"));
+        assertTrue(RwImageStoreExtractor.supportsAtomics("screen_img"));
+        assertTrue(RwImageStoreExtractor.supportsAtomics("colortex0"));
+        assertEquals("", RwImageStoreExtractor.atomicCustomImagesKey());
+    }
+
+    @Test
+    void clearingTheImagesClearsTheList() {
+        doReturn(false).when(backend).supportsComputeImageAtomics();
+        RwImageStoreExtractor.setActiveCustomImages(images());
+        RwImageStoreExtractor.setActiveCustomImages(null);
+
+        assertFalse(RwImageStoreExtractor.supportsAtomics("small_img"));
+        assertEquals("", RwImageStoreExtractor.atomicCustomImagesKey());
+    }
+}

src/test/java/net/coderbot/iris/pipeline/transform/ShaderTransformerTest.java

@@ -246,7 +246,7 @@ class ShaderTransformerTest {
 
         final EnumMap<PatchShaderType, String> result = new EnumMap<>(out);
         final String fragmentBefore = result.get(PatchShaderType.FRAGMENT);
-        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, artifacts, true);
+        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, artifacts, image -> true);
 
         assertNotNull(result.get(PatchShaderType.COMPUTE), "fixture must actually exercise the extraction");
         assertFalse(fragmentBefore.equals(result.get(PatchShaderType.FRAGMENT)), "fixture must actually rewrite the fragment stage");
@@ -263,7 +263,7 @@ class ShaderTransformerTest {
         assertNotNull(out);
 
         final EnumMap<PatchShaderType, String> result = new EnumMap<>(out);
-        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, artifacts, true);
+        ShaderTransformer.extractRwImageStores(result, Patch.COMPOSITE, artifacts, image -> true);
 
         assertFalse(artifacts.isEmpty(), "fixture must leave at least one artifact to check");
         for (Map.Entry<PatchShaderType, ShaderTransformer.StageArtifact> e : artifacts.entrySet()) {

src/test/java/net/coderbot/iris/pipeline/transform/VertexReplayExtractionTest.java

@@ -237,7 +237,7 @@ class VertexReplayExtractionTest {
     void atomicReplayIsDroppedWithoutBackendSupport() {
         final EnumMap<PatchShaderType, String> atomic = new EnumMap<>(PatchShaderType.class);
         atomic.put(PatchShaderType.VERTEX, TransformPatcher.patchAttributes(CRYSTAL_VERTEX, null, FRAGMENT, TEX_LM).get(PatchShaderType.VERTEX));
-        ShaderTransformer.extractRwImageStores(atomic, Patch.ATTRIBUTES, null, false);
+        ShaderTransformer.extractRwImageStores(atomic, Patch.ATTRIBUTES, null, image -> false);
         assertFalse(atomic.containsKey(PatchShaderType.COMPUTE));
         assertFalse(atomic.get(PatchShaderType.VERTEX).contains("imageAtomic"), atomic.get(PatchShaderType.VERTEX));
 
@@ -251,10 +251,25 @@ class VertexReplayExtractionTest {
                 gl_Position = iris_Vertex;
             }
             """);
-        ShaderTransformer.extractRwImageStores(plain, Patch.ATTRIBUTES, null, false);
+        ShaderTransformer.extractRwImageStores(plain, Patch.ATTRIBUTES, null, image -> false);
         assertTrue(plain.containsKey(PatchShaderType.COMPUTE), "plain stores need no atomics, so the replay stays");
     }
 
+    @Test
+    void atomicReplayFollowsTheImageItWrites() {
+        final String vertex = TransformPatcher.patchAttributes(CRYSTAL_VERTEX, null, FRAGMENT, TEX_LM).get(PatchShaderType.VERTEX);
+
+        final EnumMap<PatchShaderType, String> supported = new EnumMap<>(PatchShaderType.class);
+        supported.put(PatchShaderType.VERTEX, vertex);
+        ShaderTransformer.extractRwImageStores(supported, Patch.ATTRIBUTES, null, "endcrystal_img"::equals);
+        assertTrue(supported.containsKey(PatchShaderType.COMPUTE));
+
+        final EnumMap<PatchShaderType, String> unsupported = new EnumMap<>(PatchShaderType.class);
+        unsupported.put(PatchShaderType.VERTEX, vertex);
+        ShaderTransformer.extractRwImageStores(unsupported, Patch.ATTRIBUTES, null, "voxel_img"::equals);
+        assertFalse(unsupported.containsKey(PatchShaderType.COMPUTE));
+    }
+
     @Test
     void matrixInputIsNotReplayed() {
         final RwImageStoreExtractor.Result result = RwImageStoreExtractor.tryExtractVertexReplay("""
@@ -293,7 +308,7 @@ class VertexReplayExtractionTest {
             }
             """);
 
-        ShaderTransformer.extractRwImageStores(result, Patch.ATTRIBUTES, null, true);
+        ShaderTransformer.extractRwImageStores(result, Patch.ATTRIBUTES, null, image -> true);
 
         assertNotNull(result.get(PatchShaderType.COMPUTE));
         assertFalse(result.get(PatchShaderType.VERTEX).contains("imageStore"));