Enable SPV_KHR_float_controls2 when the device accepts it

lib/cl/device.jl

@@ -112,6 +112,17 @@ end
         return String[string(s) for s in split(bs)]
     end
 
+    # cl_khr_spirv_queries: the SPIR-V extensions the device accepts. Empty if the device
+    # doesn't support the query.
+    if s == :spirv_extensions
+        "cl_khr_spirv_queries" in d.extensions || return String[]
+        size = Ref{Csize_t}()
+        clGetDeviceInfo(d, CL_DEVICE_SPIRV_EXTENSIONS_KHR, 0, C_NULL, size)
+        result = Vector{Cstring}(undef, size[] ÷ sizeof(Cstring))
+        clGetDeviceInfo(d, CL_DEVICE_SPIRV_EXTENSIONS_KHR, size[], result, C_NULL)
+        return String[unsafe_string(ext) for ext in result]
+    end
+
     if s == :platform
         result = Ref{cl_platform_id}()
         clGetDeviceInfo(d, CL_DEVICE_PLATFORM, sizeof(cl_platform_id), result, C_NULL)
@@ -187,9 +198,9 @@ function Base.propertynames(::Device, private::Bool=false)
         :max_work_group_size, :max_parameter_size, :profiling_timer_resolution,
         # boolean properties
         :has_image_support, :has_local_mem, :host_unified_memory, :available,
-        :compiler_available, :non_uniform_work_group_support
+        :compiler_available, :non_uniform_work_group_support,
         # other
-        :extensions, :platform, :device_type, :max_work_item_size, :sub_group_sizes,
+        :extensions, :spirv_extensions, :platform, :device_type, :max_work_item_size, :sub_group_sizes,
         :max_image2d_shape, :max_image3d_shape,
     )
     return private ? (fieldnames(Device)..., props...) : props

src/compiler/compilation.jl

@@ -184,6 +184,12 @@ function spirv_driver(dev::cl.Device)
     end
 end
 
+# SPIR-V extensions that let the back-ends express more of the IR, enabled when the device
+# accepts them:
+# - SPV_KHR_float_controls2 carries per-instruction fast-math flags, such as `contract` (e.g.
+#   from `muladd`), which base SPIR-V can't express, so drivers never fuse those into FMAs.
+const OPTIONAL_SPIRV_EXTENSIONS = ["SPV_KHR_float_controls2"]
+
 @noinline function _compiler_config(dev, backend; kernel=true, name=nothing, always_inline=false,
                                      debug_level=Base.JLOptions().debug_level,
                                      sub_group_size::Union{Nothing,Int}=_sub_group_size(dev),
@@ -210,8 +216,13 @@ end
         error("Device does not support cl_intel_required_subgroup_size")
     end
 
+    if extensions === nothing
+        # spirv2clc doesn't implement them, so only for programs built from SPIR-V
+        extensions = source ? String[] : filter(in(dev.spirv_extensions), OPTIONAL_SPIRV_EXTENSIONS)
+    end
+
     # GPUCompiler adds the extensions that the atomics need
-    spirv_ext = extensions === nothing ? "" : join(("+$ext" for ext in extensions), ",")
+    spirv_ext = join(("+$ext" for ext in extensions), ",")
 
     # create GPUCompiler objects
     target = SPIRVCompilerTarget(; version=SPIRV_VERSION, supports_fp16, supports_fp64,

test/device.jl

@@ -31,6 +31,7 @@ end
             :driver_version,
             :version,
             :extensions,
+            :spirv_extensions,
             :platform,
             :name,
             :device_type,
@@ -78,6 +79,8 @@ end
             if length(v) > 0
                 @test isa(v, Array{String, 1})
             end
+        elseif k == :spirv_extensions
+            @test isa(v, Vector{String})
         elseif k == :platform
             @test v == cl.platform()
         elseif k == :max_work_item_sizes

test/execution.jl

@@ -456,6 +456,27 @@ end
     end
 end
 
+@testset "fast-math contraction" begin
+    # `muladd` is `fmul contract` + `fadd contract`, which only SPV_KHR_float_controls2 can
+    # carry into SPIR-V (JuliaGPU/KernelAbstractions.jl#865)
+    function contract_kernel(c, a, b)
+        i = get_global_id()
+        @inbounds c[i] = muladd(a[i], b[i], c[i])
+        return
+    end
+    T = CLDeviceArray{Float32, 1, AS.CrossWorkgroup}
+    # (enabled when the device accepts it, except on the OpenCL C source backend)
+    enabled = occursin("SPV_KHR_float_controls2",
+                       OpenCL.compiler_config(cl.device()).target.extensions)
+    spirv = sprint(io -> OpenCL.code_native(io, contract_kernel, Tuple{T, T, T}; kernel = true))
+    @test occursin("FPFastMathMode AllowContract", spirv) == enabled
+
+    # explicitly passing the extensions disables the default ones
+    spirv = sprint(io -> OpenCL.code_native(io, contract_kernel, Tuple{T, T, T};
+                                            kernel = true, extensions = String[]))
+    @test !occursin("AllowContract", spirv)
+end
+
 @testset "has_feature folding" begin
     # has_feature must fold at compile time: the feature-bitset global gets baked in and DCE'd,
     # leaving only the branch matching the device.