22168646772

Committed 18 Feb 2026 06:09PM UTC coverage: 64.742%. First build

Build # 22168646772

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #526 from daisytuner/native-ndarray

Python - Native Tensor Support: Update operations to use tensor type

Coverage Stats

2783 of 4104 new or added lines in 42 files covered. (67.81%)

23724 of 36644 relevant lines covered (64.74%)

368.07 hits per line

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

34.48

/python/docc/python/functions/scipy.py

import ast
from docc.sdfg import Scalar, PrimitiveType, Pointer, Tensor
from docc.python.ast_utils import get_debug_info


class SciPyHandler:
    """Handler for SciPy functions (scipy.special, scipy.signal, etc.)."""

    def __init__(self, expression_visitor):
        self._ev = expression_visitor
        # Nested structure: submodule -> {func_name -> handler}
        self.function_handlers = {
            "special": {
                "softmax": self._handle_softmax,
            },
            "signal": {
                "correlate2d": self._handle_correlate2d_expr,
            },
        }

    def has_handler(self, submodule, func_name):
        """Check if this handler can handle the given submodule.func_name."""
        return (
            submodule in self.function_handlers
            and func_name in self.function_handlers[submodule]
        )

    def handle_scipy_call(self, node, submodule, func_name):
        """Handle a call to a SciPy function."""
        if self.has_handler(submodule, func_name):
            return self.function_handlers[submodule][func_name](node, func_name)
        raise NotImplementedError(
            f"SciPy function scipy.{submodule}.{func_name} not supported"
        )

    # Expose parent properties for convenience
    @property
    def tensor_table(self):
        return self._ev.tensor_table

    @property
    def builder(self):
        return self._ev.builder

    @property
    def container_table(self):
        return self._ev.container_table

    def _get_unique_id(self):
        return self._ev._get_unique_id()

    def visit(self, node):
        return self._ev.visit(node)

    def _create_array_temp(self, shape, dtype):
        """Create a temporary array with the given shape and dtype."""
        return self._ev.numpy_visitor._create_array_temp(shape, dtype)

    # ========== scipy.special Functions ==========

    def _handle_softmax(self, node, func_name):
        """Handle scipy.special.softmax."""
        args = node.args
        keywords = {kw.arg: kw.value for kw in node.keywords}

        array_node = args[0]
        array_name = self.visit(array_node)

        if array_name not in self.tensor_table:
            raise ValueError(f"Softmax input must be an array, got {array_name}")

        input_tensor = self.tensor_table[array_name]
        input_shape = input_tensor.shape
        ndim = len(input_shape)

        axis = None
        if len(args) > 1:
            axis = args[1]
        elif "axis" in keywords:
            axis = keywords["axis"]

        axes = []
        if axis is None:
            axes = list(range(ndim))
        elif isinstance(axis, ast.Constant):
            val = axis.value
            if val < 0:
                val += ndim
            axes = [val]
        elif isinstance(axis, ast.Tuple):
            for elt in axis.elts:
                if isinstance(elt, ast.Constant):
                    val = elt.value
                    if val < 0:
                        val += ndim
                    axes.append(val)
        elif (
            isinstance(axis, ast.UnaryOp)
            and isinstance(axis.op, ast.USub)
            and isinstance(axis.operand, ast.Constant)
        ):
            val = -axis.operand.value
            if val < 0:
                val += ndim
            axes = [val]
        else:
            try:
                val = int(self.visit(axis))
                if val < 0:
                    val += ndim
                axes = [val]
            except:
                raise NotImplementedError("Dynamic axis not supported")

        dtype = Scalar(PrimitiveType.Double)

        tmp_name = self._create_array_temp(input_shape, dtype)

        output_tensor = self.tensor_table[tmp_name]
        self.builder.add_reduce_op(
            func_name, array_name, input_tensor, tmp_name, output_tensor, axes, False
        )

        return tmp_name

    # ========== scipy.signal Functions ==========

    def is_correlate2d(self, node):
        """Check if a node represents a scipy.signal.correlate2d call."""
        if not isinstance(node, ast.Call):
            return False

        if isinstance(node.func, ast.Attribute):
            if node.func.attr == "correlate2d":
                return True
        elif isinstance(node.func, ast.Name):
            if node.func.id == "correlate2d":
                return True

        return False

    def handle_correlate2d(self, target, value_node):
        """Handle scipy.signal.correlate2d (2D correlation/convolution).

        Args:
            target: The assignment target (ast.Name or string)
            value_node: The correlate2d call node

        Returns:
            True if handled successfully, False otherwise
        """
        if not self.is_correlate2d(value_node):
            return False

        args = value_node.args
        if len(args) < 2:
            return False

        in1_node = args[0]
        in2_node = args[1]

        in1_name = self.visit(in1_node)
        in2_name = self.visit(in2_node)

        if in1_name not in self.tensor_table:
            return False
        if in2_name not in self.tensor_table:
            return False

        in1_info = self.tensor_table[in1_name]
        in2_info = self.tensor_table[in2_name]

        # Check dimensions
        if len(in1_info.shape) != 2 or len(in2_info.shape) != 2:
            raise NotImplementedError(
                "Only 2D convolution is currently supported via scipy.signal mapping"
            )

        in1_shape = in1_info.shape
        in2_shape = in2_info.shape

        # Scipy Correlate2d / Convolve2d
        # Default mode is 'full', boundary 'fill', fillvalue 0

        mode = "full"
        # Parse kwargs
        for keyword in value_node.keywords:
            if keyword.arg == "mode" and isinstance(keyword.value, ast.Constant):
                mode = keyword.value.value

        # Also check positional args for mode
        if len(args) > 2 and isinstance(args[2], ast.Constant):
            mode = args[2].value

        if mode != "valid" and mode != "full" and mode != "same":
            raise NotImplementedError(f"Unsupported convolution mode: {mode}")

        # Map to ConvNode
        # Treat as N=1, C_in=1, C_out=1

        shape_strs = ["1", "1"] + [str(s) for s in in1_shape]
        kernel_shape_strs = [str(s) for s in in2_shape]

        # Default strides 1
        strides = ["1", "1"]
        dilations = ["1", "1"]
        group = "1"
        output_channels = "1"

        pads = ["0", "0", "0", "0"]

        if mode == "valid":
            pads = ["0", "0", "0", "0"]
        elif mode == "full":
            # Padding is kernel_size - 1 on both sides
            h_k = kernel_shape_strs[0]
            w_k = kernel_shape_strs[1]
            pad_h = f"({h_k} - 1)"
            pad_w = f"({w_k} - 1)"
            pads = [pad_h, pad_w, pad_h, pad_w]
        elif mode == "same":
            # Padding is kernel_size // 2
            h_k = kernel_shape_strs[0]
            w_k = kernel_shape_strs[1]
            pad_h = f"idiv({h_k}, 2)"
            pad_w = f"idiv({w_k}, 2)"
            pads = [pad_h, pad_w, pad_h, pad_w]

        target_name = ""
        if isinstance(target, ast.Name):
            target_name = target.id
        elif isinstance(target, str):
            target_name = target

        if not target_name:
            return False

        if self.builder.exists(target_name):
            # Ensure shape is inferred
            pass
        else:
            # Infer shape
            out_shape = []
            H1 = str(in1_shape[0])
            W1 = str(in1_shape[1])
            H2 = str(in2_shape[0])
            W2 = str(in2_shape[1])

            if mode == "valid":
                out_shape = [f"({H1} - {H2} + 1)", f"({W1} - {W2} + 1)"]
            elif mode == "same":
                out_shape = [H1, W1]
            elif mode == "full":
                out_shape = [f"({H1} + {H2} - 1)", f"({W1} + {W2} - 1)"]

            # Use Double type (float)
            dtype = Scalar(PrimitiveType.Double)
            ptr_type = Pointer(dtype)

            self.builder.add_container(target_name, ptr_type, False)

            # Update parser state
            self.container_table[target_name] = ptr_type
            self.tensor_table[target_name] = Tensor(dtype, out_shape)

            # Allocate memory for the result
            block_alloc = self.builder.add_block()

            # Calculate size: shape[0] * shape[1] * sizeof(double)
            # Assuming double (8 bytes)
            size_expr = f"(({out_shape[0]}) * ({out_shape[1]}))"
            total_size_expr = f"({size_expr} * 8)"

            t_malloc = self.builder.add_malloc(block_alloc, total_size_expr)
            t_ptr = self.builder.add_access(block_alloc, target_name)
            self.builder.add_memlet(
                block_alloc, t_malloc, "_ret", t_ptr, "void", "", ptr_type
            )

        debug_info = get_debug_info(
            value_node, getattr(self.builder, "filename", ""), ""
        )

        self.builder.add_conv(
            in1_name,
            in2_name,
            target_name,
            shape_strs,
            kernel_shape_strs,
            strides,
            pads,
            dilations,
            output_channels,
            group,
            debug_info,
        )
        return True

    def _handle_correlate2d_expr(self, node, func_name):
        """Handle scipy.signal.correlate2d as an expression (creates temp array).

        This wrapper is used when correlate2d appears in an expression context
        rather than a direct assignment.
        """
        # Create a temporary name for the result
        tmp_name = self.builder.find_new_name("_corr2d_")
        # Delegate to the main handler
        success = self.handle_correlate2d(tmp_name, node)
        if not success:
            raise NotImplementedError("Failed to handle correlate2d expression")
        return tmp_name

