Inspect extended kernel properties.

Author: maleadt

lib/level-zero/memory.jl

@@ -50,11 +50,11 @@ function device_alloc(ctx::ZeContext, dev::ZeDevice, bytesize::Integer,
         flags = ZE_RELAXED_ALLOCATION_LIMITS_EXP_FLAG_MAX_SIZE,
     ))
     GC.@preserve relaxed_allocation_ref begin
-        desc_ref = if bytesize > properties(dev).maxMemAllocSize
-            pNext = Base.unsafe_convert(Ptr{Cvoid}, relaxed_allocation_ref)
-            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal, pNext))
+        if bytesize > properties(dev).maxMemAllocSize
+            desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
+            link_extensions(desc_ref, relaxed_allocation_ref)
         else
-            Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
+            desc_ref = Ref(ze_device_mem_alloc_desc_t(; flags, ordinal))
         end
 
         ptr_ref = Ref{Ptr{Cvoid}}()

lib/level-zero/module.jl

@@ -234,6 +234,16 @@ export properties
 
 function properties(kernel::ZeKernel)
     props_ref = Ref(ze_kernel_properties_t())
+    preferred_group_size_props_ref = Ref(ze_kernel_preferred_group_size_properties_t())
+    link_extensions(props_ref, preferred_group_size_props_ref)
+    if haskey(oneL0.extension_properties(kernel.mod.context.driver),
+              "ZE_extension_kernel_max_group_size_properties")
+        # TODO: memoize
+        max_group_size_props_ref = Ref(ze_kernel_max_group_size_properties_ext_t())
+        link_extensions(preferred_group_size_props_ref, max_group_size_props_ref)
+    else
+        max_group_size_props_ref = nothing
+    end
     zeKernelGetProperties(kernel, props_ref)
 
     props = props_ref[]
@@ -251,6 +261,9 @@ function properties(kernel::ZeKernel)
         spillMemSize=Int(props.spillMemSize),
         kernel_uuid=Base.UUID(reinterpret(UInt128, [props.uuid.kid...])[1]),
         module_uuid=Base.UUID(reinterpret(UInt128, [props.uuid.mid...])[1]),
+        preferredGroupSize=Int(preferred_group_size_props_ref[].preferredMultiple),
+        maxGroupSize=max_group_size_props_ref === nothing ? missing :
+            Int(max_group_size_props_ref[].maxGroupSize)
     )
 end
 

lib/level-zero/oneL0.jl

@@ -69,6 +69,20 @@ zeroinit(::Type{_ze_native_kernel_uuid_t}) =
 zeroinit(::Type{ze_kernel_uuid_t}) =
     ze_kernel_uuid_t(ntuple(_->zero(UInt8), 16), ntuple(_->zero(UInt8), 16))
 
+# link extension objects in pNext
+function link_extensions(refs...)
+    length(refs) >= 2 || return
+    for (parent, child) in zip(refs[1:end-1], refs[2:end])
+        pNext = Base.unsafe_convert(Ptr{Cvoid}, child)
+        typ = eltype(parent)
+        @assert fieldnames(typ)[2] == :pNext
+        field = Base.unsafe_convert(Ptr{Cvoid}, parent) + fieldoffset(typ, 2)
+        field = convert(Ptr{Ptr{Cvoid}}, field)
+        unsafe_store!(field, pNext)
+    end
+    return
+end
+
 # core wrappers
 include("error.jl")
 include("common.jl")