Add work group memory to libsyclinterface

Author: sommerlukas

dpctl/_backend.pxd

@@ -469,3 +469,19 @@ 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

@@ -835,7 +835,7 @@ cdef class SyclQueue(_SyclQueue):
                 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)
+                kargs[idx] = <void*>(<size_t>arg._ref)
                 kargty[idx] = _arg_data_type._WORK_GROUP_MEMORY
             else:
                 ret = -1

dpctl/experimental/_work_group_memory.pxd

@@ -17,7 +17,15 @@
 # distutils: language = c++
 # cython: language_level=3
 
-cdef public api class WorkGroupMemory [object PyWorkGroupMemoryObject, type PyWorkGroupMemoryType]:
-    cdef Py_ssize_t nbytes
+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

@@ -18,6 +18,17 @@
 # cython: language_level=3
 # cython: linetrace=True
 
+from .._backend cimport (
+    DPCTLWorkGroupMemory_Create,
+    DPCTLWorkGroupMemory_Delete,
+    DPCTLWorkGroupMemory_Available
+)
+
+cdef class _WorkGroupMemory:
+    def __dealloc__(self):
+        if(self._mem_ref):
+            DPCTLWorkGroupMemory_Delete(self._mem_ref)
+
 cdef class WorkGroupMemory:
     """
     WorkGroupMemory(nbytes)
@@ -31,11 +42,20 @@ cdef class WorkGroupMemory:
             Expected to be positive.
     """
     def __cinit__(self, Py_ssize_t nbytes):
-        self.nbytes = 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 nbytes:
-        """Local memory size in bytes."""
+    property _ref:
+        """Returns the address of the C API ``DPCTLWorkGroupMemoryRef``
+        pointer as a ``size_t``.
+        """
         def __get__(self):
-            return self.nbytes
+            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)

libsyclinterface/include/syclinterface/dpctl_sycl_extension_interface.h

@@ -0,0 +1,55 @@
+//===---- dpctl_sycl_extension_interface.h - C API for SYCL ext  -*-C++-*- ===//
+//
+//                      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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This header declares a C interface to SYCL language extensions defined by
+/// DPC++.
+///
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include "Support/DllExport.h"
+#include "Support/ExternC.h"
+#include "Support/MemOwnershipAttrs.h"
+#include "dpctl_data_types.h"
+#include "dpctl_error_handler_type.h"
+#include "dpctl_sycl_enum_types.h"
+#include "dpctl_sycl_types.h"
+
+DPCTL_C_EXTERN_C_BEGIN
+
+typedef struct RawWorkGroupMemoryTy {
+  size_t nbytes;
+} RawWorkGroupMemory;
+
+typedef struct DPCTLOpaqueSyclWorkGroupMemory *DPCTLSyclWorkGroupMemoryRef;
+
+DPCTL_API
+__dpctl_give DPCTLSyclWorkGroupMemoryRef
+DPCTLWorkGroupMemory_Create(size_t nbytes);
+
+DPCTL_API
+void DPCTLWorkGroupMemory_Delete(__dpctl_take DPCTLSyclWorkGroupMemoryRef Ref);
+
+DPCTL_API
+bool DPCTLWorkGroupMemory_Available();
+
+DPCTL_C_EXTERN_C_END

libsyclinterface/include/syclinterface/dpctl_sycl_type_casters.hpp

@@ -80,6 +80,11 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(std::vector<DPCTLSyclPlatformRef>,
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(std::vector<DPCTLSyclEventRef>,
                                    DPCTLEventVectorRef)
 
+#include "dpctl_sycl_extension_interface.h"
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(
+    RawWorkGroupMemory,
+    DPCTLSyclWorkGroupMemoryRef)
+
 #endif
 
 } // namespace dpctl::syclinterface

libsyclinterface/source/dpctl_sycl_extension_interface.cpp

@@ -0,0 +1,65 @@
+//===---- dpctl_sycl_extension_interface.cpp - Implements C API for SYCL ext =//
+//
+//                      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.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This file implements the data types and functions declared in
+/// dpctl_sycl_extension_interface.h.
+///
+//===----------------------------------------------------------------------===//
+
+#include "dpctl_sycl_extension_interface.h"
+
+#include "dpctl_error_handlers.h"
+#include "dpctl_sycl_type_casters.hpp"
+
+#include <sycl/sycl.hpp>
+
+using namespace dpctl::syclinterface;
+
+DPCTL_API
+__dpctl_give DPCTLSyclWorkGroupMemoryRef
+DPCTLWorkGroupMemory_Create(size_t nbytes)
+{
+    DPCTLSyclWorkGroupMemoryRef wgm = nullptr;
+    try {
+        auto WorkGroupMem = new RawWorkGroupMemory{nbytes};
+        wgm = wrap<RawWorkGroupMemory>(WorkGroupMem);
+    } catch (std::exception const &e) {
+        error_handler(e, __FILE__, __func__, __LINE__);
+    }
+    return wgm;
+}
+
+DPCTL_API
+void DPCTLWorkGroupMemory_Delete(__dpctl_take DPCTLSyclWorkGroupMemoryRef Ref)
+{
+    delete unwrap<RawWorkGroupMemory>(Ref);
+}
+
+DPCTL_API
+bool DPCTLWorkGroupMemory_Available()
+{
+#ifdef SYCL_EXT_ONEAPI_WORK_GROUP_MEMORY
+    return true;
+#else
+    return false;
+#endif
+}
+

libsyclinterface/source/dpctl_sycl_queue_interface.cpp

@@ -42,6 +42,10 @@
 #include <sycl/sycl.hpp> /* SYCL headers   */
 #include <utility>
 
+#ifdef SYCL_EXT_ONEAPI_WORK_GROUP_MEMORY
+#include "dpctl_sycl_extension_interface.h"
+#endif
+
 using namespace sycl;
 
 #define SET_LOCAL_ACCESSOR_ARG(CGH, NDIM, ARGTY, R, IDX)                       \
@@ -216,14 +220,18 @@ bool set_kernel_arg(handler &cgh,
     case DPCTL_LOCAL_ACCESSOR:
         arg_set = set_local_accessor_arg(cgh, idx, (MDLocalAccessor *)Arg);
         break;
+#ifdef SYCL_EXT_ONEAPI_WORK_GROUP_MEMORY
     case DPCTL_WORK_GROUP_MEMORY:
     {
-        size_t num_bytes = reinterpret_cast<std::uintptr_t>(Arg);
+        auto ref = static_cast<DPCTLSyclWorkGroupMemoryRef>(Arg);
+        RawWorkGroupMemory *raw_mem = unwrap<RawWorkGroupMemory>(ref);
+        size_t num_bytes = raw_mem->nbytes;
         sycl::ext::oneapi::experimental::work_group_memory<char[]> mem{
             num_bytes, cgh};
         cgh.set_arg(idx, mem);
         break;
     }
+#endif
     default:
         arg_set = false;
         break;