27885489129

Committed 20 Jun 2026 10:13PM UTC coverage: 61.76% (-0.07%) from 61.832%

Build # 27885489129

Build Type

Pull #789

github

Committed by

web-flow

Commit Message

Merge 86a899853 into b7103c21a

Pull Request Pull Request #789: Device Residency

Coverage Stats

81 of 204 new or added lines in 6 files covered. (39.71%)

2 existing lines in 2 files now uncovered.

37023 of 59947 relevant lines covered (61.76%)

1015.71 hits per line

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

69.44

/python/docc/compiler/docc_program.py

from abc import ABC, abstractmethod
import sys
from typing import Any, Dict, Optional
import os
import re

from docc.sdfg import StructuredSDFG, TargetOptions
from docc.sdfg._sdfg import (
    _enable_statistics,
    _statistics_enabled_by_env,
    _statistics_summary,
)
from docc.compiler.compiled_sdfg import CompiledSDFG
from docc.compiler.target_registry import (
    get_target_schedule_fn,
    get_target_compile_fn,
    get_target_expand_fn,
    register_target_overrides,
)


def _parse_docc_debug() -> dict[str, str]:
    debug_env = os.environ.get("DOCC_DEBUG", "")
    debug_dict = {}
    if debug_env:
        for entry in re.split(r"[;:]", debug_env):
            if not entry:
                continue
            parts = entry.split("=", 1)
            key = parts[0].strip()
            value = parts[1].strip() if len(parts) > 1 else ""
            debug_dict[key] = value
    return debug_dict


def _is_debug_dump(flags: dict[str, str]) -> bool:
    return "dump" in flags


def _is_debug_compile(flags: dict[str, str]) -> bool:
    return "build" in flags


def _get_build_thread_count(flags: dict[str, str]) -> int:
    return int(flags.get("build_threads", "0"))


