28855385396

Committed 07 Jul 2026 09:17AM UTC coverage: 62.959% (-0.001%) from 62.96%

Build # 28855385396

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #810 from daisytuner/ReuseSources

Add DOCC_REUSE_SOURCES envar to recompile the shared library from locally edited sources

Coverage Stats

37 of 46 new or added lines in 1 file covered. (80.43%)

3 existing lines in 1 file now uncovered.

40560 of 64423 relevant lines covered (62.96%)

962.71 hits per line

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

73.26

/python/docc/compiler/docc_program.py

from abc import ABC, abstractmethod
import sys
from typing import Any, Dict, Optional
import json
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,
        reuse_sources: bool = False,
    ) -> str:

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

            if not output_folder is None:
                sdfg.output_dir = output_folder

            # 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)

            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,
                reuse_sources=reuse_sources,
            )

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

        # Record the device-residency decision in the persisted (py4.norm) SDFG
        # metadata. It is computed here (not stored in metadata by the pass) and
        # decides host vs device argument marshalling. Binary-reuse paths
        # (DOCC_REUSE_BINARIES) load only the cached .so + normalized SDFG and
        # never re-run scheduling/promotion, so without this they default to
        # host execution and feed host pointers into a device-resident binary
        # -> heap corruption / double free.
        if output_folder:
            self._persist_device_residency(output_folder, sdfg)

        return lib_path

    def _persist_device_residency(
        self, output_folder: str, sdfg: StructuredSDFG
    ) -> None:
        """Stamp the device-residency decision into the persisted SDFG metadata.

        Patches only the ``metadata`` object of the already-written
        ``py4.norm.json`` (the file the reuse path loads), leaving the SDFG
        structure and element IDs untouched so instrumentation references stay
        valid.
        """
        json_path = os.path.join(output_folder, f"{sdfg.name}.py4.norm.json")
        try:
            with open(json_path) as f:
                data = json.load(f)
            metadata = data.setdefault("metadata", {})
            metadata["device_resident"] = "1" if self._device_resident else "0"
            metadata["device_backend"] = self._device_backend or ""
            with open(json_path, "w") as f:
                json.dump(data, f)
        except (OSError, ValueError):
            pass

    @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 json	4✔
