Python test work group memory

Author: sommerlukas

dpctl/tests/input_files/work-group-memory-kernel.spv

@@ -0,0 +1,70 @@
+#                      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 SyclProgram and SyclKernel classes
+"""
+
+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
+
+
+def test_submit_work_group_memory():
+    if not dpctl.experimental.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.experimental.WorkGroupMemory(local_size * x.itemsize),
+            ],
+            [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)