Add support for work_group_memory extension

dpctl/__init__.py

@@ -52,6 +52,7 @@
     SyclKernelSubmitError,
     SyclQueue,
     SyclQueueCreationError,
+    WorkGroupMemory,
 )
 from ._sycl_queue_manager import get_device_cached_queue
 from ._sycl_timer import SyclTimer
@@ -100,6 +101,7 @@
     "SyclKernelInvalidRangeError",
     "SyclKernelSubmitError",
     "SyclQueueCreationError",
+    "WorkGroupMemory",
 ]
 __all__ += [
     "get_device_cached_queue",

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,17 @@ 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.pxd

@@ -22,7 +22,12 @@
 
 from libcpp cimport bool as cpp_bool
 
-from ._backend cimport DPCTLSyclDeviceRef, DPCTLSyclQueueRef, _arg_data_type
+from ._backend cimport (
+    DPCTLSyclDeviceRef,
+    DPCTLSyclQueueRef,
+    DPCTLSyclWorkGroupMemoryRef,
+    _arg_data_type,
+)
 from ._sycl_context cimport SyclContext
 from ._sycl_device cimport SyclDevice
 from ._sycl_event cimport SyclEvent
@@ -98,3 +103,13 @@ cdef public api class SyclQueue (_SyclQueue) [
     cpdef prefetch(self, ptr, size_t count=*)
     cpdef mem_advise(self, ptr, size_t count, int mem)
     cpdef SyclEvent submit_barrier(self, dependent_events=*)
+
+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/_sycl_queue.pyx

@@ -54,13 +54,17 @@ from ._backend cimport (  # noqa: E211
     DPCTLSyclContextRef,
     DPCTLSyclDeviceSelectorRef,
     DPCTLSyclEventRef,
+    DPCTLWorkGroupMemory_Available,
+    DPCTLWorkGroupMemory_Create,
+    DPCTLWorkGroupMemory_Delete,
     _arg_data_type,
     _backend_type,
     _queue_property_type,
 )
 from .memory._memory cimport _Memory
 
 import ctypes
+import numbers
 
 from .enum_types import backend_type
 
@@ -250,6 +254,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 +862,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
@@ -1524,3 +1540,89 @@ cdef api SyclQueue SyclQueue_Make(DPCTLSyclQueueRef QRef):
     """
     cdef DPCTLSyclQueueRef copied_QRef = DPCTLQueue_Copy(QRef)
     return SyclQueue._create(copied_QRef)
+
+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.
+
+    This class is intended be used as kernel argument when launching kernels.
+
+    This is based on a DPC++ SYCL extension and only available in newer
+    versions. Use ``is_available()`` to check availability in your build.
+
+    There are multiple ways to create a `WorkGroupMemory`.
+
+    - If the constructor is invoked with just a single argument, this argument
+      is interpreted as the number of bytes to allocated in the shared local
+      memory.
+
+    - If the constructor is invoked with two arguments, the first argument is
+      interpreted as the datatype of the local memory, using the numpy type
+      naming scheme.
+      The second argument is interpreted as the number of elements to allocate.
+      The number of bytes to allocate is then computed from the byte size of
+      the data type and the element count.
+
+    Args:
+        args:
+            Variadic argument, see class documentation.
+
+    Raises:
+        TypeError: In case of incorrect arguments given to constructors,
+                   unexpected types of input arguments.
+    """
+    def __cinit__(self, *args):
+        cdef size_t nbytes
+        if not DPCTLWorkGroupMemory_Available():
+            raise RuntimeError("Workgroup memory extension not available")
+
+        if not (0 < len(args) < 3):
+            raise TypeError("WorkGroupMemory constructor takes 1 or 2 "
+                            f"arguments, but {len(args)} were given")
+
+        if len(args) == 1:
+            if not isinstance(args[0], numbers.Integral):
+                raise TypeError("WorkGroupMemory single argument constructor"
+                                "expects first argument to be `int`",
+                                f"but got {type(args[0])}")
+            nbytes = <size_t>(args[0])
+        else:
+            if not isinstance(args[0], str):
+                raise TypeError("WorkGroupMemory constructor expects first"
+                                f"argument to be `str`, but got {type(args[0])}")
+            if not isinstance(args[1], numbers.Integral):
+                raise TypeError("WorkGroupMemory constructor expects second"
+                                f"argument to be `int`, but got {type(args[1])}")
+            dtype = <str>(args[0])
+            count = <size_t>(args[1])
+            if not dtype[0] in ["i", "u", "f"]:
+                raise TypeError(f"Unrecognized type value: '{dtype}'")
+            try:
+                bit_width = int(dtype[1:])
+            except ValueError:
+                raise TypeError(f"Unrecognized type value: '{dtype}'")
+
+            byte_size = <size_t>bit_width
+            nbytes = count * byte_size
+
+        self._mem_ref = DPCTLWorkGroupMemory_Create(nbytes)
+
+    """Check whether the work_group_memory extension is available"""
+    @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/apis/include/dpctl_capi.h

@@ -25,9 +25,11 @@
 #pragma once
 
 // clang-format off
-// Ordering of includes is important here. dpctl_sycl_types defines types
-// used by dpctl's Python C-API headers.
+// Ordering of includes is important here. dpctl_sycl_types and
+// dpctl_sycl_extension_interface define types used by dpctl's Python
+// C-API headers.
 #include "syclinterface/dpctl_sycl_types.h"
+#include "syclinterface/dpctl_sycl_extension_interface.h"
 #ifdef __cplusplus
 #define CYTHON_EXTERN_C extern "C"
 #else

dpctl/sycl.pxd

@@ -42,6 +42,10 @@ cdef extern from "sycl/sycl.hpp" namespace "sycl":
         "sycl::kernel_bundle<sycl::bundle_state::executable>":
         pass
 
+cdef extern from "syclinterface/dpctl_sycl_extension_interface.h":
+    cdef struct RawWorkGroupMemoryTy
+    ctypedef RawWorkGroupMemoryTy RawWorkGroupMemory
+
 cdef extern from "syclinterface/dpctl_sycl_type_casters.hpp" \
     namespace "dpctl::syclinterface":
     # queue
@@ -67,3 +71,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)

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)

