diff --git a/src/array.jl b/src/array.jl index eba51c84..f2a8b3b5 100644 --- a/src/array.jl +++ b/src/array.jl @@ -46,22 +46,39 @@ function _device_supports_bfloat16(dev=device()) return false end +# `check_eltype` runs for every array, and querying the device allocates, so cache what it +# needs. Device handles are process-local, so `__init__` empties this. +const EltypeSupport = @NamedTuple{float16::Bool, float64::Bool, bfloat16::Bool} +const eltype_support_cache = Dict{ZeDevice, EltypeSupport}() +const eltype_support_lock = ReentrantLock() + +function eltype_support(dev::ZeDevice=device()) + Base.@lock eltype_support_lock begin + get!(eltype_support_cache, dev) do + flags = oneL0.module_properties(dev).flags + (; float16 = flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 != 0, + float64 = flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 != 0, + bfloat16 = _device_supports_bfloat16(dev)) + end + end +end + function check_eltype(T) Base.allocatedinline(T) || error("oneArray only supports element types that are stored inline") Base.isbitsunion(T) && error("oneArray does not yet support isbits-union arrays") - if oneL0.module_properties(device()).fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 != - oneL0.ZE_DEVICE_MODULE_FLAG_FP16 + support = eltype_support() + if !support.float16 contains_eltype(T, Float16) && error("Float16 is not supported on this device") end - if oneL0.module_properties(device()).fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 != - oneL0.ZE_DEVICE_MODULE_FLAG_FP64 + if !support.float64 contains_eltype(T, Float64) && error("Float64 is not supported on this device") end - return @static if isdefined(Core, :BFloat16) - if !_device_supports_bfloat16() - contains_eltype(T, Core.BFloat16) && error("BFloat16 is not supported on this device") - end + @static if isdefined(Core, :BFloat16) + if !support.bfloat16 + contains_eltype(T, Core.BFloat16) && error("BFloat16 is not supported on this device") end + end + return end """ diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 7d2a9524..034ebc48 100644 --- a/src/compiler/compilation.jl +++ b/src/compiler/compilation.jl @@ -200,8 +200,9 @@ end @noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, sub_group_size=nothing, kwargs...) properties = oneL0.module_properties(dev) - supports_fp16 = properties.fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16 - supports_fp64 = properties.fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64 + support = eltype_support(dev) + supports_fp16 = support.float16 + supports_fp64 = support.float64 if sub_group_size !== nothing sizes = oneL0.compute_properties(dev).subGroupSizes diff --git a/src/oneAPI.jl b/src/oneAPI.jl index da72d9f9..fb8ddbfd 100644 --- a/src/oneAPI.jl +++ b/src/oneAPI.jl @@ -136,8 +136,9 @@ function __init__() end if oneL0.functional[] - # slabs are process-local + # slabs and device handles are process-local empty!(pools) + empty!(eltype_support_cache) POOL[] = oneL0.parse_env_bool("ONEAPI_MEMORY_POOL", true) BATCHING[] = oneL0.parse_env_bool("ONEAPI_COMMAND_BATCHING", true) diff --git a/test/array.jl b/test/array.jl index 22bb29db..33d2686f 100644 --- a/test/array.jl +++ b/test/array.jl @@ -18,6 +18,20 @@ import Adapt @test collect(oneAPI.fill(1, 2, 2)) == ones(Int, 2, 2) end +@testset "element type checks" begin + if !float64_supported + @test_throws "Float64 is not supported" oneArray{Float64}(undef, 1) + @test_throws "Float64 is not supported" oneArray{ComplexF64}(undef, 1) + end + if !float16_supported + @test_throws "Float16 is not supported" oneArray{Float16}(undef, 1) + end + + # every array construction performs this check, so it shouldn't query the device + oneAPI.check_eltype(Float32) + @test @allocated(oneAPI.check_eltype(Float32)) == 0 +end + @testset "adapt" begin A = rand(Float32, 3, 3) dA = oneArray(A) diff --git a/test/runtests.jl b/test/runtests.jl index ea4dce3d..eed10f89 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -85,11 +85,11 @@ init_worker_code = quote Float16, Float32, ComplexF32] - const float16_supported = oneL0.module_properties(device()).fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16 + const float16_supported = oneL0.module_properties(device()).flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 != 0 if float16_supported append!(eltypes, [#=Float16,=# ComplexF16]) end - const float64_supported = oneL0.module_properties(device()).fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64 + const float64_supported = oneL0.module_properties(device()).flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 != 0 if float64_supported append!(eltypes, [Float64, ComplexF64]) end