21112243720

Committed 18 Jan 2026 01:08PM UTC coverage: 64.188% (+0.03%) from 64.154%

Build # 21112243720

Build Type

Pull #462

github

Committed by

web-flow

Commit Message

Merge d6bf7a605 into 92e9cbdc3

Pull Request Pull Request #462: adds syntax support for multi-assignments and np.empty_like

Run Details

31 of 33 new or added lines in 2 files covered. (93.94%)

197 existing lines in 5 files now uncovered.

19497 of 30375 relevant lines covered (64.19%)

387.69 hits per line

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

84.76

/python/docc/__init__.py

import inspect
import shutil
import textwrap
import ast
import os
import getpass
import hashlib
import numpy as np
from typing import Annotated, get_origin, get_args
from ._sdfg import *
from .compiled_sdfg import CompiledSDFG
from .ast_parser import ASTParser


def _compile_wrapper(self, output_folder=None):
    lib_path = self._compile(output_folder)
    return CompiledSDFG(lib_path, self)


StructuredSDFG.compile = _compile_wrapper


def _map_python_type(dtype):
    # If it is already a sdfg Type, return it
    if isinstance(dtype, Type):
        return dtype

    # Handle Annotated for Arrays
    if get_origin(dtype) is Annotated:
        args = get_args(dtype)
        base_type = args[0]
        metadata = args[1:]

        if base_type is np.ndarray:
            # Convention: Annotated[np.ndarray, shape, dtype]
            shape = metadata[0]
            elem_type = Scalar(PrimitiveType.Double)  # Default

            if len(metadata) > 1:
                possible_dtype = metadata[1]
                elem_type = _map_python_type(possible_dtype)

            return Pointer(elem_type)

    # Handle numpy.ndarray[Shape, DType]
    if get_origin(dtype) is np.ndarray:
        args = get_args(dtype)
        # args[0] is shape, args[1] is dtype
        if len(args) >= 2:
            elem_type = _map_python_type(args[1])
            return Pointer(elem_type)

    # Simple mapping for python types
    if dtype is float or dtype is np.float64:
        return Scalar(PrimitiveType.Double)
    elif dtype is int or dtype is np.int64:
        return Scalar(PrimitiveType.Int64)
    elif dtype is bool or dtype is np.bool_:
        return Scalar(PrimitiveType.Bool)
    elif dtype is np.float32:
        return Scalar(PrimitiveType.Float)
    elif dtype is np.int32:
        return Scalar(PrimitiveType.Int32)

    # Handle Python classes - map to Structure type
    if inspect.isclass(dtype):
        # Use the class name as the structure name
        return Pointer(Structure(dtype.__name__))

    return dtype


