fix affinity mask setting

Signed-off-by: Mateusz P. Nowak <[email protected]>

Author: mateuszpn

sycl/test-e2e/AddressSanitizer/lit.local.cfg

@@ -31,4 +31,7 @@ unsupported_san_flags = [
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
 
-config.environment["ZE_AFFINITY_MASK"] = "0"
+if hasattr(config, 'ze_affinity_mask') and config.ze_affinity_mask is not None:
+    config.environment["ZE_AFFINITY_MASK"] = config.ze_affinity_mask
+else:
+    config.environment["ZE_AFFINITY_MASK"] = "0"

sycl/test-e2e/MemorySanitizer/lit.local.cfg

@@ -39,4 +39,9 @@ unsupported_san_flags = [
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
 
-config.environment["ZE_AFFINITY_MASK"] = "0"
+# Set ZE_AFFINITY_MASK for single-GPU execution (required by device sanitizer).
+# Use the first Level Zero GPU device number if available, otherwise default to "0".
+if hasattr(config, 'ze_affinity_mask') and config.ze_affinity_mask is not None:
+    config.environment["ZE_AFFINITY_MASK"] = config.ze_affinity_mask
+else:
+    config.environment["ZE_AFFINITY_MASK"] = "0"

sycl/test-e2e/ThreadSanitizer/lit.local.cfg

@@ -34,4 +34,7 @@ unsupported_san_flags = [
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
 
-config.environment["ZE_AFFINITY_MASK"] = "0"
+if hasattr(config, 'ze_affinity_mask') and config.ze_affinity_mask is not None:
+    config.environment["ZE_AFFINITY_MASK"] = config.ze_affinity_mask
+else:
+    config.environment["ZE_AFFINITY_MASK"] = "0"

sycl/test-e2e/format.py

@@ -283,6 +283,11 @@ def get_extra_env(sycl_devices):
             # so that device might still be accessible to some of the tests yet
             # we won't set the environment variable below for such scenario.
             extra_env = []
+
+            # Include ZE_AFFINITY_MASK if set by lit.local.cfg (e.g., for sanitizer tests)
+            if "ZE_AFFINITY_MASK" in test.config.environment:
+                extra_env.append("ZE_AFFINITY_MASK={}".format(test.config.environment["ZE_AFFINITY_MASK"]))
+
             if "level_zero:gpu" in sycl_devices and litConfig.params.get("ur_l0_debug"):
                 extra_env.append("UR_L0_DEBUG={}".format(test.config.ur_l0_debug))
 
@@ -334,14 +339,6 @@ def get_extra_env(sycl_devices):
                 expanded = "env"
 
                 extra_env = get_extra_env([parsed_dev_name])
-                backend, device = parsed_dev_name.split(":", 1)
-                device_selector = parsed_dev_name
-                if backend == "level_zero" and device.isdigit():
-                    # level_zero:0 maps to a selector for the first GPU, so pin
-                    # execution to the requested device via ZE_AFFINITY_MASK.
-                    extra_env.append(f"ZE_AFFINITY_MASK={device}")
-                    device_selector = f"{backend}:0"
-
                 if extra_env:
                     expanded += " {}".format(" ".join(extra_env))
 
@@ -350,6 +347,12 @@ def get_extra_env(sycl_devices):
                 elif "level_zero_v1" in full_dev_name:
                     expanded += " env UR_LOADER_USE_LEVEL_ZERO_V2=0"
 
+                # If ZE_AFFINITY_MASK is set, it filters devices so we should use :0
+                device_selector = parsed_dev_name
+                if "ZE_AFFINITY_MASK" in test.config.environment:
+                    backend, _ = parsed_d ev_name.split(":", 1)
+                    device_selector = f"{backend}:0"
+
                 expanded += " ONEAPI_DEVICE_SELECTOR={} {}".format(
                     device_selector, test.config.run_launcher
                 )

sycl/test-e2e/lit.cfg.py

@@ -9,9 +9,6 @@
 import shlex
 import shutil
 
-import lit.formats
-import lit.util
-
 from lit.llvm import llvm_config
 from lit.llvm.subst import ToolSubst, FindTool
 
@@ -950,6 +947,16 @@ def get_sycl_ls_verbose(sycl_device, env):
 
 config.sycl_devices = filtered_sycl_devices
 
+# Determine ZE_AFFINITY_MASK for Level Zero devices.
+# Sanitizer tests need to set ZE_AFFINITY_MASK to a single device index
+config.ze_affinity_mask = None
+for sycl_device in remove_level_zero_suffix(config.sycl_devices):
+    be, dev = sycl_device.split(":")
+    if be == "level_zero" and dev.isdigit():
+        if config.ze_affinity_mask is None:
+            config.ze_affinity_mask = dev
+            break
+
 for sycl_device in remove_level_zero_suffix(config.sycl_devices):
     be, dev = sycl_device.split(":")
     config.available_features.add("any-device-is-" + dev)