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 <lukas.sommer@codeplay.com>

Author: sommerlukas

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'

dpctl/_sycl_queue.pyx

@@ -59,6 +59,7 @@ from ._backend cimport (  # noqa: E211
     _queue_property_type,
 )
 from .memory._memory cimport _Memory
+from .experimental._work_group_memory cimport WorkGroupMemory
 
 import ctypes
 
@@ -242,6 +243,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()
 
@@ -824,6 +834,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.nbytes)
+                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,29 @@
+#                      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,23 @@
+#                      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
+
+cdef public api class WorkGroupMemory [object PyWorkGroupMemoryObject, type PyWorkGroupMemoryType]:
+    cdef Py_ssize_t nbytes
+
+

dpctl/experimental/_work_group_memory.pyx

@@ -0,0 +1,41 @@
+#                      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
+
+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):
+        self.nbytes = nbytes
+
+    property nbytes:
+        """Local memory size in bytes."""
+        def __get__(self):
+            return self.nbytes
+
+

dpctl/tests/test_sycl_kernel_submit.py

@@ -278,3 +278,4 @@ def test_kernel_arg_type():
     _check_kernel_arg_type_instance(kernel_arg_type.dpctl_float64)
     _check_kernel_arg_type_instance(kernel_arg_type.dpctl_void_ptr)
     _check_kernel_arg_type_instance(kernel_arg_type.dpctl_local_accessor)
+    _check_kernel_arg_type_instance(kernel_arg_type.dpctl_work_group_memory)

libsyclinterface/include/syclinterface/dpctl_sycl_enum_types.h

@@ -100,6 +100,7 @@ typedef enum
     DPCTL_FLOAT64_T,
     DPCTL_VOID_PTR,
     DPCTL_LOCAL_ACCESSOR,
+    DPCTL_WORK_GROUP_MEMORY,
     DPCTL_UNSUPPORTED_KERNEL_ARG
 } DPCTLKernelArgType;
 

libsyclinterface/source/dpctl_sycl_queue_interface.cpp

@@ -216,6 +216,14 @@ bool set_kernel_arg(handler &cgh,
     case DPCTL_LOCAL_ACCESSOR:
         arg_set = set_local_accessor_arg(cgh, idx, (MDLocalAccessor *)Arg);
         break;
+    case DPCTL_WORK_GROUP_MEMORY:
+    {
+        size_t num_bytes = reinterpret_cast<std::uintptr_t>(Arg);
+        sycl::ext::oneapi::experimental::work_group_memory<char[]> mem{
+            num_bytes, cgh};
+        cgh.set_arg(idx, mem);
+        break;
+    }
     default:
         arg_set = false;
         break;