30433648873

Committed 29 Jul 2026 07:58AM UTC coverage: 64.62% (+0.3%) from 64.324%

Build # 30433648873

Build Type

Pull #900

github

Committed by

web-flow

Commit Message

Merge 9ffae1a3f into f1a0f70d9

Pull Request Pull Request #900: [PyTorch] Add support for aten.slice

Coverage Stats

191 of 222 new or added lines in 3 files covered. (86.04%)

176 existing lines in 5 files now uncovered.

44062 of 68186 relevant lines covered (64.62%)

719.5 hits per line

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

84.85

/pytorch/docc/pytorch/graph_parser/normalization.py

"""
GraphParser modules for parsing normalization layers.
"""

import torch.fx

from docc.sdfg import StructuredSDFGBuilder, Tensor, Scalar, DebugInfo

from docc.pytorch.graph_parser.utils import (
    GraphParserModule,
    ContainerInfos,
    GraphParserError,
    register_module,
)


class BatchNormNoTrainingParser(GraphParserModule):
    def parse(
        self,
        node: torch.fx.Node,
        builder: StructuredSDFGBuilder,
        container_info: ContainerInfos,
    ) -> None:
        if len(node.args) != 7:
            raise GraphParserError(
                self,
                node,
                "Expected exactly 7 arguments but got " + str(len(node.args)),
            )
        if len(node.kwargs) != 0:
            raise GraphParserError(
                self, node, "Unsupported kwargs: " + str(node.kwargs)
            )
        input_container: str = self.get_arg_container(node, container_info, 0)
        input_tensor: Tensor = self.get_tensor_type(
            node, container_info, input_container
        )
        if node.args[1] is None:
            raise GraphParserError(
                self, node, "Currently the weight argument is required but got none"
            )
        weight_container: str = self.get_arg_container(node, container_info, 1)
        weight_tensor: Tensor = self.get_tensor_type(
            node, container_info, weight_container
        )
        if node.args[2] is None:
            raise GraphParserError(
                self, node, "Currently the bias argument is required but got none"
            )
        bias_container: str = self.get_arg_container(node, container_info, 2)
        bias_tensor: Tensor = self.get_tensor_type(node, container_info, bias_container)
        running_mean_container: str = self.get_arg_container(node, container_info, 3)
        running_mean_tensor: Tensor = self.get_tensor_type(
            node, container_info, running_mean_container
        )
        running_var_container: str = self.get_arg_container(node, container_info, 4)
        running_var_tensor: Tensor = self.get_tensor_type(
            node, container_info, running_var_container
        )
        # We just ignore momentum for now (argument 5)
        eps: str | tuple[str, Scalar] = self.get_arg_sdfg_value(node, container_info, 6)
        if isinstance(eps, str):
            eps_container: str = eps
            eps_type: Scalar = self.get_scalar_type(node, container_info, eps_container)
        else:
            eps_container: str = eps[0]
            eps_type: Scalar = self.align_constant_type(
                node, eps, input_tensor.element_type
            )
        result_containers: tuple[str, ...] = self.get_result_containers(
            3, node, builder, container_info
        )
        result_container: str = result_containers[0]
        result_tensor: Tensor = self.get_tensor_type(
            node, container_info, result_container
        )
        debug_info: DebugInfo = self.get_debug_info(node)
        builder.add_batchnorm_with_bias(
            input_container,
            input_tensor,
            running_var_container,
            running_var_tensor,
            running_mean_container,
            running_mean_tensor,
            weight_container,
            weight_tensor,
            bias_container,
            bias_tensor,
            eps_container,
            eps_type,
            result_container,
            result_tensor,
            debug_info,
        )


register_module(
    "aten._native_batch_norm_legit_no_training.default", BatchNormNoTrainingParser()
)