dpctl/tests/test_work_group_memory.py

@@ -0,0 +1,90 @@
+#                      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.
+
+"""Defines unit test cases for the work_group_memory in a SYCL kernel"""
+
+import os
+
+import pytest
+
+import dpctl
+import dpctl.tensor
+
+
+def get_spirv_abspath(fn):
+    curr_dir = os.path.dirname(os.path.abspath(__file__))
+    spirv_file = os.path.join(curr_dir, "input_files", fn)
+    return spirv_file
+
+
+# The kernel in the SPIR-V file used in this test was generated from the
+# following SYCL source code:
+# #include <sycl/sycl.hpp>
+# using namespace sycl;
+# namespace syclexp = sycl::ext::oneapi::experimental;
+# namespace syclext = sycl::ext::oneapi;
+# using data_t = int32_t;
+#
+# extern "C" SYCL_EXTERNAL
+# SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((syclexp::nd_range_kernel<1>))
+# void local_mem_kernel(data_t* in, data_t* out,
+#                           syclexp::work_group_memory<data_t> mem){
+#   auto* local_mem = &mem;
+#   auto item = syclext::this_work_item::get_nd_item<1>();
+#   size_t global_id = item.get_global_linear_id();
+#   size_t local_id = item.get_local_linear_id();
+#   local_mem[local_id] = in[global_id];
+#   out[global_id] = local_mem[local_id];
+# }
+
+
+def test_submit_work_group_memory():
+    if not dpctl.WorkGroupMemory.is_available():
+        pytest.skip("Work group memory extension not supported")
+
+    try:
+        q = dpctl.SyclQueue("level_zero")
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("LevelZero queue could not be created")
+    spirv_file = get_spirv_abspath("work-group-memory-kernel.spv")
+    with open(spirv_file, "br") as spv:
+        spv_bytes = spv.read()
+    prog = dpctl.program.create_program_from_spirv(q, spv_bytes)
+    kernel = prog.get_sycl_kernel("__sycl_kernel_local_mem_kernel")
+    local_size = 16
+    global_size = local_size * 8
+
+    x = dpctl.tensor.ones(global_size, dtype="int32")
+    y = dpctl.tensor.zeros(global_size, dtype="int32")
+    x.sycl_queue.wait()
+    y.sycl_queue.wait()
+
+    try:
+        q.submit(
+            kernel,
+            [
+                x.usm_data,
+                y.usm_data,
+                dpctl.WorkGroupMemory("i4", local_size),
+            ],
+            [global_size],
+            [local_size],
+        )
+        q.wait()
+    except dpctl._sycl_queue.SyclKernelSubmitError:
+        pytest.skip(f"Kernel submission to {q.sycl_device} failed")
+
+    assert dpctl.tensor.all(x == y)