Skip to content

Commit 7e85c4b

Browse files
committed
fix python test to work on cpu
1 parent c51bab7 commit 7e85c4b

File tree

1 file changed

+10
-6
lines changed

1 file changed

+10
-6
lines changed

tests/units/solver/pywrap/solver_utils.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,20 @@
11
import kokkos
22
import numpy as np
33

4-
# Dynamically determine the correct memory space and layout based on how Kokkos was compiled.
5-
# GPU builds (CUDA) expect LayoutLeft and UVM space to share memory with NumPy without segfaulting.
6-
# CPU builds expect LayoutRight and HostSpace.
7-
if hasattr(kokkos, "CudaUVMSpace"):
4+
# Dynamically determine the correct memory space and layout based on what
5+
# was actually compiled into the C++ pybind11 module, preventing CI/CD false positives.
6+
compiled_views = dir(kokkos.libpykokkos)
7+
has_uvm = any("CudaUVMSpace" in v for v in compiled_views)
8+
has_cuda = any("CudaSpace" in v for v in compiled_views) and not has_uvm
9+
has_hip = any("HIPManagedSpace" in v for v in compiled_views)
10+
11+
if has_uvm:
812
default_memspace = kokkos.CudaUVMSpace
913
default_layout = kokkos.LayoutLeft
10-
elif hasattr(kokkos, "CudaSpace"):
14+
elif has_cuda:
1115
default_memspace = kokkos.CudaSpace
1216
default_layout = kokkos.LayoutLeft
13-
elif hasattr(kokkos, "HIPManagedSpace"):
17+
elif has_hip:
1418
default_memspace = kokkos.HIPManagedSpace
1519
default_layout = kokkos.LayoutLeft
1620
else:

0 commit comments

Comments
 (0)