class LayerNormNoTrainingParser(GraphParserModule):
    def parse(
        self,
        node: torch.fx.Node,
        builder: StructuredSDFGBuilder,
        container_info: ContainerInfos,
    ) -> None:
        if len(node.args) != 5:
            raise GraphParserError(
                self,
                node,
                "Expected exactly 5 arguments but got " + str(len(node.args)),
            )
        if len(node.kwargs) != 0:
            raise GraphParserError(
                self, node, "Unsupported kwargs: " + str(node.kwargs)
            )
        input_container: str = self.get_arg_container(node, container_info, 0)
        input_tensor: Tensor = self.get_tensor_type(
            node, container_info, input_container
        )
        # normalized_shape is a list of ints; its length is the number of trailing
        # dimensions that are normalized over.
        num_normalized_dims: int = len(self.get_arg_multi_expr(node, 1))
        # weight (Gamma) is optional: elementwise_affine=False -> None
        if node.args[2] is None:
            weight_container: str = ""
            weight_tensor: Tensor = input_tensor
        else:
            weight_container = self.get_arg_container(node, container_info, 2)
            weight_tensor = self.get_tensor_type(node, container_info, weight_container)
        # bias (Beta) is optional: bias=False or elementwise_affine=False -> None
        if node.args[3] is None:
            bias_container: str = ""
            bias_tensor: Tensor = input_tensor
        else:
            bias_container = self.get_arg_container(node, container_info, 3)
            bias_tensor = self.get_tensor_type(node, container_info, bias_container)
        eps: str | tuple[str, Scalar] = self.get_arg_sdfg_value(node, container_info, 4)
        if isinstance(eps, str):
            eps_container: str = eps
            eps_type: Scalar = self.get_scalar_type(node, container_info, eps_container)
        else:
            eps_container: str = eps[0]
            eps_type: Scalar = self.align_constant_type(
                node, eps, input_tensor.element_type
            )
        # native_layer_norm returns a tuple of three results: (output, mean, rstd).
        # Only the first (the normalized output) is consumed here; the mean and rstd
        # result containers are requested to satisfy the metadata but left unwritten,
        # analogous to the BatchNorm parser above.
        result_containers: tuple[str, ...] = self.get_result_containers(
            3, node, builder, container_info
        )
        result_container: str = result_containers[0]
        result_tensor: Tensor = self.get_tensor_type(
            node, container_info, result_container
        )
        debug_info: DebugInfo = self.get_debug_info(node)
        builder.add_layernorm_with_bias(
            input_container,
            input_tensor,
            weight_container,
            weight_tensor,
            bias_container,
            bias_tensor,
            eps_container,
            eps_type,
            result_container,
            result_tensor,
            num_normalized_dims,
            debug_info,
        )


register_module("aten.native_layer_norm.default", LayerNormNoTrainingParser())

