26804877914

Committed 02 Jun 2026 07:22AM UTC coverage: 60.835% (-0.03%) from 60.869%

Build # 26804877914

Build Type

Pull #725

github

Committed by

web-flow

Commit Message

Merge ad1f0d80a into cd25c9278

Pull Request Pull Request #725: Tensor node backport

Coverage Stats

599 of 1251 new or added lines in 50 files covered. (47.88%)

538 existing lines in 45 files now uncovered.

35099 of 57695 relevant lines covered (60.84%)

11081.73 hits per line

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

69.85

/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.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)

        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.cache: dict = {}	4✔
65	debug_flags = _parse_docc_debug()	4✔
66	self.debug_dump: bool = _is_debug_dump(debug_flags)	4✔
67	self.debug_build: bool = _is_debug_compile(debug_flags)	4✔
68	self.build_thread_count: int = _get_build_thread_count(debug_flags)	4✔
69
70	# Check environment variable DOCC_CI
71	docc_ci = os.environ.get("DOCC_CI", "")	4✔
72	if docc_ci:	4✔
73	if docc_ci == "regions":	×
74	if instrumentation_mode is None:	×
75	instrumentation_mode = "ols"	×
76	elif docc_ci == "arg-capture":	×
77	if capture_args is None:	×
78	capture_args = True	×
79	else:
80	# Full mode (or unknown value treated as full)
81	if instrumentation_mode is None:	×
82	instrumentation_mode = "ols"	×
83	if capture_args is None:	×
84	capture_args = True	×
85
86	self.instrumentation_mode = instrumentation_mode	4✔
87	self.capture_args = capture_args	4✔
88
89	@abstractmethod	4✔
90	def __call__(self, *args: Any) -> Any:	4✔
91	pass	×
92
93	@abstractmethod	4✔
94	def compile(self, *args: Any, output_folder: Optional[str] = None) -> CompiledSDFG:	4✔
95	pass	×
96
97	def _resolve_compile_options(	4✔
98	self,
99	instrumentation_mode: Optional[str] = None,
100	capture_args: Optional[bool] = None,
101	remote_tuning: Optional[bool] = None,
102	) -> tuple[str, bool, bool]:
103	"""Resolve compile-time options, falling back to instance defaults and env vars."""
104	if instrumentation_mode is None:	4✔
105	instrumentation_mode = self.instrumentation_mode	4✔
106	if capture_args is None:	4✔
107	capture_args = self.capture_args	4✔
108	if remote_tuning is None:	4✔
109	remote_tuning = self.remote_tuning	4✔
110
111	# Check environment variable DOCC_CI
112	docc_ci = os.environ.get("DOCC_CI", "")	4✔
113	if docc_ci:	4✔
UNCOV 114	if docc_ci == "regions":	×
UNCOV 115	if instrumentation_mode is None:	×
UNCOV 116	instrumentation_mode = "ols"	×
UNCOV 117	elif docc_ci == "arg-capture":	×
UNCOV 118	if capture_args is None:	×
UNCOV 119	capture_args = True	×
120	else:
121	# Full mode (or unknown value treated as full)
122	if instrumentation_mode is None:	×
UNCOV 123	instrumentation_mode = "ols"	×
UNCOV 124	if capture_args is None:	×
UNCOV 125	capture_args = True	×
126
127	# Defaults
128	if instrumentation_mode is None:	4✔
129	instrumentation_mode = ""	4✔
130	if capture_args is None:	4✔
131	capture_args = False	4✔
132
133	return instrumentation_mode, capture_args, remote_tuning	4✔
134
135	def sdfg_pipe(	4✔
136	self,
137	sdfg: StructuredSDFG,
138	output_folder: Optional[str],
139	instrumentation_mode: str,
140	capture_args: bool,
141	remote_tuning: Optional[bool] = None,
142	) -> str:
143
144	if self.debug_dump:	4✔
UNCOV 145	sdfg.dump(output_folder, "py0.parsed", dump_dot=True)	×
146
147	# Enable statistics if envvar is set
148	if _statistics_enabled_by_env():	4✔
UNCOV 149	_enable_statistics()	×
150
151	sdfg.validate()	4✔
152
153	if remote_tuning is None:	4✔
UNCOV 154	remote_tuning = self.remote_tuning	×
155
156	target_options = TargetOptions()	4✔
157	target_options.target = self.target	4✔
158	target_options.category = self.category	4✔
159	target_options.remote_tuning = remote_tuning	4✔
160
161	# Einsum detection
162	sdfg.einsum()	4✔
163	if self.debug_dump:	4✔
UNCOV 164	sdfg.dump(output_folder, "py1.einsum", dump_dot=True)	×
165
166	# Tensor targets keep tensor nodes
167	custom_expand_fn = get_target_expand_fn(self.target)	4✔
168	if custom_expand_fn is not None:	4✔
169	custom_expand_fn(sdfg, self.category, {})	4✔
170	else:
171	sdfg.expand(target_options)	4✔
172	if self.debug_dump:	4✔
UNCOV 173	sdfg.dump(output_folder, "py2.expanded", dump_dot=True)	×
174
175	# Simplify pipelines
176	sdfg.simplify()	4✔
177	if self.debug_dump:	4✔
UNCOV 178	sdfg.dump(output_folder, "py3.opt", dump_dot=True)	×
179
180	# Normalization for scheduling
181	if self.target != "none":	4✔
182	sdfg.normalize()	4✔
183
184	if self.debug_dump or instrumentation_mode or capture_args:	4✔
185	sdfg.dump(	4✔
186	output_folder,
187	"py4.norm",
188	dump_dot=self.debug_dump,
189	dump_json=True,
190	record_for_instrumentation=True,
191	)
192
193	# Schedule if target is specified
194
195	custom_schedule_fn = get_target_schedule_fn(self.target)	4✔
196	if custom_schedule_fn is not None:	4✔
197	custom_schedule_fn(sdfg, self.category, {"remote_tuning": remote_tuning})	4✔
198	else:
199	sdfg.schedule(target_options)	4✔
200
201	if self.debug_dump:	4✔
UNCOV 202	sdfg.dump(output_folder, "py5.post_sched", dump_dot=True)	×
203
204	self.last_sdfg = sdfg	4✔
205
206	custom_compile_fn = get_target_compile_fn(self.target)	4✔
207	if custom_compile_fn is not None:	4✔
208	lib_path = custom_compile_fn(	4✔
209	sdfg,
210	output_folder,
211	instrumentation_mode,
212	capture_args,
213	{"debug_build": self.debug_build, "threads": self.build_thread_count},
214	)
215	else:
216	lib_path = sdfg._compile(	4✔
217	output_folder=output_folder,
218	target=self.target,
219	instrumentation_mode=instrumentation_mode,
220	capture_args=capture_args,
221	debug_build=self.debug_build,
222	threads=self.build_thread_count,
223	)
224
225	# Dump statistics after compile
226	if _statistics_enabled_by_env():	4✔
UNCOV 227	print(_statistics_summary(), file=sys.stderr)	×
228
229	return lib_path	4✔
230
231	@abstractmethod	4✔
232	def to_sdfg(self, *args: Any) -> StructuredSDFG:	4✔
UNCOV 233	pass	×
234
235	@abstractmethod	4✔
236	def _convert_inputs(self, args: tuple) -> tuple:	4✔
UNCOV 237	pass	×
238
239	@abstractmethod	4✔
240	def _convert_outputs(self, result: Any, original_args: tuple) -> Any:	4✔
UNCOV 241	pass	×
242
243	def _get_cache_key(self, *args: Any) -> str:	4✔
UNCOV 244	return ""	×

daisytuner / docc / 26804877914

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