diff --git a/legate/core/__init__.py b/legate/core/__init__.py index cc3bc4027f..4ad4c308b3 100644 --- a/legate/core/__init__.py +++ b/legate/core/__init__.py @@ -68,6 +68,7 @@ # Import select types for Legate library construction from .allocation import DistributedAllocation +from .context import track_provenance from .legate import ( Array, Library, diff --git a/legate/core/_legion/operation.py b/legate/core/_legion/operation.py index 9a52843796..2ccde2f5cd 100644 --- a/legate/core/_legion/operation.py +++ b/legate/core/_legion/operation.py @@ -133,6 +133,7 @@ def __init__( future: Future, mapper: int = 0, tag: int = 0, + provenance: Optional[str] = None, ) -> None: """ A Fill object provides a mechanism for launching fill operations @@ -169,6 +170,10 @@ def __init__( mapper, tag, ) + if provenance is not None: + legion.legion_fill_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_fill_launcher_destroy ) @@ -225,6 +230,7 @@ def __init__( mapper: int = 0, tag: int = 0, space: Optional[Union[IndexSpace, Domain]] = None, + provenance: Optional[str] = None, ) -> None: """ An IndexFill object provides a mechanism for launching index space fill @@ -308,6 +314,10 @@ def __init__( mapper, tag, ) + if provenance is not None: + legion.legion_index_fill_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_index_fill_launcher_destroy ) @@ -342,7 +352,12 @@ def launch( class Copy(Dispatchable[None]): - def __init__(self, mapper: int = 0, tag: int = 0) -> None: + def __init__( + self, + mapper: int = 0, + tag: int = 0, + provenance: Optional[str] = None, + ) -> None: """ A Copy object provides a mechanism for launching explicit region-to-region copy operations. Note: you should NOT use @@ -359,6 +374,10 @@ def __init__(self, mapper: int = 0, tag: int = 0) -> None: self.launcher = legion.legion_copy_launcher_create( legion.legion_predicate_true(), mapper, tag ) + if provenance is not None: + legion.legion_copy_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_copy_launcher_destroy ) @@ -643,7 +662,13 @@ def launch( class IndexCopy(Dispatchable[None]): - def __init__(self, domain: Rect, mapper: int = 0, tag: int = 0) -> None: + def __init__( + self, + domain: Rect, + mapper: int = 0, + tag: int = 0, + provenance: Optional[str] = None, + ) -> None: """ An IndexCopy object provides a mechanism for launching explicit region-to-region copies between many different subregions @@ -663,6 +688,10 @@ def __init__(self, domain: Rect, mapper: int = 0, tag: int = 0) -> None: self.launcher = legion.legion_index_copy_launcher_create( domain.raw(), legion.legion_predicate_true(), mapper, tag ) + if provenance is not None: + legion.legion_index_copy_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_index_copy_launcher_destroy ) diff --git a/legate/core/_legion/task.py b/legate/core/_legion/task.py index 4bc50f18e5..7e463abec6 100644 --- a/legate/core/_legion/task.py +++ b/legate/core/_legion/task.py @@ -36,6 +36,7 @@ def __init__( size: int = 0, mapper: int = 0, tag: int = 0, + provenance: Optional[str] = None, ) -> None: """ A Task object provides a mechanism for launching individual sub-tasks @@ -78,6 +79,10 @@ def __init__( mapper, tag, ) + if provenance is not None: + legion.legion_task_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_task_launcher_destroy ) @@ -478,6 +483,7 @@ def __init__( size: int = 0, mapper: int = 0, tag: int = 0, + provenance: Optional[str] = None, ) -> None: """ An IndexTask object provides a mechnanism for launching a collection @@ -541,6 +547,10 @@ def __init__( mapper, tag, ) + if provenance is not None: + legion.legion_index_launcher_set_provenance( + self.launcher, provenance.encode() + ) self._launcher = ffi.gc( self.launcher, legion.legion_index_launcher_destroy ) diff --git a/legate/core/context.py b/legate/core/context.py index fe3aa37001..d67796da01 100644 --- a/legate/core/context.py +++ b/legate/core/context.py @@ -14,7 +14,17 @@ # from __future__ import annotations -from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union, cast +import traceback +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Optional, + Protocol, + TypeVar, + Union, + cast, +) import numpy as np @@ -38,6 +48,24 @@ T = TypeVar("T") +class AnyCallable(Protocol): + def __call__(self, *args: Any, **kwargs: Any) -> Any: + ... + + +def find_last_user_frame(libname: str) -> str: + for (frame, _) in traceback.walk_stack(None): + if "__name__" not in frame.f_globals: + continue + if not any( + frame.f_globals["__name__"].startswith(prefix) + for prefix in (libname, "legate") + ): + break + + return f"{frame.f_code.co_filename}:{frame.f_lineno}" + + class Context: def __init__( self, @@ -96,7 +124,9 @@ def _create_scope( config.max_shardings, ) + self._libname = library.get_name() self._unique_op_id = 0 + self._provenance: list[Union[str, None]] = [None] def destroy(self) -> None: self._library.destroy() @@ -133,6 +163,10 @@ def empty_argmap(self) -> ArgumentMap: def type_system(self) -> TypeSystem: return self._type_system + @property + def provenance(self) -> Optional[str]: + return self._provenance[-1] + def get_task_id(self, task_id: int) -> int: return self._task_scope.translate(task_id) @@ -175,6 +209,44 @@ def get_tunable( def get_unique_op_id(self) -> int: return self._runtime.get_unique_op_id() + def set_provenance(self, provenance: str) -> None: + self._provenance[-1] = provenance + + def reset_provenance(self) -> None: + self._provenance[-1] = None + + def push_provenance(self, provenance: str) -> None: + self._provenance.append(provenance) + + def pop_provenance(self) -> None: + if len(self._provenance) == 1: + raise ValueError("Provenance stack underflow") + self._provenance.pop(-1) + + def track_provenance( + self, func: AnyCallable, nested: bool = False + ) -> AnyCallable: + if nested: + + def wrapper(*args: Any, **kwargs: Any) -> Any: + self.push_provenance(find_last_user_frame(self._libname)) + result = func(*args, **kwargs) + self.pop_provenance() + return result + + else: + + def wrapper(*args: Any, **kwargs: Any) -> Any: + if self.provenance is None: + self.set_provenance(find_last_user_frame(self._libname)) + result = func(*args, **kwargs) + self.reset_provenance() + else: + result = func(*args, **kwargs) + return result + + return wrapper + def create_task( self, task_id: int, @@ -292,3 +364,13 @@ def tree_reduce( task.add_output(result) task.execute() return result + + +def track_provenance( + context: Context, + nested: bool = False, +) -> Callable[[AnyCallable], AnyCallable]: + def decorator(func: AnyCallable) -> AnyCallable: + return context.track_provenance(func, nested=nested) + + return decorator diff --git a/legate/core/launcher.py b/legate/core/launcher.py index b2a6748a59..e36a8e2b94 100644 --- a/legate/core/launcher.py +++ b/legate/core/launcher.py @@ -692,6 +692,7 @@ def __init__( tag: int = 0, error_on_interference: bool = True, side_effect: bool = False, + provenance: Optional[str] = None, ) -> None: assert type(tag) != bool self._context = context @@ -715,6 +716,7 @@ def __init__( self._has_side_effect = side_effect self._insert_barrier = False self._can_raise_exception = False + self._provenance = provenance @property def library_task_id(self) -> int: @@ -903,6 +905,7 @@ def build_task( argbuf.get_size(), mapper=self.legion_mapper_id, tag=self._tag, + provenance=self._provenance, ) if self._sharding_space is not None: task.set_sharding_space(self._sharding_space) @@ -942,6 +945,7 @@ def build_single_task(self, argbuf: BufferBuilder) -> SingleTask: argbuf.get_size(), mapper=self.legion_mapper_id, tag=self._tag, + provenance=self._provenance, ) for (req, fields) in self._req_analyzer.requirements: req.proj.add_single(task, req, fields, _single_task_calls) @@ -983,6 +987,7 @@ def __init__( target_oor: bool = True, mapper_id: int = 0, tag: int = 0, + provenance: Optional[str] = None, ) -> None: assert type(tag) != bool self._context = context @@ -993,6 +998,7 @@ def __init__( self._point: Union[Point, None] = None self._source_oor = source_oor self._target_oor = target_oor + self._provenance = provenance @property def library_mapper_id(self) -> int: @@ -1064,6 +1070,7 @@ def build_copy(self, launch_domain: Rect) -> IndexCopy: launch_domain, mapper=self.legion_mapper_id, tag=self._tag, + provenance=self._provenance, ) for (req, fields) in self._req_analyzer.requirements: if req.permission in ( @@ -1085,6 +1092,7 @@ def build_single_copy(self) -> SingleCopy: copy = SingleCopy( mapper=self.legion_mapper_id, tag=self._tag, + provenance=self._provenance, ) for (req, fields) in self._req_analyzer.requirements: if req.permission in ( diff --git a/legate/core/operation.py b/legate/core/operation.py index e27df60afa..407784ac66 100644 --- a/legate/core/operation.py +++ b/legate/core/operation.py @@ -134,6 +134,10 @@ def __init__( self._launch_domain: Union[Rect, None] = None self._error_on_interference = True + @property + def provenance(self) -> Optional[str]: + return self._context.provenance + def get_all_stores(self) -> OrderedSet[Store]: result: OrderedSet[Store] = OrderedSet() result.update(self._inputs) @@ -483,7 +487,12 @@ def __init__( ) def launch(self, strategy: Strategy) -> None: - launcher = TaskLauncher(self.context, self._task_id, self.mapper_id) + launcher = TaskLauncher( + self.context, + self._task_id, + self.mapper_id, + provenance=self.provenance, + ) def get_requirement( store: Store, part_symb: PartSym @@ -557,9 +566,14 @@ def __init__( launch_domain: Rect, mapper_id: int = 0, op_id: int = 0, + **kwargs: Any, ) -> None: super().__init__( - context=context, task_id=task_id, mapper_id=mapper_id, op_id=op_id + context=context, + task_id=task_id, + mapper_id=mapper_id, + op_id=op_id, + **kwargs, ) self._launch_domain: Rect = launch_domain self._input_projs: list[Union[ProjFn, None]] = [] @@ -658,6 +672,7 @@ def launch(self, strategy: Strategy) -> None: self.mapper_id, error_on_interference=False, tag=tag, + provenance=self.provenance, ) for part, proj_fn in zip(self._input_parts, self._input_projs): @@ -699,8 +714,13 @@ def launch(self, strategy: Strategy) -> None: class Copy(AutoOperation): - def __init__(self, context: Context, mapper_id: int = 0) -> None: - super().__init__(context=context, mapper_id=mapper_id) + def __init__( + self, + context: Context, + mapper_id: int = 0, + **kwargs: Any, + ) -> None: + super().__init__(context=context, mapper_id=mapper_id, **kwargs) self._source_indirects: list[Store] = [] self._target_indirects: list[Store] = [] self._source_indirect_parts: list[PartSym] = [] @@ -823,6 +843,7 @@ def launch(self, strategy: Strategy) -> None: source_oor=self._source_indirect_out_of_range, target_oor=self._target_indirect_out_of_range, mapper_id=self.mapper_id, + provenance=self.provenance, ) assert len(self._inputs) == len(self._outputs) or len( @@ -903,8 +924,14 @@ def __init__( radix: int, mapper_id: int, op_id: int, + **kwargs: Any, ) -> None: - super().__init__(context=context, mapper_id=mapper_id, op_id=op_id) + super().__init__( + context=context, + mapper_id=mapper_id, + op_id=op_id, + **kwargs, + ) self._runtime = context.runtime self._radix = radix self._task_id = task_id @@ -935,7 +962,11 @@ def launch(self, strategy: Strategy) -> None: tag = self.context.core_library.LEGATE_CORE_TREE_REDUCE_TAG launcher = TaskLauncher( - self.context, self._task_id, self.mapper_id, tag=tag + self.context, + self._task_id, + self.mapper_id, + tag=tag, + provenance=self.provenance, ) for proj_fn in proj_fns: diff --git a/typings/legion_cffi/lib.pyi b/typings/legion_cffi/lib.pyi index 2bea06a0b2..bce4097e0b 100644 --- a/typings/legion_cffi/lib.pyi +++ b/typings/legion_cffi/lib.pyi @@ -122,6 +122,7 @@ def legion_copy_launcher_set_possible_src_indirect_out_of_range( *args: Any, ) -> Any: ... def legion_copy_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_copy_launcher_set_provenance(*args: Any) -> Any: ... def legion_detach_external_resources(*args: Any) -> Any: ... def legion_domain_affine_transform_identity(*args: Any) -> Any: ... def legion_domain_empty(*args: Any) -> Any: ... @@ -143,6 +144,7 @@ def legion_fill_launcher_destroy(*args: Any) -> Any: ... def legion_fill_launcher_execute(*args: Any) -> Any: ... def legion_fill_launcher_set_point(*args: Any) -> Any: ... def legion_fill_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_fill_launcher_set_provenance(*args: Any) -> Any: ... def legion_future_destroy(*args: Any) -> Any: ... def legion_future_from_untyped_pointer(*args: Any) -> Any: ... def legion_future_get_untyped_pointer(*args: Any) -> Any: ... @@ -204,6 +206,7 @@ def legion_index_copy_launcher_set_possible_src_indirect_out_of_range( *args: Any, ) -> Any: ... def legion_index_copy_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_index_copy_launcher_set_provenance(*args: Any) -> Any: ... def legion_index_fill_launcher_create_from_future_with_domain( *args: Any, ) -> Any: ... @@ -213,6 +216,7 @@ def legion_index_fill_launcher_create_from_future_with_space( def legion_index_fill_launcher_destroy(*args: Any) -> Any: ... def legion_index_fill_launcher_execute(*args: Any) -> Any: ... def legion_index_fill_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_index_fill_launcher_set_provenance(*args: Any) -> Any: ... def legion_index_launcher_add_field(*args: Any) -> Any: ... def legion_index_launcher_add_flags(*args: Any) -> Any: ... def legion_index_launcher_add_future(*args: Any) -> Any: ... @@ -238,6 +242,7 @@ def legion_index_launcher_execute_deterministic_reduction( def legion_index_launcher_execute_outputs(*args: Any) -> Any: ... def legion_index_launcher_execute_reduction_and_outputs(*args: Any) -> Any: ... def legion_index_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_index_launcher_set_provenance(*args: Any) -> Any: ... def legion_index_partition_create_by_domain(*args: Any) -> Any: ... def legion_index_partition_create_by_domain_future_map(*args: Any) -> Any: ... def legion_index_partition_create_by_image(*args: Any) -> Any: ... @@ -322,6 +327,7 @@ def legion_task_launcher_execute_outputs(*args: Any) -> Any: ... def legion_task_launcher_set_local_function_task(*args: Any) -> Any: ... def legion_task_launcher_set_point(*args: Any) -> Any: ... def legion_task_launcher_set_sharding_space(*args: Any) -> Any: ... +def legion_task_launcher_set_provenance(*args: Any) -> Any: ... def legion_unordered_detach_external_resource(*args: Any) -> Any: ... def legion_runtime_get_runtime(*args: Any) -> Any: ... def legion_runtime_generate_library_task_ids(*args: Any) -> Any: ... @@ -407,6 +413,7 @@ __all__ = ( "legion_copy_launcher_set_possible_dst_indirect_out_of_range", "legion_copy_launcher_set_possible_src_indirect_out_of_range", "legion_copy_launcher_set_sharding_space", + "legion_copy_launcher_set_provenance", "legion_detach_external_resources", "legion_domain_affine_transform_identity", "legion_domain_empty", @@ -428,6 +435,7 @@ __all__ = ( "legion_fill_launcher_execute", "legion_fill_launcher_set_point", "legion_fill_launcher_set_sharding_space", + "legion_fill_launcher_set_provenance", "legion_future_destroy", "legion_future_from_untyped_pointer", "legion_future_get_untyped_pointer", @@ -463,11 +471,13 @@ __all__ = ( "legion_index_copy_launcher_set_possible_dst_indirect_out_of_range", "legion_index_copy_launcher_set_possible_src_indirect_out_of_range", "legion_index_copy_launcher_set_sharding_space", + "legion_index_copy_launcher_set_provenance", "legion_index_fill_launcher_create_from_future_with_domain", "legion_index_fill_launcher_create_from_future_with_space", "legion_index_fill_launcher_destroy", "legion_index_fill_launcher_execute", "legion_index_fill_launcher_set_sharding_space", + "legion_index_fill_launcher_set_provenance", "legion_index_launcher_add_field", "legion_index_launcher_add_flags", "legion_index_launcher_add_future", @@ -483,6 +493,7 @@ __all__ = ( "legion_index_launcher_execute_outputs", "legion_index_launcher_execute_reduction_and_outputs", "legion_index_launcher_set_sharding_space", + "legion_index_launcher_set_provenance", "legion_index_partition_create_by_domain", "legion_index_partition_create_by_domain_future_map", "legion_index_partition_create_by_image", @@ -552,6 +563,7 @@ __all__ = ( "legion_task_launcher_set_local_function_task", "legion_task_launcher_set_point", "legion_task_launcher_set_sharding_space", + "legion_task_launcher_set_provenance", "legion_unordered_detach_external_resource", "legion_runtime_get_runtime", "legion_runtime_generate_library_task_ids",