1	"""
2	GraphParser modules for parsing normalization layers.
3	"""
4
5	import torch.fx	1✔
6
7	from docc.sdfg import StructuredSDFGBuilder, Tensor, Scalar, DebugInfo	1✔
8
9	from docc.pytorch.graph_parser.utils import (	1✔
10	GraphParserModule,
11	ContainerInfos,
12	GraphParserError,
13	register_module,
14	)
15
16
17	class BatchNormNoTrainingParser(GraphParserModule):	1✔
18	def parse(	1✔
19	self,
20	node: torch.fx.Node,
21	builder: StructuredSDFGBuilder,
22	container_info: ContainerInfos,
23	) -> None:
24	if len(node.args) != 7:	1✔
25	raise GraphParserError(	×
26	self,
27	node,
28	"Expected exactly 7 arguments but got " + str(len(node.args)),
29	)
30	if len(node.kwargs) != 0:	1✔
31	raise GraphParserError(	×
32	self, node, "Unsupported kwargs: " + str(node.kwargs)
33	)
34	input_container: str = self.get_arg_container(node, container_info, 0)	1✔
35	input_tensor: Tensor = self.get_tensor_type(	1✔
36	node, container_info, input_container
37	)
38	if node.args[1] is None:	1✔
39	raise GraphParserError(	×
40	self, node, "Currently the weight argument is required but got none"
41	)
42	weight_container: str = self.get_arg_container(node, container_info, 1)	1✔
43	weight_tensor: Tensor = self.get_tensor_type(	1✔
44	node, container_info, weight_container
45	)
46	if node.args[2] is None:	1✔
47	raise GraphParserError(	×
48	self, node, "Currently the bias argument is required but got none"
49	)
50	bias_container: str = self.get_arg_container(node, container_info, 2)	1✔
51	bias_tensor: Tensor = self.get_tensor_type(node, container_info, bias_container)	1✔
52	running_mean_container: str = self.get_arg_container(node, container_info, 3)	1✔
53	running_mean_tensor: Tensor = self.get_tensor_type(	1✔
54	node, container_info, running_mean_container
55	)
56	running_var_container: str = self.get_arg_container(node, container_info, 4)	1✔
57	running_var_tensor: Tensor = self.get_tensor_type(	1✔
58	node, container_info, running_var_container
59	)
60	# We just ignore momentum for now (argument 5)
61	eps: str \| tuple[str, Scalar] = self.get_arg_sdfg_value(node, container_info, 6)	1✔
62	if isinstance(eps, str):	1✔
63	eps_container: str = eps	×
64	eps_type: Scalar = self.get_scalar_type(node, container_info, eps_container)	×
65	else:
66	eps_container: str = eps[0]	1✔
67	eps_type: Scalar = self.align_constant_type(	1✔
68	node, eps, input_tensor.element_type
69	)
70	result_containers: tuple[str, ...] = self.get_result_containers(	1✔
71	3, node, builder, container_info
72	)
73	result_container: str = result_containers[0]	1✔
74	result_tensor: Tensor = self.get_tensor_type(	1✔
75	node, container_info, result_container
76	)
77	debug_info: DebugInfo = self.get_debug_info(node)	1✔
78	builder.add_batchnorm_with_bias(	1✔
79	input_container,
80	input_tensor,
81	running_var_container,
82	running_var_tensor,
83	running_mean_container,
84	running_mean_tensor,
85	weight_container,
86	weight_tensor,
87	bias_container,
88	bias_tensor,
89	eps_container,
90	eps_type,
91	result_container,
92	result_tensor,
93	debug_info,
94	)
95
96
97	register_module(	1✔
98	"aten._native_batch_norm_legit_no_training.default", BatchNormNoTrainingParser()
99	)
100
101
102	class LayerNormNoTrainingParser(GraphParserModule):	1✔
103	def parse(	1✔
104	self,
105	node: torch.fx.Node,
106	builder: StructuredSDFGBuilder,
107	container_info: ContainerInfos,
108	) -> None:
109	if len(node.args) != 5:	1✔
UNCOV 110	raise GraphParserError(	×
111	self,
112	node,
113	"Expected exactly 5 arguments but got " + str(len(node.args)),
114	)
115	if len(node.kwargs) != 0:	1✔
UNCOV 116	raise GraphParserError(	×
117	self, node, "Unsupported kwargs: " + str(node.kwargs)
118	)
119	input_container: str = self.get_arg_container(node, container_info, 0)	1✔
120	input_tensor: Tensor = self.get_tensor_type(	1✔
121	node, container_info, input_container
122	)
123	# normalized_shape is a list of ints; its length is the number of trailing
124	# dimensions that are normalized over.
125	num_normalized_dims: int = len(self.get_arg_multi_expr(node, 1))	1✔
126	# weight (Gamma) is optional: elementwise_affine=False -> None
127	if node.args[2] is None:	1✔
128	weight_container: str = ""	1✔
129	weight_tensor: Tensor = input_tensor	1✔
130	else:
131	weight_container = self.get_arg_container(node, container_info, 2)	1✔
132	weight_tensor = self.get_tensor_type(node, container_info, weight_container)	1✔
133	# bias (Beta) is optional: bias=False or elementwise_affine=False -> None
134	if node.args[3] is None:	1✔
135	bias_container: str = ""	1✔
136	bias_tensor: Tensor = input_tensor	1✔
137	else:
138	bias_container = self.get_arg_container(node, container_info, 3)	1✔
139	bias_tensor = self.get_tensor_type(node, container_info, bias_container)	1✔
140	eps: str \| tuple[str, Scalar] = self.get_arg_sdfg_value(node, container_info, 4)	1✔
141	if isinstance(eps, str):	1✔
UNCOV 142	eps_container: str = eps	×
UNCOV 143	eps_type: Scalar = self.get_scalar_type(node, container_info, eps_container)	×
144	else:
145	eps_container: str = eps[0]	1✔
146	eps_type: Scalar = self.align_constant_type(	1✔
147	node, eps, input_tensor.element_type
148	)
149	# native_layer_norm returns a tuple of three results: (output, mean, rstd).
150	# Only the first (the normalized output) is consumed here; the mean and rstd
151	# result containers are requested to satisfy the metadata but left unwritten,
152	# analogous to the BatchNorm parser above.
153	result_containers: tuple[str, ...] = self.get_result_containers(	1✔
154	3, node, builder, container_info
155	)
156	result_container: str = result_containers[0]	1✔
157	result_tensor: Tensor = self.get_tensor_type(	1✔
158	node, container_info, result_container
159	)
160	debug_info: DebugInfo = self.get_debug_info(node)	1✔
161	builder.add_layernorm_with_bias(	1✔
162	input_container,
163	input_tensor,
164	weight_container,
165	weight_tensor,
166	bias_container,
167	bias_tensor,
168	eps_container,
169	eps_type,
170	result_container,
171	result_tensor,
172	num_normalized_dims,
173	debug_info,
174	)
175
176
177	register_module("aten.native_layer_norm.default", LayerNormNoTrainingParser())	1✔

daisytuner / docc / 30433648873

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