From 4cc0e6e3b29cee062b3d6d4bed5e7b04f8f27295 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Mon, 27 Jul 2026 16:15:12 +0200
Subject: [PATCH] [pocl] Enable SPV_EXT_shader_atomic_float_add where supported

The SPIR-V backend refuses to translate a module containing an atomic fadd
unless the extension guarding it has been listed:

    LLVM ERROR: The atomic float instruction requires the following SPIR-V
    extension: SPV_EXT_shader_atomic_float_add

Enzyme's reverse mode runs into this because it accumulates gradients with
atomic fadd, so reverse-mode AD over a POCL kernel fails to compile outright.

Derive the extension list from the device the way `supports_fp16`/`supports_fp64`
already are, keyed off the corresponding OpenCL extension (`cl_ext_float_atomics`,
which PoCL advertises). Listing an extension only permits it -- nothing is
emitted unless a module needs those instructions -- so kernels that do not use
float atomics are unaffected. An explicit `extensions=` keyword still wins, and
`default_spirv_extensions` gives later extensions an obvious home.

Verified end-to-end: with this change, reverse-mode Enzyme over a POCL kernel
compiles, runs and produces correct gradients.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 src/pocl/compiler/compilation.jl | 35 ++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl
index 1fbc06dfe..bad14a710 100644
--- a/src/pocl/compiler/compilation.jl
+++ b/src/pocl/compiler/compilation.jl
@@ -142,6 +142,29 @@ end
 
 ## compiler implementation (configure, compile, and link)
 
+"""
+    default_spirv_extensions(dev)
+
+SPIR-V extensions to permit for `dev`, as the `+`-prefixed, comma-separated string
+`SPIRVCompilerTarget` passes on to the backend via `-spirv-ext`.
+
+Listing an extension only *permits* it: nothing is emitted unless a module actually
+needs the instructions it guards, so this costs nothing for kernels that don't.
+"""
+function default_spirv_extensions(dev)
+    exts = String[]
+
+    # Atomic float add. Without this the backend refuses to translate the module at all:
+    #   LLVM ERROR: The atomic float instruction requires the following SPIR-V
+    #   extension: SPV_EXT_shader_atomic_float_add
+    # Enzyme's reverse mode hits this because it accumulates gradients with atomic fadd.
+    if "cl_ext_float_atomics" in dev.extensions
+        push!(exts, "+SPV_EXT_shader_atomic_float_add")
+    end
+
+    return join(exts, ",")
+end
+
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
 const _toolchain = Ref{Any}()
 const _compiler_configs = Dict{UInt, OpenCLCompilerConfig}()
@@ -154,7 +177,11 @@ function compiler_config(dev::cl.Device; kwargs...)
     end
     return config
 end
-@noinline function _compiler_config(dev; kernel = true, name = nothing, always_inline = false, sub_group_size::Union{Nothing, Int} = 32, kwargs...)
+@noinline function _compiler_config(
+        dev; kernel = true, name = nothing, always_inline = false,
+        sub_group_size::Union{Nothing, Int} = 32,
+        extensions::Union{Nothing, String} = nothing, kwargs...
+    )
     supports_fp16 = "cl_khr_fp16" in dev.extensions
     supports_fp64 = "cl_khr_fp64" in dev.extensions
 
@@ -162,8 +189,12 @@ end
         error("$sub_group_size is not a valid sub-group size for this device.")
     end
 
+    if extensions === nothing
+        extensions = default_spirv_extensions(dev)
+    end
+
     # create GPUCompiler objects
-    target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, validate = true, kwargs...)
+    target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, extensions, validate = true, kwargs...)
     params = OpenCLCompilerParams(; sub_group_size)
     return CompilerConfig(target, params; kernel, name, always_inline)
 end