Do not place in experimental

ndgrigorian · Mar 7, 2025 · Dec 9, 2024 · Feb 11, 2025 · Feb 12, 2025 · Feb 25, 2025
commit 5611abfc517aeed65246443f57a457c89cd6e50c
diff --git a/dpctl/CMakeLists.txt b/dpctl/CMakeLists.txt
@@ -207,4 +207,3 @@ add_subdirectory(program)
 add_subdirectory(memory)
 add_subdirectory(tensor)
 add_subdirectory(utils)
-add_subdirectory(experimental)
@@ -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",

@@ -474,7 +474,6 @@ cdef extern from "syclinterface/dpctl_sycl_extension_interface.h":
     cdef struct RawWorkGroupMemoryTy
     ctypedef RawWorkGroupMemoryTy RawWorkGroupMemory
 
-
     cdef struct DPCTLOpaqueWorkGroupMemory
     ctypedef DPCTLOpaqueWorkGroupMemory *DPCTLSyclWorkGroupMemoryRef;
 

@@ -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
@@ -54,11 +54,13 @@ from ._backend cimport (  # noqa: E211
     DPCTLSyclContextRef,
     DPCTLSyclDeviceSelectorRef,
     DPCTLSyclEventRef,
+    DPCTLWorkGroupMemory_Available,
+    DPCTLWorkGroupMemory_Create,
+    DPCTLWorkGroupMemory_Delete,
     _arg_data_type,
     _backend_type,
     _queue_property_type,
 )
-from .experimental._work_group_memory cimport WorkGroupMemory
 from .memory._memory cimport _Memory
 
 import ctypes
@@ -1537,3 +1539,41 @@ 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 is based on a DPC++ SYCL extension and only available in newer
+    versions. Use ``is_available()`` to check availability in your build.
+
+    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)
+
+    """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
@@ -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

diff --git a/dpctl/experimental/CMakeLists.txt b/dpctl/experimental/CMakeLists.txt
diff --git a/dpctl/experimental/__init__.pxd b/dpctl/experimental/__init__.pxd
diff --git a/dpctl/experimental/__init__.py b/dpctl/experimental/__init__.py
diff --git a/dpctl/experimental/_work_group_memory.pxd b/dpctl/experimental/_work_group_memory.pxd
diff --git a/dpctl/experimental/_work_group_memory.pyx b/dpctl/experimental/_work_group_memory.pyx
@@ -52,7 +52,7 @@ def get_spirv_abspath(fn):
 
 
 def test_submit_work_group_memory():
-    if not dpctl.experimental.WorkGroupMemory.is_available():
+    if not dpctl.WorkGroupMemory.is_available():
         pytest.skip("Work group memory extension not supported")
 
     try:
@@ -78,7 +78,7 @@ def test_submit_work_group_memory():
             [
                 x.usm_data,
                 y.usm_data,
-                dpctl.experimental.WorkGroupMemory(local_size * x.itemsize),
+                dpctl.WorkGroupMemory(local_size * x.itemsize),
             ],
             [global_size],
             [local_size],

@@ -38,7 +38,7 @@
 
 
 def test_submit_work_group_memory_opencl():
-    if not dpctl.experimental.WorkGroupMemory.is_available():
+    if not dpctl.WorkGroupMemory.is_available():
         pytest.skip("Work group memory extension not supported")
 
     try:
@@ -65,7 +65,7 @@ def test_submit_work_group_memory_opencl():
             [
                 x_dev,
                 y_dev,
-                dpctl.experimental.WorkGroupMemory(local_size * x.itemsize),
+                dpctl.WorkGroupMemory(local_size * x.itemsize),
             ],
             [global_size],
             [local_size],