class DoccProgram(ABC):

    def __init__(
        self,
        name: str,
        target: str = "none",
        category: str = "server",
        instrumentation_mode: Optional[str] = None,
        capture_args: Optional[bool] = None,
        remote_tuning: bool = False,
    ):
        self.name = name
        self.target = target
        self.category = category
        self.remote_tuning = remote_tuning
        self.last_sdfg: Optional[StructuredSDFG] = None
        self._device_resident: bool = False
        self._device_backend: Optional[str] = None
        self.cache: dict = {}
        debug_flags = _parse_docc_debug()
        self.debug_dump: bool = _is_debug_dump(debug_flags)
        self.debug_build: bool = _is_debug_compile(debug_flags)
        self.build_thread_count: int = _get_build_thread_count(debug_flags)

        # Check environment variable DOCC_CI
        docc_ci = os.environ.get("DOCC_CI", "")
        if docc_ci:
            if docc_ci == "regions":
                if instrumentation_mode is None:
                    instrumentation_mode = "ols"
            elif docc_ci == "arg-capture":
                if capture_args is None:
                    capture_args = True
            else:
                # Full mode (or unknown value treated as full)
                if instrumentation_mode is None:
                    instrumentation_mode = "ols"
                if capture_args is None:
                    capture_args = True

        self.instrumentation_mode = instrumentation_mode
        self.capture_args = capture_args

    @abstractmethod
    def __call__(self, *args: Any) -> Any:
        pass

    @abstractmethod
    def compile(self, *args: Any, output_folder: Optional[str] = None) -> CompiledSDFG:
        pass

    def _resolve_compile_options(
        self,
        instrumentation_mode: Optional[str] = None,
        capture_args: Optional[bool] = None,
        remote_tuning: Optional[bool] = None,
    ) -> tuple[str, bool, bool]:
        """Resolve compile-time options, falling back to instance defaults and env vars."""
        if instrumentation_mode is None:
            instrumentation_mode = self.instrumentation_mode
        if capture_args is None:
            capture_args = self.capture_args
        if remote_tuning is None:
            remote_tuning = self.remote_tuning

        # Check environment variable DOCC_CI
        docc_ci = os.environ.get("DOCC_CI", "")
        if docc_ci:
            if docc_ci == "regions":
                if instrumentation_mode is None:
                    instrumentation_mode = "ols"
            elif docc_ci == "arg-capture":
                if capture_args is None:
                    capture_args = True
            else:
                # Full mode (or unknown value treated as full)
                if instrumentation_mode is None:
                    instrumentation_mode = "ols"
                if capture_args is None:
                    capture_args = True

        # Defaults
        if instrumentation_mode is None:
            instrumentation_mode = ""
        if capture_args is None:
            capture_args = False

        return instrumentation_mode, capture_args, remote_tuning

    def sdfg_pipe(
        self,
        sdfg: StructuredSDFG,
        output_folder: Optional[str],
        instrumentation_mode: str,
        capture_args: bool,
        remote_tuning: Optional[bool] = None,
    ) -> str:

        if self.debug_dump:
            sdfg.dump(output_folder, "py0.parsed", dump_dot=True)

        # Enable statistics if envvar is set
        if _statistics_enabled_by_env():
            _enable_statistics()

        sdfg.validate()

        if remote_tuning is None:
            remote_tuning = self.remote_tuning

        target_options = TargetOptions()
        target_options.target = self.target
        target_options.category = self.category
        target_options.remote_tuning = remote_tuning

        # Einsum detection
        sdfg.einsum()
        if self.debug_dump:
            sdfg.dump(output_folder, "py1.einsum", dump_dot=True)

        # Tensor targets keep tensor nodes
        custom_expand_fn = get_target_expand_fn(self.target)
        if custom_expand_fn is not None:
            custom_expand_fn(sdfg, self.category, {})
        else:
            sdfg.expand(target_options)
        if self.debug_dump:
            sdfg.dump(output_folder, "py2.expanded", dump_dot=True)

        # Simplify pipelines
        sdfg.simplify()
        if self.debug_dump:
            sdfg.dump(output_folder, "py3.opt", dump_dot=True)

        # Normalization for scheduling
        if self.target != "none":
            sdfg.normalize()

        if self.debug_dump or instrumentation_mode or capture_args:
            sdfg.dump(
                output_folder,
                "py4.norm",
                dump_dot=self.debug_dump,
                dump_json=True,
                record_for_instrumentation=True,
            )

        # Schedule if target is specified

        custom_schedule_fn = get_target_schedule_fn(self.target)
        if custom_schedule_fn is not None:
            custom_schedule_fn(sdfg, self.category, {"remote_tuning": remote_tuning})
        else:
            sdfg.schedule(target_options)

        if self.debug_dump:
            sdfg.dump(output_folder, "py5.post_sched", dump_dot=True)

        # Promote pointer arguments to device residency when the whole program keeps
        # data on device. Communicated explicitly via the pass return value (bool),
        # not through SDFG metadata.
        self._device_resident = False
        self._device_backend = None
        if self.target in ("cuda", "rocm"):
            if sdfg.promote_device_residency(self.target == "rocm"):
                self._device_resident = True
                self._device_backend = self.target

        self.last_sdfg = sdfg

        custom_compile_fn = get_target_compile_fn(self.target)
        if custom_compile_fn is not None:
            lib_path = custom_compile_fn(
                sdfg,
                output_folder,
                instrumentation_mode,
                capture_args,
                {"debug_build": self.debug_build, "threads": self.build_thread_count},
            )
        else:
            lib_path = sdfg._compile(
                output_folder=output_folder,
                target=self.target,
                instrumentation_mode=instrumentation_mode,
                capture_args=capture_args,
                debug_build=self.debug_build,
                threads=self.build_thread_count,
            )

        # Dump statistics after compile
        if _statistics_enabled_by_env():
            print(_statistics_summary(), file=sys.stderr)

        return lib_path

    @abstractmethod
    def to_sdfg(self, *args: Any) -> StructuredSDFG:
        pass

    @abstractmethod
    def _convert_inputs(self, args: tuple) -> tuple:
        pass

    @abstractmethod
    def _convert_outputs(self, result: Any, original_args: tuple) -> Any:
        pass

    def _get_cache_key(self, *args: Any) -> str:
        return ""

