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

81.95

/python/docc/python/ast_utils.py

import ast
from docc.sdfg import DebugInfo


def is_negative_index(node):
    """Check if an AST node represents a negative constant index.

    Returns (True, abs_value) if the node is a negative constant,
    (False, None) otherwise.
    """
    # Handle -1, -2, etc. (UnaryOp with USub)
    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        if isinstance(node.operand, ast.Constant) and isinstance(
            node.operand.value, int
        ):
            return True, node.operand.value
    # Handle negative constants directly (rare but possible)
    if (
        isinstance(node, ast.Constant)
        and isinstance(node.value, int)
        and node.value < 0
    ):
        return True, -node.value
    return False, None


def normalize_negative_index(idx_node, dim_size_str):
    """Create an AST node that normalizes a negative index.

    If idx_node is negative, returns an AST Name node with expression
    "(dim_size - abs_value)". Otherwise returns the original node.
    """
    is_neg, abs_val = is_negative_index(idx_node)
    if is_neg:
        # Create: (dim_size - abs_value)
        return ast.Name(id=f"({dim_size_str} - {abs_val})", ctx=ast.Load())
    return idx_node


def contains_ufunc_outer(node):
    """Check if an AST node contains a ufunc outer call (e.g., np.add.outer).

    Returns (True, ufunc_name, outer_node) if found, (False, None, None) otherwise.
    """

    class UfuncOuterFinder(ast.NodeVisitor):
        def __init__(self):
            self.found = False
            self.ufunc_name = None
            self.outer_node = None

        def visit_Call(self, node):
            # Check for np.add.outer, np.subtract.outer, etc.
            if (
                isinstance(node.func, ast.Attribute)
                and node.func.attr == "outer"
                and isinstance(node.func.value, ast.Attribute)
                and isinstance(node.func.value.value, ast.Name)
                and node.func.value.value.id in ["numpy", "np"]
            ):
                self.found = True
                self.ufunc_name = node.func.value.attr
                self.outer_node = node
                return  # Don't visit children once found
            self.generic_visit(node)

    finder = UfuncOuterFinder()
    finder.visit(node)
    return finder.found, finder.ufunc_name, finder.outer_node


def get_debug_info(node, filename, function_name=""):
    if hasattr(node, "lineno"):
        return DebugInfo(
            filename,
            function_name,
            node.lineno,
            node.col_offset + 1,
            (
                node.end_lineno
                if hasattr(node, "end_lineno") and node.end_lineno is not None
                else node.lineno
            ),
            (
                node.end_col_offset + 1
                if hasattr(node, "end_col_offset") and node.end_col_offset is not None
                else node.col_offset + 1
            ),
        )
    return DebugInfo()


class ArrayToElementRewriter(ast.NodeTransformer):
    def __init__(self, loop_vars, tensor_table):
        self.loop_vars = loop_vars
        self.tensor_table = tensor_table

    def visit_Name(self, node):
        if node.id in self.tensor_table:
            # Replace with subscript
            indices = [ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars]
            return ast.Subscript(
                value=ast.Name(id=node.id, ctx=ast.Load()),
                slice=(
                    ast.Tuple(elts=indices, ctx=ast.Load())
                    if len(indices) > 1
                    else indices[0]
                ),
                ctx=ast.Load(),
            )
        return node


