Cache per-device element-type support

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
 
 """

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

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)
 

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)

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