5	import os	4✔
6	import re	4✔
7
8	from docc.sdfg import StructuredSDFG, TargetOptions	4✔
9	from docc.sdfg._sdfg import (	4✔
10	_enable_statistics,
11	_statistics_enabled_by_env,
12	_statistics_summary,
13	)
14	from docc.compiler.compiled_sdfg import CompiledSDFG	4✔
15	from docc.compiler.target_registry import (	4✔
16	get_target_schedule_fn,
17	get_target_compile_fn,
18	get_target_expand_fn,
19	register_target_overrides,
20	)
21
22
23	def _parse_docc_debug() -> dict[str, str]:	4✔
24	debug_env = os.environ.get("DOCC_DEBUG", "")	4✔
25	debug_dict = {}	4✔
26	if debug_env:	4✔
27	for entry in re.split(r"[;:]", debug_env):	×
28	if not entry:	×
29	continue	×
30	parts = entry.split("=", 1)	×
31	key = parts[0].strip()	×
32	value = parts[1].strip() if len(parts) > 1 else ""	×
33	debug_dict[key] = value	×
34	return debug_dict	4✔
35
36
37	def _is_debug_dump(flags: dict[str, str]) -> bool:	4✔
38	return "dump" in flags	4✔
39
40
41	def _is_debug_compile(flags: dict[str, str]) -> bool:	4✔
42	return "build" in flags	4✔
43
44
45	def _get_build_thread_count(flags: dict[str, str]) -> int:	4✔
46	return int(flags.get("build_threads", "0"))	4✔
47
48
49	class DoccProgram(ABC):	4✔
50
51	def __init__(	4✔
52	self,
53	name: str,
54	target: str = "none",
55	category: str = "server",
56	instrumentation_mode: Optional[str] = None,
57	capture_args: Optional[bool] = None,
58	remote_tuning: bool = False,
59	):
60	self.name = name	4✔
61	self.target = target	4✔
62	self.category = category	4✔
63	self.remote_tuning = remote_tuning	4✔
64	self.last_sdfg: Optional[StructuredSDFG] = None	4✔
65	self._device_resident: bool = False	4✔
66	self._device_backend: Optional[str] = None	4✔
67	self.cache: dict = {}	4✔
68	debug_flags = _parse_docc_debug()	4✔
69	self.debug_dump: bool = _is_debug_dump(debug_flags)	4✔
70	self.debug_build: bool = _is_debug_compile(debug_flags)	4✔
71	self.build_thread_count: int = _get_build_thread_count(debug_flags)	4✔
72
73	# Check environment variable DOCC_CI
74	docc_ci = os.environ.get("DOCC_CI", "")	4✔
75	if docc_ci:	4✔
76	if docc_ci == "regions":	×
77	if instrumentation_mode is None:	×
78	instrumentation_mode = "ols"	×
79	elif docc_ci == "arg-capture":	×
80	if capture_args is None:	×
81	capture_args = True	×
82	else:
83	# Full mode (or unknown value treated as full)
84	if instrumentation_mode is None:	×
85	instrumentation_mode = "ols"	×
86	if capture_args is None:	×
87	capture_args = True	×
88
89	self.instrumentation_mode = instrumentation_mode	4✔
90	self.capture_args = capture_args	4✔
91
92	@abstractmethod	4✔
93	def __call__(self, *args: Any) -> Any:	4✔
94	pass	×
95
96	@abstractmethod	4✔
97	def compile(self, *args: Any, output_folder: Optional[str] = None) -> CompiledSDFG:	4✔
98	pass	×
99
100	def _resolve_compile_options(	4✔
101	self,
102	instrumentation_mode: Optional[str] = None,
103	capture_args: Optional[bool] = None,
104	remote_tuning: Optional[bool] = None,
105	) -> tuple[str, bool, bool]:
106	"""Resolve compile-time options, falling back to instance defaults and env vars."""
107	if instrumentation_mode is None:	4✔
108	instrumentation_mode = self.instrumentation_mode	4✔
109	if capture_args is None:	4✔
110	capture_args = self.capture_args	4✔
111	if remote_tuning is None:	4✔
112	remote_tuning = self.remote_tuning	4✔
113
114	# Check environment variable DOCC_CI
115	docc_ci = os.environ.get("DOCC_CI", "")	4✔
116	if docc_ci:	4✔
117	if docc_ci == "regions":	×
118	if instrumentation_mode is None:	×
119	instrumentation_mode = "ols"	×
120	elif docc_ci == "arg-capture":	×
121	if capture_args is None:	×
122	capture_args = True	×
123	else:
124	# Full mode (or unknown value treated as full)
125	if instrumentation_mode is None:	×
126	instrumentation_mode = "ols"	×
127	if capture_args is None:	×
128	capture_args = True	×
129
130	# Defaults
131	if instrumentation_mode is None:	4✔
132	instrumentation_mode = ""	4✔
133	if capture_args is None:	4✔
134	capture_args = False	4✔
135
136	return instrumentation_mode, capture_args, remote_tuning	4✔
137
138	def sdfg_pipe(	4✔
139	self,
140	sdfg: StructuredSDFG,
141	output_folder: Optional[str],
142	instrumentation_mode: str,
143	capture_args: bool,
144	remote_tuning: Optional[bool] = None,
145	reuse_sources: bool = False,
146	) -> str:
147
148	if not reuse_sources and output_folder:	4✔
149	if self.debug_dump:	4✔
NEW 150	sdfg.dump(output_folder, "py0.parsed", dump_dot=True)	×
151
152	if not output_folder is None:	4✔
153	sdfg.output_dir = output_folder	4✔
154
155	# Enable statistics if envvar is set
156	if _statistics_enabled_by_env():	4✔
NEW 157	_enable_statistics()	×
158
159	sdfg.validate()	4✔
160
161	if remote_tuning is None:	4✔
NEW 162	remote_tuning = self.remote_tuning	×
163
164	target_options = TargetOptions()	4✔
165	target_options.target = self.target	4✔
166	target_options.category = self.category	4✔
167	target_options.remote_tuning = remote_tuning	4✔
168
169	# Einsum detection
170	sdfg.einsum()	4✔
171	if self.debug_dump:	4✔
NEW 172	sdfg.dump(output_folder, "py1.einsum", dump_dot=True)	×
173
174	sdfg.validate()	4✔
175
176	if remote_tuning is None:	4✔
NEW 177	remote_tuning = self.remote_tuning	×
178
179	target_options = TargetOptions()	4✔
180	target_options.target = self.target	4✔
181	target_options.category = self.category	4✔
182	target_options.remote_tuning = remote_tuning	4✔
183
184	# Einsum detection
185	sdfg.einsum()	4✔
186	if self.debug_dump:	4✔
NEW 187	sdfg.dump(output_folder, "py1.einsum", dump_dot=True)	×
188
189	# Tensor targets keep tensor nodes
190	custom_expand_fn = get_target_expand_fn(self.target)	4✔
191	if custom_expand_fn is not None:	4✔
192	custom_expand_fn(sdfg, self.category, {})	4✔
193	else:
194	sdfg.expand(target_options)	4✔
195	if self.debug_dump:	4✔
NEW 196	sdfg.dump(output_folder, "py2.expanded", dump_dot=True)	×
197
198	# Simplify pipelines
199	sdfg.simplify()	4✔
200	if self.debug_dump:	4✔
NEW 201	sdfg.dump(output_folder, "py3.opt", dump_dot=True)	×
202
203	# Normalization for scheduling
204	if self.target != "none":	4✔
205	sdfg.normalize()	4✔
206
207	if self.debug_dump or instrumentation_mode or capture_args:	4✔
208	sdfg.dump(	4✔
209	output_folder,
210	"py4.norm",
211	dump_dot=self.debug_dump,
212	dump_json=True,
213	record_for_instrumentation=True,
214	)
215
216	# Schedule if target is specified
217
218	custom_schedule_fn = get_target_schedule_fn(self.target)	4✔
219	if custom_schedule_fn is not None:	4✔
220	custom_schedule_fn(	4✔
221	sdfg, self.category, {"remote_tuning": remote_tuning}
222	)
223	else:
224	sdfg.schedule(target_options)	4✔
225
226	if self.debug_dump:	4✔
NEW 227	sdfg.dump(output_folder, "py5.post_sched", dump_dot=True)	×
228
229	# Promote pointer arguments to device residency when the whole program keeps
230	# data on device. Communicated explicitly via the pass return value (bool),
231	# not through SDFG metadata.
232	self._device_resident = False	4✔
233	self._device_backend = None	4✔
234	if self.target in ("cuda", "rocm"):	4✔
235	if sdfg.promote_device_residency(self.target == "rocm"):	×
236	self._device_resident = True	×
237	self._device_backend = self.target	×
238
239	self.last_sdfg = sdfg	4✔
240
241	custom_compile_fn = get_target_compile_fn(self.target)	4✔
242	if custom_compile_fn is not None:	4✔
243	lib_path = custom_compile_fn(	4✔
244	sdfg,
245	output_folder,
246	instrumentation_mode,
247	capture_args,
248	{"debug_build": self.debug_build, "threads": self.build_thread_count},
249	)
250	else:
251	lib_path = sdfg._compile(	4✔
252	output_folder=output_folder,
253	target=self.target,
254	instrumentation_mode=instrumentation_mode,
255	capture_args=capture_args,
256	debug_build=self.debug_build,
257	threads=self.build_thread_count,
258	reuse_sources=reuse_sources,
259	)
260
261	# Dump statistics after compile
262	if _statistics_enabled_by_env():	4✔
263	print(_statistics_summary(), file=sys.stderr)	×
264
265	# Record the device-residency decision in the persisted (py4.norm) SDFG
266	# metadata. It is computed here (not stored in metadata by the pass) and
267	# decides host vs device argument marshalling. Binary-reuse paths
268	# (DOCC_REUSE_BINARIES) load only the cached .so + normalized SDFG and
269	# never re-run scheduling/promotion, so without this they default to
270	# host execution and feed host pointers into a device-resident binary
271	# -> heap corruption / double free.
272	if output_folder:	4✔
273	self._persist_device_residency(output_folder, sdfg)	4✔
274
275	return lib_path	4✔
276
277	def _persist_device_residency(	4✔
278	self, output_folder: str, sdfg: StructuredSDFG
279	) -> None:
280	"""Stamp the device-residency decision into the persisted SDFG metadata.
281
282	Patches only the ``metadata`` object of the already-written
283	``py4.norm.json`` (the file the reuse path loads), leaving the SDFG
284	structure and element IDs untouched so instrumentation references stay
285	valid.
286	"""
287	json_path = os.path.join(output_folder, f"{sdfg.name}.py4.norm.json")	4✔
288	try:	4✔
289	with open(json_path) as f:	4✔
290	data = json.load(f)	4✔
291	metadata = data.setdefault("metadata", {})	4✔
292	metadata["device_resident"] = "1" if self._device_resident else "0"	4✔
293	metadata["device_backend"] = self._device_backend or ""	4✔
294	with open(json_path, "w") as f:	4✔
295	json.dump(data, f)	4✔
296	except (OSError, ValueError):	4✔
297	pass	4✔
298
299	@abstractmethod	4✔
300	def to_sdfg(self, *args: Any) -> StructuredSDFG:	4✔
301	pass	×
302
303	@abstractmethod	4✔
304	def _convert_inputs(self, args: tuple) -> tuple:	4✔
305	pass	×
306
307	@abstractmethod	4✔
308	def _convert_outputs(self, result: Any, original_args: tuple) -> Any:	4✔
309	pass	×
310
311	def _get_cache_key(self, *args: Any) -> str:	4✔
312	return ""	×

daisytuner / docc / 28855385396

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