class SliceRewriter(ast.NodeTransformer):
    # Functions that operate on whole arrays, not element-wise
    # Their arguments should NOT be transformed to indexed access
    ARRAY_WISE_FUNCTIONS = {
        "flip",
        "fliplr",
        "flipud",  # Array reversal
        "dot",
        "matmul",
        "outer",
        "inner",  # Linear algebra
        "sum",
        "max",
        "min",
        "mean",
        "std",
        "var",
        "prod",  # Reductions
        "transpose",
        "reshape",
        "ravel",
        "flatten",  # Shape operations
        "sort",
        "argsort",
        "argmax",
        "argmin",  # Sorting/searching
        "cumsum",
        "cumprod",  # Cumulative operations
        "concatenate",
        "stack",
        "vstack",
        "hstack",  # Stacking
    }

    def __init__(self, loop_vars, tensor_table, expr_visitor):
        self.loop_vars = loop_vars
        self.tensor_table = tensor_table
        self.expr_visitor = expr_visitor

    def visit_Name(self, node):
        if node.id in self.tensor_table and self.loop_vars:
            ndim = len(self.tensor_table[node.id].shape)
            if ndim <= len(self.loop_vars) and ndim > 0:
                # For broadcasting: use the LAST ndim loop vars
                # e.g., for 1D array with 2 loop vars [i, j], use [j]
                indices = [
                    ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars[-ndim:]
                ]
                return ast.Subscript(
                    value=ast.Name(id=node.id, ctx=ast.Load()),
                    slice=(
                        ast.Tuple(elts=indices, ctx=ast.Load())
                        if len(indices) > 1
                        else indices[0]
                    ),
                    ctx=ast.Load(),
                )
        return node

    def visit_BinOp(self, node):
        if isinstance(node.op, ast.MatMult):
            if self.loop_vars:
                indices = [ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars]
                return ast.Subscript(
                    value=node,
                    slice=(
                        ast.Tuple(elts=indices, ctx=ast.Load())
                        if len(indices) > 1
                        else indices[0]
                    ),
                    ctx=ast.Load(),
                )
        return self.generic_visit(node)

    def visit_Call(self, node):
        """Handle function calls - don't transform arguments of array-wise functions."""
        func_name = None
        if isinstance(node.func, ast.Attribute):
            func_name = node.func.attr
        elif isinstance(node.func, ast.Name):
            func_name = node.func.id

        if func_name in self.ARRAY_WISE_FUNCTIONS:
            # Don't transform arguments - return node as-is
            # The function handler will process the slice correctly
            return node

        # For element-wise functions, transform arguments normally
        return self.generic_visit(node)

    def visit_Subscript(self, node):
        # First check if this subscript has any slice indices that need loop vars
        indices = []
        if isinstance(node.slice, ast.Tuple):
            indices = node.slice.elts
        else:
            indices = [node.slice]

        has_slice = any(isinstance(idx, ast.Slice) for idx in indices)

        # If no slices (all point indices), don't transform this subscript at all
        # The array is already fully indexed (e.g., _fict_[0])
        if not has_slice:
            return node

        # We need to visit the value to get its string representation for tensor_table lookup
        # But DO NOT transform it - the slices in THIS subscript will be replaced by loop vars
        value_str = self.expr_visitor.visit(node.value)
        if value_str not in self.tensor_table:
            return node

        ndim = len(self.tensor_table[value_str].shape)
        if len(indices) < ndim:
            indices = list(indices)
            for _ in range(ndim - len(indices)):
                indices.append(ast.Slice(lower=None, upper=None, step=None))

        new_indices = []
        slice_counter = 0

        for i, idx in enumerate(indices):
            if isinstance(idx, ast.Slice):
                if slice_counter >= len(self.loop_vars):
                    raise ValueError("Rank mismatch in slice assignment")

                loop_var = self.loop_vars[slice_counter]
                slice_counter += 1

                # Determine step value and check if negative
                step_str = "1"
                step_is_negative = False
                if idx.step:
                    step_str = self.expr_visitor.visit(idx.step)
                    # Check for negative step
                    if step_str.startswith("-") or step_str.startswith("(-"):
                        step_is_negative = True
                    elif isinstance(idx.step, ast.UnaryOp) and isinstance(
                        idx.step.op, ast.USub
                    ):
                        step_is_negative = True

                shapes = self.tensor_table[value_str].shape
                dim_size = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"

                # Compute start based on step direction
                if step_is_negative:
                    # Negative step: default start is (shape - 1)
                    if idx.lower:
                        start_str = self.expr_visitor.visit(idx.lower)
                        if start_str.startswith("-"):
                            start_str = f"({dim_size} {start_str})"
                    else:
                        start_str = f"({dim_size} - 1)"
                else:
                    # Positive step: default start is 0
                    start_str = "0"
                    if idx.lower:
                        start_str = self.expr_visitor.visit(idx.lower)
                        if start_str.startswith("-"):
                            start_str = f"({dim_size} {start_str})"

                if step_str == "1":
                    if start_str == "0":
                        term = loop_var
                    else:
                        term = f"({start_str} + {loop_var})"
                else:
                    term = f"({start_str} + {loop_var} * {step_str})"
                new_indices.append(ast.Name(id=term, ctx=ast.Load()))
            else:
                # Handle non-slice indices - need to normalize negative indices
                shapes = self.tensor_table[value_str].shape
                dim_size = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"
                normalized_idx = normalize_negative_index(idx, dim_size)
                new_indices.append(self.visit(normalized_idx))

        if len(new_indices) == 1:
            node.slice = new_indices[0]
        else:
            node.slice = ast.Tuple(elts=new_indices, ctx=ast.Load())

        return node


