23284927681

Committed 19 Mar 2026 07:46AM UTC coverage: 63.839%. Remained the same

Build # 23284927681

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #595 from daisytuner/expand-override-hook

python target override for expand phase & combined sdfg pipeline code…

Coverage Stats

48 of 63 new or added lines in 4 files covered. (76.19%)

26096 of 40878 relevant lines covered (63.84%)

401.57 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

91.18

/python/docc/compiler/target_registry.py

from typing import Callable, Optional, Dict, Any
from docc.sdfg import StructuredSDFG

TargetScheduleFn = Callable[[StructuredSDFG, str, Dict[str, Any]], None]
LegacyTargetScheduleFn = Callable[[StructuredSDFG, str], None]
TargetCompileFn = Callable[[StructuredSDFG, str, str, bool, Dict[str, Any]], str]
TargetExpandFn = Callable[[StructuredSDFG, str, Dict[str, Any]], None]

_target_expand_registry: dict[str, TargetExpandFn] = {}
_target_schedule_registry: dict[str, TargetScheduleFn] = {}
_target_compile_registry: dict[str, TargetCompileFn] = {}


def register_target(name: str, schedule_fn: Optional[LegacyTargetScheduleFn]) -> None:
    def wrapper(
        sdfg: StructuredSDFG, category: str, options: Optional[Dict[str, Any]] = None
    ) -> None:
        schedule_fn(sdfg, category)  # type: ignore

    register_target_overrides(name, schedule_fn=wrapper, compile_fn=None)


def register_target_overrides(
    name: str,
    schedule_fn: Optional[TargetScheduleFn],
    compile_fn: Optional[TargetCompileFn] = None,
    expand_fn: Optional[TargetExpandFn] = None,
) -> None:
    """Override the scheduling or compile step for this target.

    The schedule function will be called with:
    - sdfg: The StructuredSDFG to schedule (has _native_ptr for native access)
    - category: The target category (e.g., "desktop", "server")
    - options: (Optional) Dictionary of options

    Args:
        name: Target name (e.g., "openmp")
        schedule_fn: Function that performs scheduling transformations
        compile_fn: Function that performs compile step
        expand_fn: Function that performs the expand step
    """
    if schedule_fn is not None:
        _target_schedule_registry[name] = schedule_fn  # type: ignore

    if compile_fn is not None:
        _target_compile_registry[name] = compile_fn

    if expand_fn is not None:
        _target_expand_registry[name] = expand_fn


def unregister_target(name: str) -> None:
    """Unregister a custom target scheduler."""
    _target_schedule_registry.pop(name, None)
    _target_compile_registry.pop(name, None)
    _target_expand_registry.pop(name, None)


def get_target_schedule_fn(name: str) -> Optional[TargetScheduleFn]:
    """Get a registered target scheduler, or None if not found."""
    return _target_schedule_registry.get(name)


def get_target_compile_fn(name: str) -> Optional[TargetCompileFn]:
    """Get a registered target compile function, or None if not found."""
    return _target_compile_registry.get(name)


def get_target_expand_fn(name: str) -> Optional[TargetExpandFn]:
    """Get a registered target expand function, or None if not found."""
    return _target_expand_registry.get(name)


def reset_target_registry() -> None:
    """Clear all global state of the registry"""
    _target_expand_registry.clear()
    _target_schedule_registry.clear()
    _target_compile_registry.clear()

1	from typing import Callable, Optional, Dict, Any	4✔
2	from docc.sdfg import StructuredSDFG	4✔
3
4	TargetScheduleFn = Callable[[StructuredSDFG, str, Dict[str, Any]], None]	4✔
5	LegacyTargetScheduleFn = Callable[[StructuredSDFG, str], None]	4✔
6	TargetCompileFn = Callable[[StructuredSDFG, str, str, bool, Dict[str, Any]], str]	4✔
7	TargetExpandFn = Callable[[StructuredSDFG, str, Dict[str, Any]], None]	4✔
8
9	_target_expand_registry: dict[str, TargetExpandFn] = {}	4✔
10	_target_schedule_registry: dict[str, TargetScheduleFn] = {}	4✔
11	_target_compile_registry: dict[str, TargetCompileFn] = {}	4✔
12
13
14	def register_target(name: str, schedule_fn: Optional[LegacyTargetScheduleFn]) -> None:	4✔
15	def wrapper(	4✔
16	sdfg: StructuredSDFG, category: str, options: Optional[Dict[str, Any]] = None
17	) -> None:
18	schedule_fn(sdfg, category) # type: ignore	4✔
19
20	register_target_overrides(name, schedule_fn=wrapper, compile_fn=None)	4✔
21
22
23	def register_target_overrides(	4✔
24	name: str,
25	schedule_fn: Optional[TargetScheduleFn],
26	compile_fn: Optional[TargetCompileFn] = None,
27	expand_fn: Optional[TargetExpandFn] = None,
28	) -> None:
29	"""Override the scheduling or compile step for this target.
30
31	The schedule function will be called with:
32	- sdfg: The StructuredSDFG to schedule (has _native_ptr for native access)
33	- category: The target category (e.g., "desktop", "server")
34	- options: (Optional) Dictionary of options
35
36	Args:
37	name: Target name (e.g., "openmp")
38	schedule_fn: Function that performs scheduling transformations
39	compile_fn: Function that performs compile step
40	expand_fn: Function that performs the expand step
41	"""
42	if schedule_fn is not None:	4✔
43	_target_schedule_registry[name] = schedule_fn # type: ignore	4✔
44
45	if compile_fn is not None:	4✔
46	_target_compile_registry[name] = compile_fn	4✔
47
48	if expand_fn is not None:	4✔
49	_target_expand_registry[name] = expand_fn	4✔
50
51
52	def unregister_target(name: str) -> None:	4✔
53	"""Unregister a custom target scheduler."""
54	_target_schedule_registry.pop(name, None)	×
55	_target_compile_registry.pop(name, None)	×
NEW 56	_target_expand_registry.pop(name, None)	×
57
58
59	def get_target_schedule_fn(name: str) -> Optional[TargetScheduleFn]:	4✔
60	"""Get a registered target scheduler, or None if not found."""
61	return _target_schedule_registry.get(name)	4✔
62
63
64	def get_target_compile_fn(name: str) -> Optional[TargetCompileFn]:	4✔
65	"""Get a registered target compile function, or None if not found."""
66	return _target_compile_registry.get(name)	4✔
67
68
69	def get_target_expand_fn(name: str) -> Optional[TargetExpandFn]:	4✔
70	"""Get a registered target expand function, or None if not found."""
71	return _target_expand_registry.get(name)	4✔
72
73
74	def reset_target_registry() -> None:	4✔
75	"""Clear all global state of the registry"""
76	_target_expand_registry.clear()	4✔
77	_target_schedule_registry.clear()	4✔
78	_target_compile_registry.clear()	4✔

daisytuner / docc / 23284927681

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous