Add support for work_group_memory extension

Extend kernel argument handling to add support for the
work_group_memory extension, allowing users to dynamically allocate
local memory for a kernel.

Signed-off-by: Lukas Sommer <[email protected]>

Author: sommerlukas

.flake8

@@ -31,6 +31,7 @@ per-file-ignores =
     dpctl/utils/_compute_follows_data.pyx: E999, E225, E227
     dpctl/utils/_onetrace_context.py: E501, W505
     dpctl/tensor/_array_api.py: E501, W505
+    dpctl/experimental/_work_group_memory.pyx: E999
     examples/cython/sycl_buffer/syclbuffer/_syclbuffer.pyx: E999, E225, E402
     examples/cython/usm_memory/blackscholes/_blackscholes_usm.pyx: E999, E225, E226, E402
     examples/cython/use_dpctl_sycl/use_dpctl_sycl/_cython_api.pyx: E999, E225, E226, E402

dpctl/CMakeLists.txt

@@ -207,3 +207,4 @@ add_subdirectory(program)
 add_subdirectory(memory)
 add_subdirectory(tensor)
 add_subdirectory(utils)
+add_subdirectory(experimental)

dpctl/_backend.pxd

@@ -69,7 +69,8 @@ cdef extern from "syclinterface/dpctl_sycl_enum_types.h":
         _FLOAT              'DPCTL_FLOAT32_T',
         _DOUBLE             'DPCTL_FLOAT64_T',
         _VOID_PTR           'DPCTL_VOID_PTR',
-        _LOCAL_ACCESSOR     'DPCTL_LOCAL_ACCESSOR'
+        _LOCAL_ACCESSOR     'DPCTL_LOCAL_ACCESSOR',
+        _WORK_GROUP_MEMORY  'DPCTL_WORK_GROUP_MEMORY'
 
     ctypedef enum _queue_property_type 'DPCTLQueuePropertyType':
         _DEFAULT_PROPERTY   'DPCTL_DEFAULT_PROPERTY'
@@ -468,3 +469,18 @@ cdef extern from "syclinterface/dpctl_sycl_usm_interface.h":
     cdef DPCTLSyclDeviceRef DPCTLUSM_GetPointerDevice(
         DPCTLSyclUSMRef MRef,
         DPCTLSyclContextRef CRef)
+
+cdef extern from "syclinterface/dpctl_sycl_extension_interface.h":
+    cdef struct RawWorkGroupMemoryTy
+    ctypedef RawWorkGroupMemoryTy RawWorkGroupMemory
+
+
+    cdef struct DPCTLOpaqueWorkGroupMemory
+    ctypedef DPCTLOpaqueWorkGroupMemory *DPCTLSyclWorkGroupMemoryRef;
+
+    cdef DPCTLSyclWorkGroupMemoryRef DPCTLWorkGroupMemory_Create(size_t nbytes);
+
+    cdef void DPCTLWorkGroupMemory_Delete(
+        DPCTLSyclWorkGroupMemoryRef Ref);
+
+    cdef bint DPCTLWorkGroupMemory_Available();

dpctl/_sycl_queue.pyx

@@ -58,6 +58,7 @@ from ._backend cimport (  # noqa: E211
     _backend_type,
     _queue_property_type,
 )
+from .experimental._work_group_memory cimport WorkGroupMemory
 from .memory._memory cimport _Memory
 
 import ctypes
@@ -250,6 +251,15 @@ cdef class _kernel_arg_type:
             _arg_data_type._LOCAL_ACCESSOR
         )
 
+    @property
+    def dpctl_work_group_memory(self):
+        cdef str p_name = "dpctl_work_group_memory"
+        return kernel_arg_type_attribute(
+            self._name,
+            p_name,
+            _arg_data_type._WORK_GROUP_MEMORY
+        )
+
 
 kernel_arg_type = _kernel_arg_type()
 
@@ -849,6 +859,9 @@ cdef class SyclQueue(_SyclQueue):
             elif isinstance(arg, _Memory):
                 kargs[idx]= <void*>(<size_t>arg._pointer)
                 kargty[idx] = _arg_data_type._VOID_PTR
+            elif isinstance(arg, WorkGroupMemory):
+                kargs[idx] = <void*>(<size_t>arg._ref)
+                kargty[idx] = _arg_data_type._WORK_GROUP_MEMORY
             else:
                 ret = -1
         return ret

