25501744090

Committed 07 May 2026 02:21PM UTC coverage: 65.649%. First build

Build # 25501744090

Build Type

Pull #681

github

Committed by

web-flow

Commit Message

Merge 444892749 into 8010510fa

Pull Request Pull Request #681: Target specific expand

Coverage Stats

58 of 71 new or added lines in 5 files covered. (81.69%)

32178 of 49015 relevant lines covered (65.65%)

12931.86 hits per line

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

71.55

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


def _cuda_expand_fn(sdfg, category: str, kwargs: Dict[str, Any]) -> None:
    sdfg.expand_cuda()
    sdfg.expand()


def _rocm_expand_fn(sdfg, category: str, kwargs: Dict[str, Any]) -> None:
    sdfg.expand_rocm()
    sdfg.expand()


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

        if target == "cuda":
            from docc.python import register_target_overrides

            register_target_overrides(
                "cuda",
                schedule_fn=None,
                compile_fn=None,
                expand_fn=_cuda_expand_fn,
            )
        elif target == "rocm":
            from docc.python import register_target_overrides

            register_target_overrides(
                "rocm",
                schedule_fn=None,
                compile_fn=None,
                expand_fn=_rocm_expand_fn,
            )

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

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

    def sdfg_pipe(
        self,
        sdfg: StructuredSDFG,
        output_folder: Optional[str],
        instrumentation_mode: str,
        capture_args: bool,
    ) -> 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()

        # 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()
        if self.debug_dump:
            sdfg.dump(output_folder, "py1.expanded", dump_dot=True)

        # Simplify pipelines
        sdfg.simplify()
        if self.debug_dump:
            sdfg.dump(output_folder, "py2.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,
                "py3.norm",
                dump_dot=self.debug_dump,
                dump_json=True,
                record_for_instrumentation=True,
            )

        # Schedule if target is specified

        if self.target != "none":
            # Check for custom registered target first
            custom_schedule_fn = get_target_schedule_fn(self.target)
            if custom_schedule_fn is not None:
                custom_schedule_fn(
                    sdfg, self.category, {"remote_tuning": self.remote_tuning}
                )
            else:
                sdfg.schedule(self.target, self.category, self.remote_tuning)

            if self.debug_dump:
                sdfg.dump(output_folder, "py4.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	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	)
19
20
21	def _cuda_expand_fn(sdfg, category: str, kwargs: Dict[str, Any]) -> None:	4✔
22	sdfg.expand_cuda()	4✔
23	sdfg.expand()	4✔
24
25
26	def _rocm_expand_fn(sdfg, category: str, kwargs: Dict[str, Any]) -> None:	4✔
NEW 27	sdfg.expand_rocm()	×
NEW 28	sdfg.expand()	×
29
30
31	def _parse_docc_debug() -> dict[str, str]:	4✔
32	debug_env = os.environ.get("DOCC_DEBUG", "")	4✔
33	debug_dict = {}	4✔
34	if debug_env:	4✔
35	for entry in re.split(r"[;:]", debug_env):	×
36	if not entry:	×
37	continue	×
38	parts = entry.split("=", 1)	×
39	key = parts[0].strip()	×
40	value = parts[1].strip() if len(parts) > 1 else ""	×
41	debug_dict[key] = value	×
42	return debug_dict	4✔
43
44
45	def _is_debug_dump(flags: dict[str, str]) -> bool:	4✔
46	return "dump" in flags	4✔
47
48
49	def _is_debug_compile(flags: dict[str, str]) -> bool:	4✔
50	return "build" in flags	4✔
51
52
53	def _get_build_thread_count(flags: dict[str, str]) -> int:	4✔
54	return int(flags.get("build_threads", "0"))	4✔
55
56
57	class DoccProgram(ABC):	4✔
58
59	def __init__(	4✔
60	self,
61	name: str,
62	target: str = "none",
63	category: str = "server",
64	instrumentation_mode: Optional[str] = None,
65	capture_args: Optional[bool] = None,
66	remote_tuning: bool = False,
67	):
68	self.name = name	4✔
69	self.target = target	4✔
70	self.category = category	4✔
71	self.remote_tuning = remote_tuning	4✔
72	self.last_sdfg: Optional[StructuredSDFG] = None	4✔
73	self.cache: dict = {}	4✔
74	debug_flags = _parse_docc_debug()	4✔
75	self.debug_dump: bool = _is_debug_dump(debug_flags)	4✔
76	self.debug_build: bool = _is_debug_compile(debug_flags)	4✔
77	self.build_thread_count: int = _get_build_thread_count(debug_flags)	4✔
78
79	# Check environment variable DOCC_CI
80	docc_ci = os.environ.get("DOCC_CI", "")	4✔
81	if docc_ci:	4✔
82	if docc_ci == "regions":	×
83	if instrumentation_mode is None:	×
84	instrumentation_mode = "ols"	×
85	elif docc_ci == "arg-capture":	×
86	if capture_args is None:	×
87	capture_args = True	×
88	else:
89	# Full mode (or unknown value treated as full)
90	if instrumentation_mode is None:	×
91	instrumentation_mode = "ols"	×
92	if capture_args is None:	×
93	capture_args = True	×
94
95	self.instrumentation_mode = instrumentation_mode	4✔
96	self.capture_args = capture_args	4✔
97
98	if target == "cuda":	4✔
99	from docc.python import register_target_overrides	4✔
100
101	register_target_overrides(	4✔
102	"cuda",
103	schedule_fn=None,
104	compile_fn=None,
105	expand_fn=_cuda_expand_fn,
106	)
107	elif target == "rocm":	4✔
NEW 108	from docc.python import register_target_overrides	×
109
NEW 110	register_target_overrides(	×
111	"rocm",
112	schedule_fn=None,
113	compile_fn=None,
114	expand_fn=_rocm_expand_fn,
115	)
116
117	@abstractmethod	4✔
118	def __call__(self, *args: Any) -> Any:	4✔
119	pass	×
120
121	@abstractmethod	4✔
122	def compile(self, *args: Any, output_folder: Optional[str] = None) -> CompiledSDFG:	4✔
123	pass	×
124
125	def sdfg_pipe(	4✔
126	self,
127	sdfg: StructuredSDFG,
128	output_folder: Optional[str],
129	instrumentation_mode: str,
130	capture_args: bool,
131	) -> str:
132
133	if self.debug_dump:	4✔
134	sdfg.dump(output_folder, "py0.parsed", dump_dot=True)	×
135
136	# Enable statistics if envvar is set
137	if _statistics_enabled_by_env():	4✔
138	_enable_statistics()	×
139
140	sdfg.validate()	4✔
141
142	# Tensor targets keep tensor nodes
143	custom_expand_fn = get_target_expand_fn(self.target)	4✔
144	if custom_expand_fn is not None:	4✔
145	custom_expand_fn(sdfg, self.category, {})	4✔
146	else:
147	sdfg.expand()	4✔
148	if self.debug_dump:	4✔
149	sdfg.dump(output_folder, "py1.expanded", dump_dot=True)	×
150
151	# Simplify pipelines
152	sdfg.simplify()	4✔
153	if self.debug_dump:	4✔
154	sdfg.dump(output_folder, "py2.opt", dump_dot=True)	×
155
156	# Normalization for scheduling
157	if self.target != "none":	4✔
158	sdfg.normalize()	4✔
159
160	if self.debug_dump or instrumentation_mode or capture_args:	4✔
161	sdfg.dump(	4✔
162	output_folder,
163	"py3.norm",
164	dump_dot=self.debug_dump,
165	dump_json=True,
166	record_for_instrumentation=True,
167	)
168
169	# Schedule if target is specified
170
171	if self.target != "none":	4✔
172	# Check for custom registered target first
173	custom_schedule_fn = get_target_schedule_fn(self.target)	4✔
174	if custom_schedule_fn is not None:	4✔
175	custom_schedule_fn(	4✔
176	sdfg, self.category, {"remote_tuning": self.remote_tuning}
177	)
178	else:
179	sdfg.schedule(self.target, self.category, self.remote_tuning)	4✔
180
181	if self.debug_dump:	4✔
182	sdfg.dump(output_folder, "py4.post_sched", dump_dot=True)	×
183
184	self.last_sdfg = sdfg	4✔
185
186	custom_compile_fn = get_target_compile_fn(self.target)	4✔
187	if custom_compile_fn is not None:	4✔
188	lib_path = custom_compile_fn(	4✔
189	sdfg,
190	output_folder,
191	instrumentation_mode,
192	capture_args,
193	{"debug_build": self.debug_build, "threads": self.build_thread_count},
194	)
195	else:
196	lib_path = sdfg._compile(	4✔
197	output_folder=output_folder,
198	target=self.target,
199	instrumentation_mode=instrumentation_mode,
200	capture_args=capture_args,
201	debug_build=self.debug_build,
202	threads=self.build_thread_count,
203	)
204
205	# Dump statistics after compile
206	if _statistics_enabled_by_env():	4✔
207	print(_statistics_summary(), file=sys.stderr)	×
208
209	return lib_path	4✔
210
211	@abstractmethod	4✔
212	def to_sdfg(self, *args: Any) -> StructuredSDFG:	4✔
213	pass	×
214
215	@abstractmethod	4✔
216	def _convert_inputs(self, args: tuple) -> tuple:	4✔
217	pass	×
218
219	@abstractmethod	4✔
220	def _convert_outputs(self, result: Any, original_args: tuple) -> Any:	4✔
221	pass	×
222
223	def _get_cache_key(self, *args: Any) -> str:	4✔
224	return ""	×

daisytuner / docc / 25501744090

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