def get_unique_id(counter_ref):
    counter_ref[0] += 1
    return counter_ref[0]

1	import ast	4✔
2	from docc.sdfg import DebugInfo	4✔
3
4
5	def is_negative_index(node):	4✔
6	"""Check if an AST node represents a negative constant index.
7
8	Returns (True, abs_value) if the node is a negative constant,
9	(False, None) otherwise.
10	"""
11	# Handle -1, -2, etc. (UnaryOp with USub)
12	if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):	4✔
13	if isinstance(node.operand, ast.Constant) and isinstance(	4✔
14	node.operand.value, int
15	):
16	return True, node.operand.value	4✔
17	# Handle negative constants directly (rare but possible)
18	if (	4✔
19	isinstance(node, ast.Constant)
20	and isinstance(node.value, int)
21	and node.value < 0
22	):
23	return True, -node.value	×
24	return False, None	4✔
25
26
27	def normalize_negative_index(idx_node, dim_size_str):	4✔
28	"""Create an AST node that normalizes a negative index.
29
30	If idx_node is negative, returns an AST Name node with expression
31	"(dim_size - abs_value)". Otherwise returns the original node.
32	"""
33	is_neg, abs_val = is_negative_index(idx_node)	4✔
34	if is_neg:	4✔
35	# Create: (dim_size - abs_value)
36	return ast.Name(id=f"({dim_size_str} - {abs_val})", ctx=ast.Load())	4✔
37	return idx_node	4✔
38
39
40	def contains_ufunc_outer(node):	4✔
41	"""Check if an AST node contains a ufunc outer call (e.g., np.add.outer).
42
43	Returns (True, ufunc_name, outer_node) if found, (False, None, None) otherwise.
44	"""
45
46	class UfuncOuterFinder(ast.NodeVisitor):	4✔
47	def __init__(self):	4✔
48	self.found = False	4✔
49	self.ufunc_name = None	4✔
50	self.outer_node = None	4✔
51
52	def visit_Call(self, node):	4✔
53	# Check for np.add.outer, np.subtract.outer, etc.
54	if (	4✔
55	isinstance(node.func, ast.Attribute)
56	and node.func.attr == "outer"
57	and isinstance(node.func.value, ast.Attribute)
58	and isinstance(node.func.value.value, ast.Name)
59	and node.func.value.value.id in ["numpy", "np"]
60	):
61	self.found = True	4✔
62	self.ufunc_name = node.func.value.attr	4✔
63	self.outer_node = node	4✔
64	return # Don't visit children once found	4✔
65	self.generic_visit(node)	4✔
66
67	finder = UfuncOuterFinder()	4✔
68	finder.visit(node)	4✔
69	return finder.found, finder.ufunc_name, finder.outer_node	4✔
70
71
72	def get_debug_info(node, filename, function_name=""):	4✔
73	if hasattr(node, "lineno"):	4✔
74	return DebugInfo(	4✔
75	filename,
76	function_name,
77	node.lineno,
78	node.col_offset + 1,
79	(
80	node.end_lineno
81	if hasattr(node, "end_lineno") and node.end_lineno is not None
82	else node.lineno
83	),
84	(
85	node.end_col_offset + 1
86	if hasattr(node, "end_col_offset") and node.end_col_offset is not None
87	else node.col_offset + 1
88	),
89	)
90	return DebugInfo()	4✔
91
92
93	class ArrayToElementRewriter(ast.NodeTransformer):	4✔
94	def __init__(self, loop_vars, tensor_table):	4✔
95	self.loop_vars = loop_vars	×
NEW 96	self.tensor_table = tensor_table	×
97
98	def visit_Name(self, node):	4✔
NEW 99	if node.id in self.tensor_table:	×
100	# Replace with subscript
101	indices = [ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars]	×
102	return ast.Subscript(	×
103	value=ast.Name(id=node.id, ctx=ast.Load()),
104	slice=(
105	ast.Tuple(elts=indices, ctx=ast.Load())
106	if len(indices) > 1
107	else indices[0]
108	),
109	ctx=ast.Load(),
110	)
111	return node	×
112
113
114	class SliceRewriter(ast.NodeTransformer):	4✔
115	# Functions that operate on whole arrays, not element-wise
116	# Their arguments should NOT be transformed to indexed access
117	ARRAY_WISE_FUNCTIONS = {	4✔
118	"flip",
119	"fliplr",
120	"flipud", # Array reversal
121	"dot",
122	"matmul",
123	"outer",
124	"inner", # Linear algebra
125	"sum",
126	"max",
127	"min",
128	"mean",
129	"std",
130	"var",
131	"prod", # Reductions
132	"transpose",
133	"reshape",
134	"ravel",
135	"flatten", # Shape operations
136	"sort",
137	"argsort",
138	"argmax",
139	"argmin", # Sorting/searching
140	"cumsum",
141	"cumprod", # Cumulative operations
142	"concatenate",
143	"stack",
144	"vstack",
145	"hstack", # Stacking
146	}
147
148	def __init__(self, loop_vars, tensor_table, expr_visitor):	4✔
149	self.loop_vars = loop_vars	4✔
150	self.tensor_table = tensor_table	4✔
151	self.expr_visitor = expr_visitor	4✔
152
153	def visit_Name(self, node):	4✔
154	if node.id in self.tensor_table and self.loop_vars:	4✔
155	ndim = len(self.tensor_table[node.id].shape)	4✔
156	if ndim <= len(self.loop_vars) and ndim > 0:	4✔
157	# For broadcasting: use the LAST ndim loop vars
158	# e.g., for 1D array with 2 loop vars [i, j], use [j]
159	indices = [	4✔
160	ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars[-ndim:]
161	]
162	return ast.Subscript(	4✔
163	value=ast.Name(id=node.id, ctx=ast.Load()),
164	slice=(
165	ast.Tuple(elts=indices, ctx=ast.Load())
166	if len(indices) > 1
167	else indices[0]
168	),
169	ctx=ast.Load(),
170	)
171	return node	4✔
172
173	def visit_BinOp(self, node):	4✔
174	if isinstance(node.op, ast.MatMult):	4✔
175	if self.loop_vars:	×
176	indices = [ast.Name(id=lv, ctx=ast.Load()) for lv in self.loop_vars]	×
177	return ast.Subscript(	×
178	value=node,
179	slice=(
180	ast.Tuple(elts=indices, ctx=ast.Load())
181	if len(indices) > 1
182	else indices[0]
183	),
184	ctx=ast.Load(),
185	)
186	return self.generic_visit(node)	4✔
187
188	def visit_Call(self, node):	4✔
189	"""Handle function calls - don't transform arguments of array-wise functions."""
190	func_name = None	4✔
191	if isinstance(node.func, ast.Attribute):	4✔
NEW 192	func_name = node.func.attr	×
193	elif isinstance(node.func, ast.Name):	4✔
194	func_name = node.func.id	4✔
195
196	if func_name in self.ARRAY_WISE_FUNCTIONS:	4✔
197	# Don't transform arguments - return node as-is
198	# The function handler will process the slice correctly
NEW 199	return node	×
200
201	# For element-wise functions, transform arguments normally
202	return self.generic_visit(node)	4✔
203
204	def visit_Subscript(self, node):	4✔
205	# First check if this subscript has any slice indices that need loop vars
206	indices = []	4✔
207	if isinstance(node.slice, ast.Tuple):	4✔
208	indices = node.slice.elts	4✔
209	else:
210	indices = [node.slice]	4✔
211
212	has_slice = any(isinstance(idx, ast.Slice) for idx in indices)	4✔
213
214	# If no slices (all point indices), don't transform this subscript at all
215	# The array is already fully indexed (e.g., _fict_[0])
216	if not has_slice:	4✔
217	return node	4✔
218
219	# We need to visit the value to get its string representation for tensor_table lookup
220	# But DO NOT transform it - the slices in THIS subscript will be replaced by loop vars
221	value_str = self.expr_visitor.visit(node.value)	4✔
222	if value_str not in self.tensor_table:	4✔
223	return node	×
224
225	ndim = len(self.tensor_table[value_str].shape)	4✔
226	if len(indices) < ndim:	4✔
227	indices = list(indices)	×
228	for _ in range(ndim - len(indices)):	×
229	indices.append(ast.Slice(lower=None, upper=None, step=None))	×
230
231	new_indices = []	4✔
232	slice_counter = 0	4✔
233
234	for i, idx in enumerate(indices):	4✔
235	if isinstance(idx, ast.Slice):	4✔
236	if slice_counter >= len(self.loop_vars):	4✔
237	raise ValueError("Rank mismatch in slice assignment")	×
238
239	loop_var = self.loop_vars[slice_counter]	4✔
240	slice_counter += 1	4✔
241
242	# Determine step value and check if negative
243	step_str = "1"	4✔
244	step_is_negative = False	4✔
245	if idx.step:	4✔
246	step_str = self.expr_visitor.visit(idx.step)	4✔
247	# Check for negative step
248	if step_str.startswith("-") or step_str.startswith("(-"):	4✔
249	step_is_negative = True	4✔
250	elif isinstance(idx.step, ast.UnaryOp) and isinstance(	4✔
251	idx.step.op, ast.USub
252	):
NEW 253	step_is_negative = True	×
254
255	shapes = self.tensor_table[value_str].shape	4✔
256	dim_size = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"	4✔
257
258	# Compute start based on step direction
259	if step_is_negative:	4✔
260	# Negative step: default start is (shape - 1)
261	if idx.lower:	4✔
NEW 262	start_str = self.expr_visitor.visit(idx.lower)	×
NEW 263	if start_str.startswith("-"):	×
NEW 264	start_str = f"({dim_size} {start_str})"	×
265	else:
266	start_str = f"({dim_size} - 1)"	4✔
267	else:
268	# Positive step: default start is 0
269	start_str = "0"	4✔
270	if idx.lower:	4✔
271	start_str = self.expr_visitor.visit(idx.lower)	4✔
272	if start_str.startswith("-"):	4✔
NEW 273	start_str = f"({dim_size} {start_str})"	×
274
275	if step_str == "1":	4✔
276	if start_str == "0":	4✔
277	term = loop_var	4✔
278	else:
279	term = f"({start_str} + {loop_var})"	4✔
280	else:
281	term = f"({start_str} + {loop_var} * {step_str})"	4✔
282	new_indices.append(ast.Name(id=term, ctx=ast.Load()))	4✔
283	else:
284	# Handle non-slice indices - need to normalize negative indices
285	shapes = self.tensor_table[value_str].shape	4✔
286	dim_size = shapes[i] if i < len(shapes) else f"_{value_str}_shape_{i}"	4✔
287	normalized_idx = normalize_negative_index(idx, dim_size)	4✔
288	new_indices.append(self.visit(normalized_idx))	4✔
289
290	if len(new_indices) == 1:	4✔
291	node.slice = new_indices[0]	4✔
292	else:
293	node.slice = ast.Tuple(elts=new_indices, ctx=ast.Load())	4✔
294
295	return node	4✔
296
297
298	def get_unique_id(counter_ref):	4✔
299	counter_ref[0] += 1	×
300	return counter_ref[0]	×

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