Support keys() and unsafe_cached_load()

Author: michel2323

lib/level-zero/device.jl

@@ -204,6 +204,8 @@ Base.length(iter::ZeDevices) = length(iter.handles)
 
 Base.IteratorSize(::ZeDevices) = Base.HasLength()
 
+Base.keys(iter::ZeDevices) = 1:length(iter)
+
 function Base.show(io::IO, ::MIME"text/plain", iter::ZeDevices)
     print(io, "ZeDevice iterator for $(length(iter)) devices")
     if !isempty(iter)

src/context.jl

@@ -103,7 +103,13 @@ See also: [`device`](@ref), [`devices`](@ref)
 function device!(drv::ZeDevice)
     task_local_storage(:ZeDevice, drv)
 end
-device!(i::Int) = device!(devices(driver())[i])
+function device!(i::Int)
+    devs = devices(driver())
+    if i < 1 || i > length(devs)
+        throw(ArgumentError("Invalid device index $i (must be between 1 and $(length(devs)))"))
+    end
+    device!(devs[i])
+end
 
 const global_contexts = Dict{ZeDriver,ZeContext}()
 

src/device/array.jl

@@ -195,6 +195,13 @@ end
     end
 end
 
+@device_function @inline function unsafe_cached_load(ptr::LLVMPtr{T,A}, i::Integer, align::Val) where {T,A}
+    # For SPIR-V/Level Zero, we don't have explicit cache control intrinsics like CUDA's __ldg
+    # So we fall back to a regular unsafe_load. The SPIR-V compiler may still apply
+    # appropriate optimizations based on context.
+    unsafe_load(ptr, i, align)
+end
+
 @device_function @inline function const_arrayref(A::oneDeviceArray{T}, index::Integer) where {T}
     # simplified bounds check (see `arrayset`)
     #@boundscheck checkbounds(A, index)