1	import ast	4✔
2	from docc.sdfg import Scalar, PrimitiveType, Pointer, Tensor	4✔
3	from docc.python.ast_utils import get_debug_info	4✔
4
5
6	class SciPyHandler:	4✔
7	"""Handler for SciPy functions (scipy.special, scipy.signal, etc.)."""
8
9	def __init__(self, expression_visitor):	4✔
10	self._ev = expression_visitor	4✔
11	# Nested structure: submodule -> {func_name -> handler}
12	self.function_handlers = {	4✔
13	"special": {
14	"softmax": self._handle_softmax,
15	},
16	"signal": {
17	"correlate2d": self._handle_correlate2d_expr,
18	},
19	}
20
21	def has_handler(self, submodule, func_name):	4✔
22	"""Check if this handler can handle the given submodule.func_name."""
23	return (	4✔
24	submodule in self.function_handlers
25	and func_name in self.function_handlers[submodule]
26	)
27
28	def handle_scipy_call(self, node, submodule, func_name):	4✔
29	"""Handle a call to a SciPy function."""
30	if self.has_handler(submodule, func_name):	4✔
31	return self.function_handlers[submodule][func_name](node, func_name)	4✔
NEW 32	raise NotImplementedError(	×
33	f"SciPy function scipy.{submodule}.{func_name} not supported"
34	)
35
36	# Expose parent properties for convenience
37	@property	4✔
38	def tensor_table(self):	4✔
39	return self._ev.tensor_table	4✔
40
41	@property	4✔
42	def builder(self):	4✔
43	return self._ev.builder	4✔
44
45	@property	4✔
46	def container_table(self):	4✔
NEW 47	return self._ev.container_table	×
48
49	def _get_unique_id(self):	4✔
NEW 50	return self._ev._get_unique_id()	×
51
52	def visit(self, node):	4✔
53	return self._ev.visit(node)	4✔
54
55	def _create_array_temp(self, shape, dtype):	4✔
56	"""Create a temporary array with the given shape and dtype."""
57	return self._ev.numpy_visitor._create_array_temp(shape, dtype)	4✔
58
59	# ========== scipy.special Functions ==========
60
61	def _handle_softmax(self, node, func_name):	4✔
62	"""Handle scipy.special.softmax."""
63	args = node.args	4✔
64	keywords = {kw.arg: kw.value for kw in node.keywords}	4✔
65
66	array_node = args[0]	4✔
67	array_name = self.visit(array_node)	4✔
68
69	if array_name not in self.tensor_table:	4✔
NEW 70	raise ValueError(f"Softmax input must be an array, got {array_name}")	×
71
72	input_tensor = self.tensor_table[array_name]	4✔
73	input_shape = input_tensor.shape	4✔
74	ndim = len(input_shape)	4✔
75
76	axis = None	4✔
77	if len(args) > 1:	4✔
NEW 78	axis = args[1]	×
79	elif "axis" in keywords:	4✔
80	axis = keywords["axis"]	4✔
81
82	axes = []	4✔
83	if axis is None:	4✔
84	axes = list(range(ndim))	4✔
85	elif isinstance(axis, ast.Constant):	4✔
86	val = axis.value	4✔
87	if val < 0:	4✔
NEW 88	val += ndim	×
89	axes = [val]	4✔
NEW 90	elif isinstance(axis, ast.Tuple):	×
NEW 91	for elt in axis.elts:	×
NEW 92	if isinstance(elt, ast.Constant):	×
NEW 93	val = elt.value	×
NEW 94	if val < 0:	×
NEW 95	val += ndim	×
NEW 96	axes.append(val)	×
NEW 97	elif (	×
98	isinstance(axis, ast.UnaryOp)
99	and isinstance(axis.op, ast.USub)
100	and isinstance(axis.operand, ast.Constant)
101	):
NEW 102	val = -axis.operand.value	×
NEW 103	if val < 0:	×
NEW 104	val += ndim	×
NEW 105	axes = [val]	×
106	else:
NEW 107	try:	×
NEW 108	val = int(self.visit(axis))	×
NEW 109	if val < 0:	×
NEW 110	val += ndim	×
NEW 111	axes = [val]	×
NEW 112	except:	×
NEW 113	raise NotImplementedError("Dynamic axis not supported")	×
114
115	dtype = Scalar(PrimitiveType.Double)	4✔
116
117	tmp_name = self._create_array_temp(input_shape, dtype)	4✔
118
119	output_tensor = self.tensor_table[tmp_name]	4✔
120	self.builder.add_reduce_op(	4✔
121	func_name, array_name, input_tensor, tmp_name, output_tensor, axes, False
122	)
123
124	return tmp_name	4✔
125
126	# ========== scipy.signal Functions ==========
127
128	def is_correlate2d(self, node):	4✔
129	"""Check if a node represents a scipy.signal.correlate2d call."""
130	if not isinstance(node, ast.Call):	4✔
131	return False	4✔
132
133	if isinstance(node.func, ast.Attribute):	4✔
134	if node.func.attr == "correlate2d":	4✔
NEW 135	return True	×
136	elif isinstance(node.func, ast.Name):	4✔
137	if node.func.id == "correlate2d":	4✔
NEW 138	return True	×
139
140	return False	4✔
141
142	def handle_correlate2d(self, target, value_node):	4✔
143	"""Handle scipy.signal.correlate2d (2D correlation/convolution).
144
145	Args:
146	target: The assignment target (ast.Name or string)
147	value_node: The correlate2d call node
148
149	Returns:
150	True if handled successfully, False otherwise
151	"""
NEW 152	if not self.is_correlate2d(value_node):	×
NEW 153	return False	×
154
NEW 155	args = value_node.args	×
NEW 156	if len(args) < 2:	×
NEW 157	return False	×
158
NEW 159	in1_node = args[0]	×
NEW 160	in2_node = args[1]	×
161
NEW 162	in1_name = self.visit(in1_node)	×
NEW 163	in2_name = self.visit(in2_node)	×
164
NEW 165	if in1_name not in self.tensor_table:	×
NEW 166	return False	×
NEW 167	if in2_name not in self.tensor_table:	×
NEW 168	return False	×
169
NEW 170	in1_info = self.tensor_table[in1_name]	×
NEW 171	in2_info = self.tensor_table[in2_name]	×
172
173	# Check dimensions
NEW 174	if len(in1_info.shape) != 2 or len(in2_info.shape) != 2:	×
NEW 175	raise NotImplementedError(	×
176	"Only 2D convolution is currently supported via scipy.signal mapping"
177	)
178
NEW 179	in1_shape = in1_info.shape	×
NEW 180	in2_shape = in2_info.shape	×
181
182	# Scipy Correlate2d / Convolve2d
183	# Default mode is 'full', boundary 'fill', fillvalue 0
184
NEW 185	mode = "full"	×
186	# Parse kwargs
NEW 187	for keyword in value_node.keywords:	×
NEW 188	if keyword.arg == "mode" and isinstance(keyword.value, ast.Constant):	×
NEW 189	mode = keyword.value.value	×
190
191	# Also check positional args for mode
NEW 192	if len(args) > 2 and isinstance(args[2], ast.Constant):	×
NEW 193	mode = args[2].value	×
194
NEW 195	if mode != "valid" and mode != "full" and mode != "same":	×
NEW 196	raise NotImplementedError(f"Unsupported convolution mode: {mode}")	×
197
198	# Map to ConvNode
199	# Treat as N=1, C_in=1, C_out=1
200
NEW 201	shape_strs = ["1", "1"] + [str(s) for s in in1_shape]	×
NEW 202	kernel_shape_strs = [str(s) for s in in2_shape]	×
203
204	# Default strides 1
NEW 205	strides = ["1", "1"]	×
NEW 206	dilations = ["1", "1"]	×
NEW 207	group = "1"	×
NEW 208	output_channels = "1"	×
209
NEW 210	pads = ["0", "0", "0", "0"]	×
211
NEW 212	if mode == "valid":	×
NEW 213	pads = ["0", "0", "0", "0"]	×
NEW 214	elif mode == "full":	×
215	# Padding is kernel_size - 1 on both sides
NEW 216	h_k = kernel_shape_strs[0]	×
NEW 217	w_k = kernel_shape_strs[1]	×
NEW 218	pad_h = f"({h_k} - 1)"	×
NEW 219	pad_w = f"({w_k} - 1)"	×
NEW 220	pads = [pad_h, pad_w, pad_h, pad_w]	×
NEW 221	elif mode == "same":	×
222	# Padding is kernel_size // 2
NEW 223	h_k = kernel_shape_strs[0]	×
NEW 224	w_k = kernel_shape_strs[1]	×
NEW 225	pad_h = f"idiv({h_k}, 2)"	×
NEW 226	pad_w = f"idiv({w_k}, 2)"	×
NEW 227	pads = [pad_h, pad_w, pad_h, pad_w]	×
228
NEW 229	target_name = ""	×
NEW 230	if isinstance(target, ast.Name):	×
NEW 231	target_name = target.id	×
NEW 232	elif isinstance(target, str):	×
NEW 233	target_name = target	×
234
NEW 235	if not target_name:	×
NEW 236	return False	×
237
NEW 238	if self.builder.exists(target_name):	×
239	# Ensure shape is inferred
NEW 240	pass	×
241	else:
242	# Infer shape
NEW 243	out_shape = []	×
NEW 244	H1 = str(in1_shape[0])	×
NEW 245	W1 = str(in1_shape[1])	×
NEW 246	H2 = str(in2_shape[0])	×
NEW 247	W2 = str(in2_shape[1])	×
248
NEW 249	if mode == "valid":	×
NEW 250	out_shape = [f"({H1} - {H2} + 1)", f"({W1} - {W2} + 1)"]	×
NEW 251	elif mode == "same":	×
NEW 252	out_shape = [H1, W1]	×
NEW 253	elif mode == "full":	×
NEW 254	out_shape = [f"({H1} + {H2} - 1)", f"({W1} + {W2} - 1)"]	×
255
256	# Use Double type (float)
NEW 257	dtype = Scalar(PrimitiveType.Double)	×
NEW 258	ptr_type = Pointer(dtype)	×
259
NEW 260	self.builder.add_container(target_name, ptr_type, False)	×
261
262	# Update parser state
NEW 263	self.container_table[target_name] = ptr_type	×
NEW 264	self.tensor_table[target_name] = Tensor(dtype, out_shape)	×
265
266	# Allocate memory for the result
NEW 267	block_alloc = self.builder.add_block()	×
268
269	# Calculate size: shape[0] * shape[1] * sizeof(double)
270	# Assuming double (8 bytes)
NEW 271	size_expr = f"(({out_shape[0]}) * ({out_shape[1]}))"	×
NEW 272	total_size_expr = f"({size_expr} * 8)"	×
273
NEW 274	t_malloc = self.builder.add_malloc(block_alloc, total_size_expr)	×
NEW 275	t_ptr = self.builder.add_access(block_alloc, target_name)	×
NEW 276	self.builder.add_memlet(	×
277	block_alloc, t_malloc, "_ret", t_ptr, "void", "", ptr_type
278	)
279
NEW 280	debug_info = get_debug_info(	×
281	value_node, getattr(self.builder, "filename", ""), ""
282	)
283
NEW 284	self.builder.add_conv(	×
285	in1_name,
286	in2_name,
287	target_name,
288	shape_strs,
289	kernel_shape_strs,
290	strides,
291	pads,
292	dilations,
293	output_channels,
294	group,
295	debug_info,
296	)
NEW 297	return True	×
298
299	def _handle_correlate2d_expr(self, node, func_name):	4✔
300	"""Handle scipy.signal.correlate2d as an expression (creates temp array).
301
302	This wrapper is used when correlate2d appears in an expression context
303	rather than a direct assignment.
304	"""
305	# Create a temporary name for the result
NEW 306	tmp_name = self.builder.find_new_name("_corr2d_")	×
307	# Delegate to the main handler
NEW 308	success = self.handle_correlate2d(tmp_name, node)	×
NEW 309	if not success:	×
NEW 310	raise NotImplementedError("Failed to handle correlate2d expression")	×
NEW 311	return tmp_name	×

daisytuner / docc / 22168646772

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