1	from abc import ABC, abstractmethod	4✔
2	import sys	4✔
3	from typing import Any, Dict, Optional	4✔
4	import os	4✔
5	import re	4✔
6
7	from docc.sdfg import StructuredSDFG, TargetOptions	4✔
8	from docc.sdfg._sdfg import (	4✔
9	_enable_statistics,
10	_statistics_enabled_by_env,
11	_statistics_summary,
12	)
13	from docc.compiler.compiled_sdfg import CompiledSDFG	4✔
14	from docc.compiler.target_registry import (	4✔
15	get_target_schedule_fn,
16	get_target_compile_fn,
17	get_target_expand_fn,
18	register_target_overrides,
19	)
20
21
22	def _parse_docc_debug() -> dict[str, str]:	4✔
23	debug_env = os.environ.get("DOCC_DEBUG", "")	4✔
24	debug_dict = {}	4✔
25	if debug_env:	4✔
26	for entry in re.split(r"[;:]", debug_env):	×
27	if not entry:	×
28	continue	×
29	parts = entry.split("=", 1)	×
30	key = parts[0].strip()	×
31	value = parts[1].strip() if len(parts) > 1 else ""	×
32	debug_dict[key] = value	×
33	return debug_dict	4✔
34
35
36	def _is_debug_dump(flags: dict[str, str]) -> bool:	4✔
37	return "dump" in flags	4✔
38
39
40	def _is_debug_compile(flags: dict[str, str]) -> bool:	4✔
41	return "build" in flags	4✔
42
43
44	def _get_build_thread_count(flags: dict[str, str]) -> int:	4✔
45	return int(flags.get("build_threads", "0"))	4✔
46
47
48	class DoccProgram(ABC):	4✔
49
50	def __init__(	4✔
51	self,
52	name: str,
53	target: str = "none",
54	category: str = "server",
55	instrumentation_mode: Optional[str] = None,
56	capture_args: Optional[bool] = None,
57	remote_tuning: bool = False,
58	):
59	self.name = name	4✔
60	self.target = target	4✔
61	self.category = category	4✔
62	self.remote_tuning = remote_tuning	4✔
63	self.last_sdfg: Optional[StructuredSDFG] = None	4✔
64	self._device_resident: bool = False	4✔
65	self._device_backend: Optional[str] = None	4✔
66	self.cache: dict = {}	4✔
67	debug_flags = _parse_docc_debug()	4✔
68	self.debug_dump: bool = _is_debug_dump(debug_flags)	4✔
69	self.debug_build: bool = _is_debug_compile(debug_flags)	4✔
70	self.build_thread_count: int = _get_build_thread_count(debug_flags)	4✔
71
72	# Check environment variable DOCC_CI
73	docc_ci = os.environ.get("DOCC_CI", "")	4✔
74	if docc_ci:	4✔
75	if docc_ci == "regions":	×
76	if instrumentation_mode is None:	×
77	instrumentation_mode = "ols"	×
78	elif docc_ci == "arg-capture":	×
79	if capture_args is None:	×
80	capture_args = True	×
81	else:
82	# Full mode (or unknown value treated as full)
83	if instrumentation_mode is None:	×
84	instrumentation_mode = "ols"	×
85	if capture_args is None:	×
86	capture_args = True	×
87
88	self.instrumentation_mode = instrumentation_mode	4✔
89	self.capture_args = capture_args	4✔
90
91	@abstractmethod	4✔
92	def __call__(self, *args: Any) -> Any:	4✔
93	pass	×
94
95	@abstractmethod	4✔
96	def compile(self, *args: Any, output_folder: Optional[str] = None) -> CompiledSDFG:	4✔
97	pass	×
98
99	def _resolve_compile_options(	4✔
100	self,
101	instrumentation_mode: Optional[str] = None,
102	capture_args: Optional[bool] = None,
103	remote_tuning: Optional[bool] = None,
104	) -> tuple[str, bool, bool]:
105	"""Resolve compile-time options, falling back to instance defaults and env vars."""
106	if instrumentation_mode is None:	4✔
107	instrumentation_mode = self.instrumentation_mode	4✔
108	if capture_args is None:	4✔
109	capture_args = self.capture_args	4✔
110	if remote_tuning is None:	4✔
111	remote_tuning = self.remote_tuning	4✔
112
113	# Check environment variable DOCC_CI
114	docc_ci = os.environ.get("DOCC_CI", "")	4✔
115	if docc_ci:	4✔
116	if docc_ci == "regions":	×
117	if instrumentation_mode is None:	×
118	instrumentation_mode = "ols"	×
119	elif docc_ci == "arg-capture":	×
120	if capture_args is None:	×
121	capture_args = True	×
122	else:
123	# Full mode (or unknown value treated as full)
124	if instrumentation_mode is None:	×
125	instrumentation_mode = "ols"	×
126	if capture_args is None:	×
127	capture_args = True	×
128
129	# Defaults
130	if instrumentation_mode is None:	4✔
131	instrumentation_mode = ""	4✔
132	if capture_args is None:	4✔
133	capture_args = False	4✔
134
135	return instrumentation_mode, capture_args, remote_tuning	4✔
136
137	def sdfg_pipe(	4✔
138	self,
139	sdfg: StructuredSDFG,
140	output_folder: Optional[str],
141	instrumentation_mode: str,
142	capture_args: bool,
143	remote_tuning: Optional[bool] = None,
144	) -> str:
145
146	if self.debug_dump:	4✔
147	sdfg.dump(output_folder, "py0.parsed", dump_dot=True)	×
148
149	# Enable statistics if envvar is set
150	if _statistics_enabled_by_env():	4✔
151	_enable_statistics()	×
152
153	sdfg.validate()	4✔
154
155	if remote_tuning is None:	4✔
156	remote_tuning = self.remote_tuning	×
157
158	target_options = TargetOptions()	4✔
159	target_options.target = self.target	4✔
160	target_options.category = self.category	4✔
161	target_options.remote_tuning = remote_tuning	4✔
162
163	# Einsum detection
164	sdfg.einsum()	4✔
165	if self.debug_dump:	4✔
166	sdfg.dump(output_folder, "py1.einsum", dump_dot=True)	×
167
168	# Tensor targets keep tensor nodes
169	custom_expand_fn = get_target_expand_fn(self.target)	4✔
170	if custom_expand_fn is not None:	4✔
171	custom_expand_fn(sdfg, self.category, {})	4✔
172	else:
173	sdfg.expand(target_options)	4✔
174	if self.debug_dump:	4✔
175	sdfg.dump(output_folder, "py2.expanded", dump_dot=True)	×
176
177	# Simplify pipelines
178	sdfg.simplify()	4✔
179	if self.debug_dump:	4✔
180	sdfg.dump(output_folder, "py3.opt", dump_dot=True)	×
181
182	# Normalization for scheduling
183	if self.target != "none":	4✔
184	sdfg.normalize()	4✔
185
186	if self.debug_dump or instrumentation_mode or capture_args:	4✔
187	sdfg.dump(	4✔
188	output_folder,
189	"py4.norm",
190	dump_dot=self.debug_dump,
191	dump_json=True,
192	record_for_instrumentation=True,
193	)
194
195	# Schedule if target is specified
196
197	custom_schedule_fn = get_target_schedule_fn(self.target)	4✔
198	if custom_schedule_fn is not None:	4✔
199	custom_schedule_fn(sdfg, self.category, {"remote_tuning": remote_tuning})	4✔
200	else:
201	sdfg.schedule(target_options)	4✔
202
203	if self.debug_dump:	4✔
204	sdfg.dump(output_folder, "py5.post_sched", dump_dot=True)	×
205
206	# Promote pointer arguments to device residency when the whole program keeps
207	# data on device. Communicated explicitly via the pass return value (bool),
208	# not through SDFG metadata.
209	self._device_resident = False	4✔
210	self._device_backend = None	4✔
211	if self.target in ("cuda", "rocm"):	4✔
NEW 212	if sdfg.promote_device_residency(self.target == "rocm"):	×
NEW 213	self._device_resident = True	×
NEW 214	self._device_backend = self.target	×
215
216	self.last_sdfg = sdfg	4✔
217
218	custom_compile_fn = get_target_compile_fn(self.target)	4✔
219	if custom_compile_fn is not None:	4✔
220	lib_path = custom_compile_fn(	4✔
221	sdfg,
222	output_folder,
223	instrumentation_mode,
224	capture_args,
225	{"debug_build": self.debug_build, "threads": self.build_thread_count},
226	)
227	else:
228	lib_path = sdfg._compile(	4✔
229	output_folder=output_folder,
230	target=self.target,
231	instrumentation_mode=instrumentation_mode,
232	capture_args=capture_args,
233	debug_build=self.debug_build,
234	threads=self.build_thread_count,
235	)
236
237	# Dump statistics after compile
238	if _statistics_enabled_by_env():	4✔
239	print(_statistics_summary(), file=sys.stderr)	×
240
241	return lib_path	4✔
242
243	@abstractmethod	4✔
244	def to_sdfg(self, *args: Any) -> StructuredSDFG:	4✔
245	pass	×
246
247	@abstractmethod	4✔
248	def _convert_inputs(self, args: tuple) -> tuple:	4✔
249	pass	×
250
251	@abstractmethod	4✔
252	def _convert_outputs(self, result: Any, original_args: tuple) -> Any:	4✔
253	pass	×
254
255	def _get_cache_key(self, *args: Any) -> str:	4✔
256	return ""	×

daisytuner / docc / 27885489129

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