22020750556

Committed 14 Feb 2026 04:38PM UTC coverage: 64.828% (-1.5%) from 66.315%

Build # 22020750556

Build Type

Pull #524

github

Committed by

web-flow

Commit Message

Merge 2784aa264 into 9d01cacd5

Pull Request Pull Request #524: Native Tensor Support - Step 2: Use tensor types on memlets of tensor nodes

Run Details

245 of 570 new or added lines in 24 files covered. (42.98%)

458 existing lines in 18 files now uncovered.

23080 of 35602 relevant lines covered (64.83%)

371.57 hits per line

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

18.27

/python/docc/python/convolution.py

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


class ConvolutionHandler:
    def __init__(self, builder, array_info, symbol_table, expr_visitor):
        self.builder = builder
        self.array_info = array_info
        self.symbol_table = symbol_table
        self.expr_visitor = expr_visitor

    def _parse_expr(self, node):
        return self.expr_visitor.visit(node)

    def is_conv(self, node):
        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_conv(self, target, value_node):
        if not self.is_conv(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._parse_expr(in1_node)
        in2_name = self._parse_expr(in2_node)

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

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

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

        in1_shape = in1_info["shapes"]
        in2_shape = in2_info["shapes"]

        # 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
            # shapes are symbolic strings, so we construct the padding string
            # This is tricky without a symbolic engine in Python.
            # But we can produce a string expression that SDFG builder parses.
            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.symbol_table[target_name] = ptr_type
            self.array_info[target_name] = {"ndim": 2, "shapes": 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", ""), ""
        )
        # Note: filename might not be accessible easily if not passed to handler.
        # But ASTParser passes filename to debug info helpers usually.
        # We'll pass a generic debug info if needed or modify init.

        # wait, ASTParser initializes Handler.

        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

1	import ast	4✔
2	from docc.sdfg import Scalar, PrimitiveType, Pointer	4✔
3	from docc.python.ast_utils import get_debug_info	4✔
4
5
6	class ConvolutionHandler:	4✔
7	def __init__(self, builder, array_info, symbol_table, expr_visitor):	4✔
8	self.builder = builder	4✔
9	self.array_info = array_info	4✔
10	self.symbol_table = symbol_table	4✔
11	self.expr_visitor = expr_visitor	4✔
12
13	def _parse_expr(self, node):	4✔
UNCOV 14	return self.expr_visitor.visit(node)	×
15
16	def is_conv(self, node):	4✔
17	if not isinstance(node, ast.Call):	4✔
18	return False	4✔
19
20	if isinstance(node.func, ast.Attribute):	4✔
21	if node.func.attr == "correlate2d":	4✔
UNCOV 22	return True	×
23	elif isinstance(node.func, ast.Name):	4✔
24	if node.func.id == "correlate2d":	4✔
25	return True	×
26
27	return False	4✔
28
29	def handle_conv(self, target, value_node):	4✔
UNCOV 30	if not self.is_conv(value_node):	×
31	return False	×
32
UNCOV 33	args = value_node.args	×
UNCOV 34	if len(args) < 2:	×
35	return False	×
36
UNCOV 37	in1_node = args[0]	×
UNCOV 38	in2_node = args[1]	×
39
UNCOV 40	in1_name = self._parse_expr(in1_node)	×
UNCOV 41	in2_name = self._parse_expr(in2_node)	×
42
UNCOV 43	if in1_name not in self.array_info:	×
44	return False	×
UNCOV 45	if in2_name not in self.array_info:	×
46	return False	×
47
UNCOV 48	in1_info = self.array_info[in1_name]	×
UNCOV 49	in2_info = self.array_info[in2_name]	×
50
51	# Check dimensions
UNCOV 52	if in1_info["ndim"] != 2 or in2_info["ndim"] != 2:	×
53	raise NotImplementedError(	×
54	"Only 2D convolution is currently supported via scipy.signal mapping"
55	)
56
UNCOV 57	in1_shape = in1_info["shapes"]	×
UNCOV 58	in2_shape = in2_info["shapes"]	×
59
60	# Scipy Correlate2d / Convolve2d
61	# Default mode is 'full', boundary 'fill', fillvalue 0
62
UNCOV 63	mode = "full"	×
64	# Parse kwargs
UNCOV 65	for keyword in value_node.keywords:	×
UNCOV 66	if keyword.arg == "mode" and isinstance(keyword.value, ast.Constant):	×
UNCOV 67	mode = keyword.value.value	×
68
69	# Also check positional args for mode
UNCOV 70	if len(args) > 2 and isinstance(args[2], ast.Constant):	×
71	mode = args[2].value	×
72
UNCOV 73	if mode != "valid" and mode != "full" and mode != "same":	×
74	raise NotImplementedError(f"Unsupported convolution mode: {mode}")	×
75
76	# Map to ConvNode
77	# Treat as N=1, C_in=1, C_out=1
78
UNCOV 79	shape_strs = ["1", "1"] + [str(s) for s in in1_shape]	×
UNCOV 80	kernel_shape_strs = [str(s) for s in in2_shape]	×
81
82	# Default strides 1
UNCOV 83	strides = ["1", "1"]	×
UNCOV 84	dilations = ["1", "1"]	×
UNCOV 85	group = "1"	×
UNCOV 86	output_channels = "1"	×
87
UNCOV 88	pads = ["0", "0", "0", "0"]	×
89
UNCOV 90	if mode == "valid":	×
UNCOV 91	pads = ["0", "0", "0", "0"]	×
UNCOV 92	elif mode == "full":	×
93	# Padding is kernel_size - 1 on both sides
94	# shapes are symbolic strings, so we construct the padding string
95	# This is tricky without a symbolic engine in Python.
96	# But we can produce a string expression that SDFG builder parses.
97	h_k = kernel_shape_strs[0]	×
98	w_k = kernel_shape_strs[1]	×
99	pad_h = f"({h_k} - 1)"	×
100	pad_w = f"({w_k} - 1)"	×
101	pads = [pad_h, pad_w, pad_h, pad_w]	×
UNCOV 102	elif mode == "same":	×
103	# Padding is kernel_size // 2
UNCOV 104	h_k = kernel_shape_strs[0]	×
UNCOV 105	w_k = kernel_shape_strs[1]	×
UNCOV 106	pad_h = f"idiv({h_k}, 2)"	×
UNCOV 107	pad_w = f"idiv({w_k}, 2)"	×
UNCOV 108	pads = [pad_h, pad_w, pad_h, pad_w]	×
109
UNCOV 110	target_name = ""	×
UNCOV 111	if isinstance(target, ast.Name):	×
UNCOV 112	target_name = target.id	×
113	elif isinstance(target, str):	×
114	target_name = target	×
115
UNCOV 116	if not target_name:	×
117	return False	×
118
UNCOV 119	if self.builder.exists(target_name):	×
120	# Ensure shape is inferred
121	pass	×
122	else:
123	# Infer shape
UNCOV 124	out_shape = []	×
UNCOV 125	H1 = str(in1_shape[0])	×
UNCOV 126	W1 = str(in1_shape[1])	×
UNCOV 127	H2 = str(in2_shape[0])	×
UNCOV 128	W2 = str(in2_shape[1])	×
129
UNCOV 130	if mode == "valid":	×
UNCOV 131	out_shape = [f"({H1} - {H2} + 1)", f"({W1} - {W2} + 1)"]	×
UNCOV 132	elif mode == "same":	×
UNCOV 133	out_shape = [H1, W1]	×
134	elif mode == "full":	×
135	out_shape = [f"({H1} + {H2} - 1)", f"({W1} + {W2} - 1)"]	×
136
137	# Use Double type (float)
UNCOV 138	dtype = Scalar(PrimitiveType.Double)	×
UNCOV 139	ptr_type = Pointer(dtype)	×
140
UNCOV 141	self.builder.add_container(target_name, ptr_type, False)	×
142
143	# Update parser state
UNCOV 144	self.symbol_table[target_name] = ptr_type	×
UNCOV 145	self.array_info[target_name] = {"ndim": 2, "shapes": out_shape}	×
146
147	# Allocate memory for the result
UNCOV 148	block_alloc = self.builder.add_block()	×
149
150	# Calculate size: shape[0] * shape[1] * sizeof(double)
151	# Assuming double (8 bytes)
UNCOV 152	size_expr = f"(({out_shape[0]}) * ({out_shape[1]}))"	×
UNCOV 153	total_size_expr = f"({size_expr} * 8)"	×
154
UNCOV 155	t_malloc = self.builder.add_malloc(block_alloc, total_size_expr)	×
UNCOV 156	t_ptr = self.builder.add_access(block_alloc, target_name)	×
UNCOV 157	self.builder.add_memlet(	×
158	block_alloc, t_malloc, "_ret", t_ptr, "void", "", ptr_type
159	)
160
UNCOV 161	debug_info = get_debug_info(	×
162	value_node, getattr(self.builder, "filename", ""), ""
163	)
164	# Note: filename might not be accessible easily if not passed to handler.
165	# But ASTParser passes filename to debug info helpers usually.
166	# We'll pass a generic debug info if needed or modify init.
167
168	# wait, ASTParser initializes Handler.
169
UNCOV 170	self.builder.add_conv(	×
171	in1_name,
172	in2_name,
173	target_name,
174	shape_strs,
175	kernel_shape_strs,
176	strides,
177	pads,
178	dilations,
179	output_channels,
180	group,
181	debug_info,
182	)
UNCOV 183	return True	×

daisytuner / docc / 22020750556

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