Allow construction with data type

Author: sommerlukas

.flake8

@@ -31,7 +31,6 @@ 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/_sycl_queue.pyx

@@ -1555,15 +1555,47 @@ cdef class WorkGroupMemory:
     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:
-        nbytes (int)
-            number of bytes to allocate in local memory.
-            Expected to be positive.
+        args:
+            Variadic argument, see class documentation.
+
+    Raises:
+        TypeError: In case of incorrect arguments given to constructors,
+                   unexpected types of input arguments.
     """
-    def __cinit__(self, Py_ssize_t nbytes):
+    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:
+            nbytes = <size_t>(args[0])
+        else:
+            dtype = <str>(args[0])
+            count = <size_t>(args[1])
+            ty = dtype[0]
+            if not ty in ["i", "u", "f"]:
+                raise TypeError(f"Unrecognized type value: '{dtype}'")
+            byte_size = <size_t>(int(dtype[1:]))
+            nbytes = count * byte_size
+
         self._mem_ref = DPCTLWorkGroupMemory_Create(nbytes)
 
     """Check whether the work_group_memory extension is available"""

dpctl/tests/test_work_group_memory.py

@@ -78,7 +78,7 @@ def test_submit_work_group_memory():
             [
                 x.usm_data,
                 y.usm_data,
-                dpctl.WorkGroupMemory(local_size * x.itemsize),
+                dpctl.WorkGroupMemory("i4", local_size),
             ],
             [global_size],
             [local_size],