def handle_args(is_for: bool, *args) -> HandledArgs:
"""
Handle the *args passed to parallel_* functions
:param is_for: whether the arguments belong to a parallel_for call
:param *args: the list of arguments being checked
:returns: a HandledArgs object containing the passed arguments
"""
unpacked: Tuple = tuple(*args)
name: Optional[str] = None
policy: Union[ExecutionPolicy, int]
workunit: Callable
view: Optional[ViewType] = None
initial_value: Union[int, float] = 0
if len(unpacked) == 2:
policy = unpacked[0]
workunit = unpacked[1]
elif len(unpacked) == 3:
if isinstance(unpacked[0], str):
name = unpacked[0]
policy = unpacked[1]
workunit = unpacked[2]
elif is_for and isinstance(unpacked[2], ViewType):
policy = unpacked[0]
workunit = unpacked[1]
view = unpacked[2]
elif isinstance(unpacked[2], (int, float)):
policy = unpacked[0]
workunit = unpacked[1]
initial_value = unpacked[2]
else:
raise TypeError(f"ERROR: wrong arguments {unpacked}")
elif len(unpacked) == 4:
if isinstance(unpacked[0], str):
name = unpacked[0]
policy = unpacked[1]
workunit = unpacked[2]
if is_for and isinstance(unpacked[3], ViewType):
view = unpacked[3]
elif isinstance(unpacked[3], (int, float)):
initial_value = unpacked[3]
else:
raise TypeError(f"ERROR: wrong arguments {unpacked}")
else:
raise TypeError(f"ERROR: wrong arguments {unpacked}")
else:
raise ValueError(f"ERROR: incorrect number of arguments {len(unpacked)}")
if isinstance(policy, int):
policy = RangePolicy(km.get_default_space(), 0, policy)
return HandledArgs(name, policy, workunit, view, initial_value)
def is_debug(self, space: ExecutionSpace) -> bool:
"""
Check if the execution space is Debug and account for Default space
:param space: the execution space
:returns: True or False
"""
return space is ExecutionSpace.Debug or (
space is ExecutionSpace.Default
and km.get_default_space() is ExecutionSpace.Debug)
def get_output_dir(self, main: Path, metadata: EntityMetadata,
space: ExecutionSpace) -> Path:
"""
Get the output directory for an execution space
:param main: the path to the main file in the current PyKokkos application
:param metadata: the metadata of the entity being compiled
:param space: the execution space to compile for
:returns: the path to the output directory for a specific execution space
"""
if space is ExecutionSpace.Default:
space = km.get_default_space()
return self.get_entity_dir(main, metadata) / space.value
def compile_entity(self, main: Path, metadata: EntityMetadata,
entity: PyKokkosEntity,
classtypes: List[PyKokkosEntity],
spaces: List[ExecutionSpace], force_uvm: bool) -> None:
"""
Compile the entity
:param main: the path to the main file in the current PyKokkos application
:param metadata: the metadata of the entity being compiled
:param entity: the parsed entity being compiled
:param classtypes: the list of parsed classtypes being compiled
:param spaces: the list of execution spaces to compile for
:param force_uvm: whether CudaUVMSpace is enabled
:returns: the PyKokkos members obtained during translation
"""
cpp_setup = CppSetup(self.module_file, self.functor_file,
self.bindings_file)
translator = StaticTranslator(self.module_name, self.functor_file)
spaces = [
km.get_default_space() if s is ExecutionSpace.Default else s
for s in spaces
]
is_compiled: List[bool] = [
CppSetup.is_compiled(self.get_output_dir(main, metadata, s))
for s in spaces
]
if not all(is_compiled):
t_start: float = time.perf_counter()
functor: List[str]
bindings: List[str]
functor, bindings = translator.translate(entity, classtypes)
t_end: float = time.perf_counter() - t_start
self.logger.info(f"translation {t_end}")
for i, s in enumerate(spaces):
if s is ExecutionSpace.Debug:
continue
if not is_compiled[i]:
output_dir: Path = self.get_output_dir(main, metadata, s)
c_start: float = time.perf_counter()
cpp_setup.compile(output_dir, functor, bindings, s,
force_uvm, self.get_compiler())
c_end: float = time.perf_counter() - c_start
self.logger.info(f"compilation {c_end}")
def _init_view(
self,
shape: List[int],
dtype: Union[DataTypeClass, type] = real,
space: MemorySpace = MemorySpace.MemorySpaceDefault,
layout: Layout = Layout.LayoutDefault,
trait: Trait = Trait.TraitDefault,
array: Optional[np.ndarray] = None
) -> None:
"""
Initialize the view
:param shape: the shape of the view as a list of integers
:param dtype: the data type of the view, either a pykokkos DataType or "int" or "float".
:param space: the memory space of the view. Will be set to the execution space of the view by default.
:param layout: the layout of the view in memory.
:param trait: the memory trait of the view
:param array: the numpy array if trait is Unmanaged
"""
self.shape: List[int] = shape
self.dtype: Optional[DataType] = self._get_type(dtype)
if self.dtype is None:
sys.exit(f"ERROR: Invalid dtype {dtype}")
if space is MemorySpace.MemorySpaceDefault:
space = get_default_memory_space(km.get_default_space())
if layout is Layout.LayoutDefault:
layout = get_default_layout(space)
if space is MemorySpace.CudaSpace:
space = MemorySpace.HostSpace
self.space: MemorySpace = space
self.layout: Layout = layout
self.trait: Trait = trait
if trait is trait.Unmanaged:
self.array = kokkos.unmanaged_array(array, self.dtype.value, self.space.value)
else:
self.array = kokkos.array("", shape, None, None, self.dtype.value, space.value, layout.value, trait.value)
self.data = np.array(self.array, copy=False)
def run_workunit_debug(policy: ExecutionPolicy,
workunit: Callable[..., None],
operation: str,
initial_value=0,
**kwargs) -> Optional[Union[float, int]]:
"""
Run a workunit in Python
:param operation: the name of the operation "for", "reduce", or "scan"
:param policy: the execution policy of the operation
:param workunit: the workunit function object
:param initial_value: the initial value of the accumulator
:param kwargs: the keyword arguments passed to the workunit
:returns: the result of the operation (None for parallel_for)
"""
acc = Acc(initial_value)
if policy.space is ExecutionSpace.Default:
policy.space = km.get_default_space()
if isinstance(policy, TeamPolicy):
for i in range(policy.league_size):
call_workunit(operation, workunit, TeamMember(i, 0), acc, **kwargs)
elif isinstance(policy, TeamThreadRange) or isinstance(
policy, ThreadVectorRange):
for i in range(policy.count):
call_workunit(operation, workunit, TeamMember(i, 0), acc, **kwargs)
else:
if isinstance(policy, MDRangePolicy):
if policy.rank > 1:
for idx in itertools.product(*[
range(*interval)
for interval in zip(policy.begin, policy.end)
]):
call_workunit(operation, workunit, idx, acc, **kwargs)
else:
for i in range(policy.begin, policy.end):
call_workunit(operation, workunit, i, acc, **kwargs)
return acc.val
def execute(space: ExecutionSpace, workload: object) -> None:
if space is ExecutionSpace.Default:
runtime_singleton.runtime.run_workload(km.get_default_space(), workload)
else:
runtime_singleton.runtime.run_workload(space, workload)