diff --git a/src/spirv.jl b/src/spirv.jl index 4c4e1dda..88a12ee7 100644 --- a/src/spirv.jl +++ b/src/spirv.jl @@ -295,6 +295,63 @@ function translate(input::Vector{UInt8}, options::Ref{LLVMSPIRVTranslateOptions} return take_buffer(backend, buffer[]) end +# Kernels may only contract floating-point operations marked `contract`, so the LLVM back-end +# gives them the `ContractionOff` execution mode. With SPV_KHR_float_controls2, which +# deprecates that, it instead gives them `FPFastMathDefault` (with no flags) for every +# floating-point type in the module, but only declares the `FloatControls2` capability and +# the extension that mode needs when it decorates an instruction with fast-math flags. A +# kernel with floating-point types but without `contract` operations is then invalid. Declare +# them in the generated binary. (Declaring the default in `spirv.ExecutionMode` metadata +# instead makes the back-end account for it, but crashes it on some `OpExtInst`s.) +function add_float_controls2_declarations(spirv::Vector{UInt8}) + OpExtension, OpCapability, OpExecutionModeId = 10, 17, 331 + ExecutionModeFPFastMathDefault, CapabilityFloatControls2 = 6028, 6029 + extension = "SPV_KHR_float_controls2" + + words = reinterpret(UInt32, spirv) + length(words) >= 5 && words[1] == 0x07230203 || return spirv + # find the end of the capabilities and of the extensions, which come first, and whether + # they need to be declared + capabilities_end = extensions_end = 6 + has_capability = has_extension = uses_mode = false + i = 6 + while i <= length(words) + count, opcode = Int(words[i] >> 16), Int(words[i] & 0xffff) + count == 0 && return spirv + if opcode == OpCapability + has_capability |= words[i+1] == CapabilityFloatControls2 + capabilities_end = extensions_end = i + count + elseif opcode == OpExtension + has_extension |= spirv_string(words, i+1, i+count-1) == extension + extensions_end = i + count + elseif opcode == OpExecutionModeId + uses_mode |= words[i+2] == ExecutionModeFPFastMathDefault + end + i += count + end + uses_mode && !(has_capability && has_extension) || return spirv + + # the extension's name, NUL-terminated and padded to whole words + name = Vector{UInt8}(extension) + append!(name, zeros(UInt8, 4 - length(name) % 4)) + declarations = UInt32[] + has_capability || append!(declarations, (UInt32(2) << 16 | OpCapability, + CapabilityFloatControls2)) + capability_words = length(declarations) + has_extension || append!(declarations, UInt32(1 + length(name) รท 4) << 16 | OpExtension, + reinterpret(UInt32, name)) + patched = vcat(words[1:capabilities_end-1], declarations[1:capability_words], + words[capabilities_end:extensions_end-1], declarations[capability_words+1:end], + words[extensions_end:end]) + return Vector{UInt8}(reinterpret(UInt8, patched)) +end + +# the NUL-terminated string literal in `words[first:last]` +function spirv_string(words, first, last) + bytes = reinterpret(UInt8, words[first:last]) + return String(bytes[1:something(findfirst(iszero, bytes), length(bytes) + 1) - 1]) +end + @unlocked function mcgen(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module, format=LLVM.CodeGenFileType.Assembly) target = job.config.target @@ -322,9 +379,10 @@ end options = Ref(SPIRVCompileOptions(Int === Int64, version.major, version.minor, Base.unsafe_convert(Cstring, extensions), #=opt_level=# 2)) - external_compile(backend, input, options, - "Failed to compile to SPIR-V with the SPIR-V back-end") + obj = external_compile(backend, input, options, + "Failed to compile to SPIR-V with the SPIR-V back-end") end + add_float_controls2_declarations(obj) elseif target.backend === :khronos # translate in-process through libllvm_spirv. like `llvm-spirv`, the translator # takes the triple from the module and rejects unknown extensions. diff --git a/test/spirv.jl b/test/spirv.jl index d7834519..d5b95050 100644 --- a/test/spirv.jl +++ b/test/spirv.jl @@ -306,6 +306,67 @@ end end end +@testset "float controls" begin + # with SPV_KHR_float_controls2, `contract` operations are decorated `AllowContract`. the + # LLVM back-end also gives every kernel floating-point defaults, which need the extension, + # so this checks that kernels without `contract` operations still validate. + mod = @eval module $(gensym()) + function contract(p::Ptr{Float32}, x::Float32, y::Float32) + unsafe_store!(p, muladd(x, y, unsafe_load(p))) + return + end + function strict(p::Ptr{Float64}, x::Float64) + unsafe_store!(p, 2 * x) + return + end + function integer(p::Ptr{Int32}, x::Int32) + unsafe_store!(p, 2 * x) + return + end + # (a floating-point type only in the signature) + function signature(p::Ptr{Int32}, x::Float32) + unsafe_store!(p, Int32(1)) + return + end + function vector(p::Ptr{Int32}, x::NTuple{4, VecElement{Float32}}) + unsafe_store!(p, Int32(1)) + return + end + end + extensions = "+SPV_KHR_float_controls2" + + @test @filecheck begin + @check "OpExtension \"SPV_KHR_float_controls2\"" + @check "FPFastMathMode AllowContract" + SPIRV.code_native(mod.contract, Tuple{Ptr{Float32}, Float32, Float32}; backend, + kernel=true, extensions) + end + @test @filecheck begin + @check "OpEntryPoint Kernel" + @check_not "AllowContract" + SPIRV.code_native(mod.strict, Tuple{Ptr{Float64}, Float64}; backend, kernel=true, + extensions) + end + @test @filecheck begin + @check "OpEntryPoint Kernel" + @check_not "AllowContract" + SPIRV.code_native(mod.integer, Tuple{Ptr{Int32}, Int32}; backend, kernel=true, + extensions) + end + @test @filecheck begin + @check "OpEntryPoint Kernel" + @check_not "AllowContract" + SPIRV.code_native(mod.signature, Tuple{Ptr{Int32}, Float32}; backend, kernel=true, + extensions) + end + @test @filecheck begin + @check "OpEntryPoint Kernel" + @check_not "AllowContract" + SPIRV.code_native(mod.vector, Tuple{Ptr{Int32}, NTuple{4, VecElement{Float32}}}; + backend, kernel=true, extensions) + end +end + @testset "nested insertvalue" begin # Intel's graphics compiler drops fields when legalizing aggregates built by nested # `insertvalue`s, so these are flattened for that driver (JuliaGPU/OpenCL.jl#502)