class DoccProgram:
    def __init__(
        self,
        func,
        target="none",
        category="desktop",
        instrumentation_mode=None,
        capture_args=None,
    ):
        self.func = func
        self.name = func.__name__
        self.target = target
        self.category = category
        self.instrumentation_mode = instrumentation_mode
        self.capture_args = capture_args
        self.last_sdfg = None
        self.cache = {}

    def __call__(self, *args):
        # JIT compile and run
        compiled = self.compile(*args)
        res = compiled(*args)

        # Handle return value conversion based on annotation
        sig = inspect.signature(self.func)
        ret_annotation = sig.return_annotation

        if ret_annotation is not inspect.Signature.empty:
            if get_origin(ret_annotation) is Annotated:
                type_args = get_args(ret_annotation)
                if len(type_args) >= 1 and type_args[0] is np.ndarray:
                    shape = None
                    if len(type_args) >= 2:
                        shape = type_args[1]

                    if shape is not None:
                        try:
                            return np.ctypeslib.as_array(res, shape=shape)
                        except Exception:
                            pass

        # Try to infer return shape from metadata
        if hasattr(compiled, "get_return_shape"):
            shape = compiled.get_return_shape(*args)
            if shape is not None:
                try:
                    return np.ctypeslib.as_array(res, shape=shape)
                except Exception:
                    pass

        return res

    def compile(
        self, *args, output_folder=None, instrumentation_mode=None, capture_args=None
    ):
        original_output_folder = output_folder

        # Resolve options
        if instrumentation_mode is None:
            instrumentation_mode = self.instrumentation_mode
        if capture_args is None:
            capture_args = self.capture_args

        # 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

        # 1. Analyze arguments and shapes
        arg_types = []
        shape_values = []  # List of unique shape values found
        shape_sources = []  # List of (arg_idx, dim_idx) for each unique shape value

        # Mapping from (arg_idx, dim_idx) -> unique_shape_idx
        arg_shape_mapping = {}

        for i, arg in enumerate(args):
            t = self._infer_type(arg)
            arg_types.append(t)

            if isinstance(arg, np.ndarray):
                for dim_idx, dim_val in enumerate(arg.shape):
                    # Check if we've seen this value
                    if dim_val in shape_values:
                        # Reuse
                        u_idx = shape_values.index(dim_val)
                    else:
                        # New
                        u_idx = len(shape_values)
                        shape_values.append(dim_val)
                        shape_sources.append((i, dim_idx))

                    arg_shape_mapping[(i, dim_idx)] = u_idx

        # 2. Signature
        mapping_sig = sorted(arg_shape_mapping.items())
        type_sig = ", ".join(self._type_to_str(t) for t in arg_types)
        signature = f"{type_sig}|{mapping_sig}"

        if output_folder is None:
            filename = inspect.getsourcefile(self.func)
            hash_input = f"{filename}|{self.name}|{self.target}|{self.category}|{self.capture_args}|{self.instrumentation_mode}|{signature}".encode(
                "utf-8"
            )
            stable_id = hashlib.sha256(hash_input).hexdigest()[:16]

            docc_tmp = os.environ.get("DOCC_TMP")
            if docc_tmp:
                output_folder = f"{docc_tmp}/{self.name}-{stable_id}"
            else:
                user = getpass.getuser()
                output_folder = f"/tmp/{user}/DOCC/{self.name}-{stable_id}"

        if original_output_folder is None and signature in self.cache:
            return self.cache[signature]

        # 3. Build SDFG
        if os.path.exists(output_folder):
            # Multiple python processes running the same code?
            shutil.rmtree(output_folder)
        sdfg = self._build_sdfg(arg_types, args, arg_shape_mapping, len(shape_values))
        sdfg.expand()
        sdfg.simplify()

        if self.target != "none":
            sdfg.normalize()

        sdfg.dump(output_folder)

        # Schedule if target is specified
        if self.target != "none":
            sdfg.schedule(self.target, self.category)

        self.last_sdfg = sdfg

        lib_path = sdfg._compile(
            output_folder=output_folder,
            instrumentation_mode=instrumentation_mode,
            capture_args=capture_args,
        )

        # 5. Create CompiledSDFG
        compiled = CompiledSDFG(
            lib_path, sdfg, shape_sources, self._last_structure_member_info
        )

        # Cache if using default output folder
        if original_output_folder is None:
            self.cache[signature] = compiled

        return compiled

    def _get_signature(self, arg_types):
        return ", ".join(self._type_to_str(t) for t in arg_types)

    def _type_to_str(self, t):
        if isinstance(t, Scalar):
            return f"Scalar({t.primitive_type})"
        elif isinstance(t, Array):
            return f"Array({self._type_to_str(t.element_type)}, {t.num_elements})"
        elif isinstance(t, Pointer):
            return f"Pointer({self._type_to_str(t.pointee_type)})"
        elif isinstance(t, Structure):
            return f"Structure({t.name})"
        return str(t)

    def _infer_type(self, arg):
        if isinstance(arg, (bool, np.bool_)):
            return Scalar(PrimitiveType.Bool)
        elif isinstance(arg, (int, np.int64)):
            return Scalar(PrimitiveType.Int64)
        elif isinstance(arg, (float, np.float64)):
            return Scalar(PrimitiveType.Double)
        elif isinstance(arg, np.int32):
            return Scalar(PrimitiveType.Int32)
        elif isinstance(arg, np.float32):
            return Scalar(PrimitiveType.Float)
        elif isinstance(arg, np.ndarray):
            # Map dtype
            if arg.dtype == np.float64:
                elem_type = Scalar(PrimitiveType.Double)
            elif arg.dtype == np.float32:
                elem_type = Scalar(PrimitiveType.Float)
            elif arg.dtype == np.int64:
                elem_type = Scalar(PrimitiveType.Int64)
            elif arg.dtype == np.int32:
                elem_type = Scalar(PrimitiveType.Int32)
            elif arg.dtype == np.bool_:
                elem_type = Scalar(PrimitiveType.Bool)
            else:
                raise ValueError(f"Unsupported numpy dtype: {arg.dtype}")

            return Pointer(elem_type)
        elif isinstance(arg, str):
            # Explicitly reject strings - they are not supported
            raise ValueError(f"Unsupported argument type: {type(arg)}")
        else:
            # Check if it's a class instance
            if hasattr(arg, "__class__") and not isinstance(arg, type):
                # It's an instance of a class, return pointer to Structure
                return Pointer(Structure(arg.__class__.__name__))
            raise ValueError(f"Unsupported argument type: {type(arg)}")

    def _build_sdfg(self, arg_types, args, arg_shape_mapping, num_unique_shapes):
        sig = inspect.signature(self.func)

        # Handle return type
        return_type = Scalar(PrimitiveType.Void)
        infer_return_type = True
        if sig.return_annotation is not inspect.Signature.empty:
            return_type = _map_python_type(sig.return_annotation)
            infer_return_type = False
            if not isinstance(return_type, Type):
                raise ValueError(
                    f"Return type has invalid type annotation: {return_type}"
                )

        builder = StructuredSDFGBuilder(f"{self.name}_sdfg", return_type)

        # Register structure types for any class arguments
        # Also track member name to index mapping for each structure
        structures_to_register = {}
        structure_member_info = {}  # Maps struct_name -> {member_name: (index, type)}
        for i, (arg, dtype) in enumerate(zip(args, arg_types)):
            if isinstance(dtype, Pointer) and dtype.has_pointee_type():
                pointee = dtype.pointee_type
                if isinstance(pointee, Structure):
                    struct_name = pointee.name
                    if struct_name not in structures_to_register:
                        # Get class from arg to introspect members
                        if hasattr(arg, "__dict__"):
                            # Use __dict__ to get only instance attributes
                            # Sort by name to ensure consistent ordering
                            # Note: This alphabetical ordering is used to define the
                            # structure layout and must match the order expected by
                            # the backend code generation
                            member_types = []
                            member_names = []
                            for attr_name, attr_value in sorted(arg.__dict__.items()):
                                if not attr_name.startswith("_"):
                                    # Infer member type from instance attribute
                                    # Check bool before int since bool is subclass of int
                                    member_type = None
                                    if isinstance(attr_value, bool):
                                        member_type = Scalar(PrimitiveType.Bool)
                                    elif isinstance(attr_value, (int, np.int64)):
                                        member_type = Scalar(PrimitiveType.Int64)
                                    elif isinstance(attr_value, (float, np.float64)):
                                        member_type = Scalar(PrimitiveType.Double)
                                    elif isinstance(attr_value, np.int32):
                                        member_type = Scalar(PrimitiveType.Int32)
                                    elif isinstance(attr_value, np.float32):
                                        member_type = Scalar(PrimitiveType.Float)
                                    # TODO: Consider using np.integer and np.floating abstract types
                                    # for more comprehensive numpy type coverage
                                    # TODO: Add support for nested structures and arrays

                                    if member_type is not None:
                                        member_types.append(member_type)
                                        member_names.append(attr_name)

                            if member_types:
                                structures_to_register[struct_name] = member_types
                                # Build member name to (index, type) mapping
                                structure_member_info[struct_name] = {
                                    name: (idx, mtype)
                                    for idx, (name, mtype) in enumerate(
                                        zip(member_names, member_types)
                                    )
                                }

        # Store structure_member_info for later use in CompiledSDFG
        self._last_structure_member_info = structure_member_info

        # Register all discovered structures with the builder
        for struct_name, member_types in structures_to_register.items():
            builder.add_structure(struct_name, member_types)

        # Register arguments
        params = list(sig.parameters.items())
        if len(params) != len(arg_types):
            raise ValueError(
                f"Argument count mismatch: expected {len(params)}, got {len(arg_types)}"
            )

        array_info = {}

        # Add regular arguments
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
            builder.add_container(name, dtype, is_argument=True)

            # If it's an array, prepare shape info
            if isinstance(arg, np.ndarray):
                shapes = []
                for dim_idx in range(arg.ndim):
                    u_idx = arg_shape_mapping[(i, dim_idx)]
                    shapes.append(f"_s{u_idx}")

                array_info[name] = {"ndim": arg.ndim, "shapes": shapes}

        # Add unified shape arguments
        for i in range(num_unique_shapes):
            builder.add_container(
                f"_s{i}", Scalar(PrimitiveType.Int64), is_argument=True
            )

        # Create symbol table for parser
        symbol_table = {}
        for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):
            symbol_table[name] = dtype

        for i in range(num_unique_shapes):
            symbol_table[f"_s{i}"] = Scalar(PrimitiveType.Int64)

        # Parse AST
        source_lines, start_line = inspect.getsourcelines(self.func)
        source = textwrap.dedent("".join(source_lines))
        tree = ast.parse(source)
        ast.increment_lineno(tree, start_line - 1)
        func_def = tree.body[0]

        filename = inspect.getsourcefile(self.func)
        function_name = self.func.__name__

        parser = ASTParser(
            builder,
            array_info,
            symbol_table,
            filename,
            function_name,
            infer_return_type=infer_return_type,
            globals_dict=self.func.__globals__,
            structure_member_info=structure_member_info,
        )
        for node in func_def.body:
            parser.visit(node)

        sdfg = builder.move()
        return sdfg