dpctl/experimental/CMakeLists.txt

@@ -0,0 +1,7 @@
+file(GLOB _cython_sources *.pyx)
+foreach(_cy_file ${_cython_sources})
+    get_filename_component(_trgt ${_cy_file} NAME_WLE)
+    build_dpctl_ext(${_trgt} ${_cy_file} "dpctl/experimental" RELATIVE_PATH "..")
+    target_include_directories(${_trgt} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
+    target_link_libraries(DpctlCAPI INTERFACE ${_trgt}_headers)
+endforeach()

dpctl/experimental/__init__.pxd

@@ -0,0 +1,25 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2025 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# distutils: language = c++
+# cython: language_level=3
+
+"""This file declares the extension types and functions for the Cython API
+implemented in dpctl.experimental.*.pyx.
+"""
+
+
+from dpctl.experimental._work_group_memory cimport *

dpctl/experimental/__init__.py

@@ -0,0 +1,27 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2025 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+    **Data Parallel Control Experimental" provides Python objects to interface
+    with different experimental SYCL language extensions defined by the DPC++
+    SYCL implementation.
+"""
+
+from ._work_group_memory import WorkGroupMemory
+
+__all__ = [
+    "WorkGroupMemory",
+]

dpctl/experimental/_work_group_memory.pxd

@@ -0,0 +1,31 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2025 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# distutils: language = c++
+# cython: language_level=3
+
+from .._backend cimport DPCTLSyclWorkGroupMemoryRef
+
+
+cdef public api class _WorkGroupMemory [
+    object Py_WorkGroupMemoryObject, type Py_WorkGroupMemoryType
+]:
+    cdef DPCTLSyclWorkGroupMemoryRef _mem_ref
+
+cdef public api class WorkGroupMemory(_WorkGroupMemory) [
+    object PyWorkGroupMemoryObject, type PyWorkGroupMemoryType
+]:
+    pass

dpctl/experimental/_work_group_memory.pyx

@@ -0,0 +1,60 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2025 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# distutils: language = c++
+# cython: language_level=3
+# cython: linetrace=True
+
+from .._backend cimport (
+    DPCTLWorkGroupMemory_Available,
+    DPCTLWorkGroupMemory_Create,
+    DPCTLWorkGroupMemory_Delete,
+)
+
+
+cdef class _WorkGroupMemory:
+    def __dealloc__(self):
+        if(self._mem_ref):
+            DPCTLWorkGroupMemory_Delete(self._mem_ref)
+
+cdef class WorkGroupMemory:
+    """
+    WorkGroupMemory(nbytes)
+    Python class representing the ``work_group_memory`` class from the
+    Workgroup Memory oneAPI SYCL extension for low-overhead allocation of local
+    memory shared by the workitems in a workgroup.
+
+    Args:
+        nbytes (int)
+            number of bytes to allocate in local memory.
+            Expected to be positive.
+    """
+    def __cinit__(self, Py_ssize_t nbytes):
+        if not DPCTLWorkGroupMemory_Available():
+            raise RuntimeError("Workgroup memory extension not available")
+
+        self._mem_ref = DPCTLWorkGroupMemory_Create(nbytes)
+
+    @staticmethod
+    def is_available():
+        return DPCTLWorkGroupMemory_Available()
+
+    property _ref:
+        """Returns the address of the C API ``DPCTLWorkGroupMemoryRef``
+        pointer as a ``size_t``.
+        """
+        def __get__(self):
+            return <size_t>self._mem_ref

dpctl/sycl.pxd

@@ -67,3 +67,12 @@ cdef extern from "syclinterface/dpctl_sycl_type_casters.hpp" \
         "dpctl::syclinterface::wrap<sycl::event>" (const event *)
     cdef event * unwrap_event "dpctl::syclinterface::unwrap<sycl::event>" (
         dpctl_backend.DPCTLSyclEventRef)
+
+    # work group memory extension[
+    cdef dpctl_backend.DPCTLSyclWorkGroupMemoryRef wrap_work_group_memory \
+            "dpctl::syclinterface::wrap<RawWorkGroupMemory>" \
+            (const RawWorkGroupMemory *)
+
+    cdef RawWorkGroupMemory * unwrap_work_group_memory \
+            "dpctl::syclinterface::unwrap<RawWorkGroupMemory>" (
+            dpctl_backend.DPCTLSyclWorkGroupMemoryRef)