def program(
    func=None,
    *,
    target="none",
    category="desktop",
    instrumentation_mode=None,
    capture_args=None,
):
    if func is None:
        return lambda f: DoccProgram(
            f,
            target=target,
            category=category,
            instrumentation_mode=instrumentation_mode,
            capture_args=capture_args,
        )
    return DoccProgram(
        func,
        target=target,
        category=category,
        instrumentation_mode=instrumentation_mode,
        capture_args=capture_args,
    )

1	import inspect	3✔
2	import shutil	3✔
3	import textwrap	3✔
4	import ast	3✔
5	import os	3✔
6	import getpass	3✔
7	import hashlib	3✔
8	import numpy as np	3✔
9	from typing import Annotated, get_origin, get_args	3✔
10	from ._sdfg import *	3✔
11	from .compiled_sdfg import CompiledSDFG	3✔
12	from .ast_parser import ASTParser	3✔
13
14
15	def _compile_wrapper(self, output_folder=None):	3✔
16	lib_path = self._compile(output_folder)	×
17	return CompiledSDFG(lib_path, self)	×
18
19
20	StructuredSDFG.compile = _compile_wrapper	3✔
21
22
23	def _map_python_type(dtype):	3✔
24	# If it is already a sdfg Type, return it
25	if isinstance(dtype, Type):	3✔
26	return dtype	×
27
28	# Handle Annotated for Arrays
29	if get_origin(dtype) is Annotated:	3✔
30	args = get_args(dtype)	3✔
31	base_type = args[0]	3✔
32	metadata = args[1:]	3✔
33
34	if base_type is np.ndarray:	3✔
35	# Convention: Annotated[np.ndarray, shape, dtype]
36	shape = metadata[0]	3✔
37	elem_type = Scalar(PrimitiveType.Double) # Default	3✔
38
39	if len(metadata) > 1:	3✔
40	possible_dtype = metadata[1]	3✔
41	elem_type = _map_python_type(possible_dtype)	3✔
42
43	return Pointer(elem_type)	3✔
44
45	# Handle numpy.ndarray[Shape, DType]
46	if get_origin(dtype) is np.ndarray:	3✔
47	args = get_args(dtype)	×
48	# args[0] is shape, args[1] is dtype
49	if len(args) >= 2:	×
50	elem_type = _map_python_type(args[1])	×
51	return Pointer(elem_type)	×
52
53	# Simple mapping for python types
54	if dtype is float or dtype is np.float64:	3✔
55	return Scalar(PrimitiveType.Double)	3✔
56	elif dtype is int or dtype is np.int64:	3✔
57	return Scalar(PrimitiveType.Int64)	3✔
58	elif dtype is bool or dtype is np.bool_:	3✔
59	return Scalar(PrimitiveType.Bool)	3✔
60	elif dtype is np.float32:	3✔
61	return Scalar(PrimitiveType.Float)	×
62	elif dtype is np.int32:	3✔
63	return Scalar(PrimitiveType.Int32)	×
64
65	# Handle Python classes - map to Structure type
66	if inspect.isclass(dtype):	3✔
67	# Use the class name as the structure name
68	return Pointer(Structure(dtype.__name__))	3✔
69
70	return dtype	×
71
72
73	class DoccProgram:	3✔
74	def __init__(	3✔
75	self,
76	func,
77	target="none",
78	category="desktop",
79	instrumentation_mode=None,
80	capture_args=None,
81	):
82	self.func = func	3✔
83	self.name = func.__name__	3✔
84	self.target = target	3✔
85	self.category = category	3✔
86	self.instrumentation_mode = instrumentation_mode	3✔
87	self.capture_args = capture_args	3✔
88	self.last_sdfg = None	3✔
89	self.cache = {}	3✔
90
91	def __call__(self, *args):	3✔
92	# JIT compile and run
93	compiled = self.compile(*args)	3✔
94	res = compiled(*args)	3✔
95
96	# Handle return value conversion based on annotation
97	sig = inspect.signature(self.func)	3✔
98	ret_annotation = sig.return_annotation	3✔
99
100	if ret_annotation is not inspect.Signature.empty:	3✔
101	if get_origin(ret_annotation) is Annotated:	3✔
102	type_args = get_args(ret_annotation)	3✔
103	if len(type_args) >= 1 and type_args[0] is np.ndarray:	3✔
104	shape = None	3✔
105	if len(type_args) >= 2:	3✔
106	shape = type_args[1]	3✔
107
108	if shape is not None:	3✔
109	try:	3✔
110	return np.ctypeslib.as_array(res, shape=shape)	3✔
111	except Exception:	×
112	pass	×
113
114	# Try to infer return shape from metadata
115	if hasattr(compiled, "get_return_shape"):	3✔
116	shape = compiled.get_return_shape(*args)	3✔
117	if shape is not None:	3✔
118	try:	3✔
119	return np.ctypeslib.as_array(res, shape=shape)	3✔
120	except Exception:	×
121	pass	×
122
123	return res	3✔
124
125	def compile(	3✔
126	self, *args, output_folder=None, instrumentation_mode=None, capture_args=None
127	):
128	original_output_folder = output_folder	3✔
129
130	# Resolve options
131	if instrumentation_mode is None:	3✔
132	instrumentation_mode = self.instrumentation_mode	3✔
133	if capture_args is None:	3✔
134	capture_args = self.capture_args	3✔
135
136	# Check environment variable DOCC_CI
137	docc_ci = os.environ.get("DOCC_CI", "")	3✔
138	if docc_ci:	3✔
139	if docc_ci == "regions":	×
140	if instrumentation_mode is None:	×
141	instrumentation_mode = "ols"	×
142	elif docc_ci == "arg-capture":	×
143	if capture_args is None:	×
144	capture_args = True	×
145	else:
146	# Full mode (or unknown value treated as full)
147	if instrumentation_mode is None:	×
148	instrumentation_mode = "ols"	×
149	if capture_args is None:	×
150	capture_args = True	×
151
152	# Defaults
153	if instrumentation_mode is None:	3✔
154	instrumentation_mode = ""	3✔
155	if capture_args is None:	3✔
156	capture_args = False	3✔
157
158	# 1. Analyze arguments and shapes
159	arg_types = []	3✔
160	shape_values = [] # List of unique shape values found	3✔
161	shape_sources = [] # List of (arg_idx, dim_idx) for each unique shape value	3✔
162
163	# Mapping from (arg_idx, dim_idx) -> unique_shape_idx
164	arg_shape_mapping = {}	3✔
165
166	for i, arg in enumerate(args):	3✔
167	t = self._infer_type(arg)	3✔
168	arg_types.append(t)	3✔
169
170	if isinstance(arg, np.ndarray):	3✔
171	for dim_idx, dim_val in enumerate(arg.shape):	3✔
172	# Check if we've seen this value
173	if dim_val in shape_values:	3✔
174	# Reuse
175	u_idx = shape_values.index(dim_val)	3✔
176	else:
177	# New
178	u_idx = len(shape_values)	3✔
179	shape_values.append(dim_val)	3✔
180	shape_sources.append((i, dim_idx))	3✔
181
182	arg_shape_mapping[(i, dim_idx)] = u_idx	3✔
183
184	# 2. Signature
185	mapping_sig = sorted(arg_shape_mapping.items())	3✔
186	type_sig = ", ".join(self._type_to_str(t) for t in arg_types)	3✔
187	signature = f"{type_sig}\|{mapping_sig}"	3✔
188
189	if output_folder is None:	3✔
190	filename = inspect.getsourcefile(self.func)	3✔
191	hash_input = f"{filename}\|{self.name}\|{self.target}\|{self.category}\|{self.capture_args}\|{self.instrumentation_mode}\|{signature}".encode(	3✔
192	"utf-8"
193	)
194	stable_id = hashlib.sha256(hash_input).hexdigest()[:16]	3✔
195
196	docc_tmp = os.environ.get("DOCC_TMP")	3✔
197	if docc_tmp:	3✔
198	output_folder = f"{docc_tmp}/{self.name}-{stable_id}"	×
199	else:
200	user = getpass.getuser()	3✔
201	output_folder = f"/tmp/{user}/DOCC/{self.name}-{stable_id}"	3✔
202
203	if original_output_folder is None and signature in self.cache:	3✔
204	return self.cache[signature]	3✔
205
206	# 3. Build SDFG
207	if os.path.exists(output_folder):	3✔
208	# Multiple python processes running the same code?
209	shutil.rmtree(output_folder)	×
210	sdfg = self._build_sdfg(arg_types, args, arg_shape_mapping, len(shape_values))	3✔
211	sdfg.expand()	3✔
212	sdfg.simplify()	3✔
213
214	if self.target != "none":	3✔
215	sdfg.normalize()	3✔
216
217	sdfg.dump(output_folder)	3✔
218
219	# Schedule if target is specified
220	if self.target != "none":	3✔
221	sdfg.schedule(self.target, self.category)	3✔
222
223	self.last_sdfg = sdfg	3✔
224
225	lib_path = sdfg._compile(	3✔
226	output_folder=output_folder,
227	instrumentation_mode=instrumentation_mode,
228	capture_args=capture_args,
229	)
230
231	# 5. Create CompiledSDFG
232	compiled = CompiledSDFG(	3✔
233	lib_path, sdfg, shape_sources, self._last_structure_member_info
234	)
235
236	# Cache if using default output folder
237	if original_output_folder is None:	3✔
238	self.cache[signature] = compiled	3✔
239
240	return compiled	3✔
241
242	def _get_signature(self, arg_types):	3✔
UNCOV 243	return ", ".join(self._type_to_str(t) for t in arg_types)	×
244
245	def _type_to_str(self, t):	3✔
246	if isinstance(t, Scalar):	3✔
247	return f"Scalar({t.primitive_type})"	3✔
248	elif isinstance(t, Array):	3✔
UNCOV 249	return f"Array({self._type_to_str(t.element_type)}, {t.num_elements})"	×
250	elif isinstance(t, Pointer):	3✔
251	return f"Pointer({self._type_to_str(t.pointee_type)})"	3✔
252	elif isinstance(t, Structure):	3✔
253	return f"Structure({t.name})"	3✔
UNCOV 254	return str(t)	×
255
256	def _infer_type(self, arg):	3✔
257	if isinstance(arg, (bool, np.bool_)):	3✔
258	return Scalar(PrimitiveType.Bool)	3✔
259	elif isinstance(arg, (int, np.int64)):	3✔
260	return Scalar(PrimitiveType.Int64)	3✔
261	elif isinstance(arg, (float, np.float64)):	3✔
262	return Scalar(PrimitiveType.Double)	3✔
263	elif isinstance(arg, np.int32):	3✔
264	return Scalar(PrimitiveType.Int32)	3✔
265	elif isinstance(arg, np.float32):	3✔
266	return Scalar(PrimitiveType.Float)	3✔
267	elif isinstance(arg, np.ndarray):	3✔
268	# Map dtype
269	if arg.dtype == np.float64:	3✔
270	elem_type = Scalar(PrimitiveType.Double)	3✔
271	elif arg.dtype == np.float32:	3✔
272	elem_type = Scalar(PrimitiveType.Float)	3✔
273	elif arg.dtype == np.int64:	3✔
274	elem_type = Scalar(PrimitiveType.Int64)	3✔
275	elif arg.dtype == np.int32:	3✔
276	elem_type = Scalar(PrimitiveType.Int32)	3✔
UNCOV 277	elif arg.dtype == np.bool_:	×
UNCOV 278	elem_type = Scalar(PrimitiveType.Bool)	×
279	else:
UNCOV 280	raise ValueError(f"Unsupported numpy dtype: {arg.dtype}")	×
281
282	return Pointer(elem_type)	3✔
283	elif isinstance(arg, str):	3✔
284	# Explicitly reject strings - they are not supported
285	raise ValueError(f"Unsupported argument type: {type(arg)}")	3✔
286	else:
287	# Check if it's a class instance
288	if hasattr(arg, "__class__") and not isinstance(arg, type):	3✔
289	# It's an instance of a class, return pointer to Structure
290	return Pointer(Structure(arg.__class__.__name__))	3✔
UNCOV 291	raise ValueError(f"Unsupported argument type: {type(arg)}")	×
292
293	def _build_sdfg(self, arg_types, args, arg_shape_mapping, num_unique_shapes):	3✔
294	sig = inspect.signature(self.func)	3✔
295
296	# Handle return type
297	return_type = Scalar(PrimitiveType.Void)	3✔
298	infer_return_type = True	3✔
299	if sig.return_annotation is not inspect.Signature.empty:	3✔
300	return_type = _map_python_type(sig.return_annotation)	3✔
301	infer_return_type = False	3✔
302	if not isinstance(return_type, Type):	3✔
303	raise ValueError(	×
304	f"Return type has invalid type annotation: {return_type}"
305	)
306
307	builder = StructuredSDFGBuilder(f"{self.name}_sdfg", return_type)	3✔
308
309	# Register structure types for any class arguments
310	# Also track member name to index mapping for each structure
311	structures_to_register = {}	3✔
312	structure_member_info = {} # Maps struct_name -> {member_name: (index, type)}	3✔
313	for i, (arg, dtype) in enumerate(zip(args, arg_types)):	3✔
314	if isinstance(dtype, Pointer) and dtype.has_pointee_type():	3✔
315	pointee = dtype.pointee_type	3✔
316	if isinstance(pointee, Structure):	3✔
317	struct_name = pointee.name	3✔
318	if struct_name not in structures_to_register:	3✔
319	# Get class from arg to introspect members
320	if hasattr(arg, "__dict__"):	3✔
321	# Use __dict__ to get only instance attributes
322	# Sort by name to ensure consistent ordering
323	# Note: This alphabetical ordering is used to define the
324	# structure layout and must match the order expected by
325	# the backend code generation
326	member_types = []	3✔
327	member_names = []	3✔
328	for attr_name, attr_value in sorted(arg.__dict__.items()):	3✔
329	if not attr_name.startswith("_"):	3✔
330	# Infer member type from instance attribute
331	# Check bool before int since bool is subclass of int
332	member_type = None	3✔
333	if isinstance(attr_value, bool):	3✔
334	member_type = Scalar(PrimitiveType.Bool)	×
335	elif isinstance(attr_value, (int, np.int64)):	3✔
336	member_type = Scalar(PrimitiveType.Int64)	×
337	elif isinstance(attr_value, (float, np.float64)):	3✔
338	member_type = Scalar(PrimitiveType.Double)	3✔
UNCOV 339	elif isinstance(attr_value, np.int32):	×
340	member_type = Scalar(PrimitiveType.Int32)	×
341	elif isinstance(attr_value, np.float32):	×
342	member_type = Scalar(PrimitiveType.Float)	×
343	# TODO: Consider using np.integer and np.floating abstract types
344	# for more comprehensive numpy type coverage
345	# TODO: Add support for nested structures and arrays
346
347	if member_type is not None:	3✔
348	member_types.append(member_type)	3✔
349	member_names.append(attr_name)	3✔
350
351	if member_types:	3✔
352	structures_to_register[struct_name] = member_types	3✔
353	# Build member name to (index, type) mapping
354	structure_member_info[struct_name] = {	3✔
355	name: (idx, mtype)
356	for idx, (name, mtype) in enumerate(
357	zip(member_names, member_types)
358	)
359	}
360
361	# Store structure_member_info for later use in CompiledSDFG
362	self._last_structure_member_info = structure_member_info	3✔
363
364	# Register all discovered structures with the builder
365	for struct_name, member_types in structures_to_register.items():	3✔
366	builder.add_structure(struct_name, member_types)	3✔
367
368	# Register arguments
369	params = list(sig.parameters.items())	3✔
370	if len(params) != len(arg_types):	3✔
UNCOV 371	raise ValueError(	×
372	f"Argument count mismatch: expected {len(params)}, got {len(arg_types)}"
373	)
374
375	array_info = {}	3✔
376
377	# Add regular arguments
378	for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):	3✔
379	builder.add_container(name, dtype, is_argument=True)	3✔
380
381	# If it's an array, prepare shape info
382	if isinstance(arg, np.ndarray):	3✔
383	shapes = []	3✔
384	for dim_idx in range(arg.ndim):	3✔
385	u_idx = arg_shape_mapping[(i, dim_idx)]	3✔
386	shapes.append(f"_s{u_idx}")	3✔
387
388	array_info[name] = {"ndim": arg.ndim, "shapes": shapes}	3✔
389
390	# Add unified shape arguments
391	for i in range(num_unique_shapes):	3✔
392	builder.add_container(	3✔
393	f"_s{i}", Scalar(PrimitiveType.Int64), is_argument=True
394	)
395
396	# Create symbol table for parser
397	symbol_table = {}	3✔
398	for i, ((name, param), dtype, arg) in enumerate(zip(params, arg_types, args)):	3✔
399	symbol_table[name] = dtype	3✔
400
401	for i in range(num_unique_shapes):	3✔
402	symbol_table[f"_s{i}"] = Scalar(PrimitiveType.Int64)	3✔
403
404	# Parse AST
405	source_lines, start_line = inspect.getsourcelines(self.func)	3✔
406	source = textwrap.dedent("".join(source_lines))	3✔
407	tree = ast.parse(source)	3✔
408	ast.increment_lineno(tree, start_line - 1)	3✔
409	func_def = tree.body[0]	3✔
410
411	filename = inspect.getsourcefile(self.func)	3✔
412	function_name = self.func.__name__	3✔
413
414	parser = ASTParser(	3✔
415	builder,
416	array_info,
417	symbol_table,
418	filename,
419	function_name,
420	infer_return_type=infer_return_type,
421	globals_dict=self.func.__globals__,
422	structure_member_info=structure_member_info,
423	)
424	for node in func_def.body:	3✔
425	parser.visit(node)	3✔
426
427	sdfg = builder.move()	3✔
428	return sdfg	3✔
429
430
431	def program(	3✔
432	func=None,
433	*,
434	target="none",
435	category="desktop",
436	instrumentation_mode=None,
437	capture_args=None,
438	):
439	if func is None:	3✔
440	return lambda f: DoccProgram(	3✔
441	f,
442	target=target,
443	category=category,
444	instrumentation_mode=instrumentation_mode,
445	capture_args=capture_args,
446	)
447	return DoccProgram(	3✔
448	func,
449	target=target,
450	category=category,
451	instrumentation_mode=instrumentation_mode,
452	capture_args=capture_args,
453	)

daisytuner / sdfglib / 21112243720

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