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

0.0

/python/docc/python/onnx_model_builder.py

"""
ONNX Model Builder

Converts JSON graph descriptions emitted by the ONNX tensor dispatchers
into valid ONNX model files that can be loaded by ONNX Runtime.
"""

import json
import os
from pathlib import Path
from typing import List, Dict, Any, Set, Tuple


def build_onnx_model(json_path: str, output_path: str) -> str:
    """
    Build an ONNX model from a JSON graph description.

    Args:
        json_path: Path to the JSON file containing node definitions
        output_path: Path where the .onnx model should be written

    Returns:
        Path to the generated ONNX model file
    """
    try:
        import onnx
        from onnx import helper, TensorProto
    except ImportError:
        raise ImportError(
            "The 'onnx' package is required for ONNX model generation. "
            "Install it with: pip install onnx"
        )

    # Read the JSON file
    with open(json_path, "r") as f:
        content = f.read().strip()

    # The JSON file contains comma-separated node definitions
    # Wrap in array brackets to make it valid JSON
    if content.endswith(","):
        content = content[:-1]  # Remove trailing comma
    json_content = f"[{content}]"

    try:
        nodes_data = json.loads(json_content)
    except json.JSONDecodeError as e:
        raise ValueError(
            f"Failed to parse ONNX graph JSON: {e}\nContent: {json_content[:500]}"
        )

    # Collect all inputs and outputs
    all_inputs: Set[str] = set()
    all_outputs: Set[str] = set()
    intermediate: Set[str] = set()
    elem_type = TensorProto.FLOAT  # Default

    # Build ONNX nodes
    onnx_nodes = []
    for node_data in nodes_data:
        op_type = node_data.get("op_type", "Unknown")
        name = node_data.get("name", "unnamed")
        inputs = node_data.get("inputs", [])
        outputs = node_data.get("outputs", [])
        attributes = node_data.get("attributes", {})

        if "elem_type" in node_data:
            elem_type = node_data["elem_type"]

        # Track inputs/outputs
        for inp in inputs:
            if inp not in intermediate:
                all_inputs.add(inp)
        for out in outputs:
            intermediate.add(out)
            all_outputs.add(out)

        # Build attribute list
        onnx_attrs = []
        for attr_name, attr_value in attributes.items():
            if isinstance(attr_value, list):
                if all(isinstance(x, int) for x in attr_value):
                    onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
                elif all(isinstance(x, float) for x in attr_value):
                    onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, int):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, float):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))
            elif isinstance(attr_value, str):
                onnx_attrs.append(helper.make_attribute(attr_name, attr_value))

        # Create the ONNX node
        node = helper.make_node(
            op_type,
            inputs=inputs,
            outputs=outputs,
            name=name,
        )
        # Add attributes
        node.attribute.extend(onnx_attrs)
        onnx_nodes.append(node)

    # Remove intermediate values from outputs (only keep final outputs)
    # Final outputs are those not consumed by any other node
    consumed = set()
    for node_data in nodes_data:
        for inp in node_data.get("inputs", []):
            consumed.add(inp)

    final_outputs = all_outputs - consumed
    graph_inputs = all_inputs - intermediate

    # If no final outputs, use all outputs
    if not final_outputs:
        final_outputs = all_outputs

    # Create graph inputs (with dynamic shapes using dim_param)
    input_tensors = []
    for inp_name in sorted(graph_inputs):
        # Use dynamic shape with symbolic dimension
        input_tensor = helper.make_tensor_value_info(
            inp_name, elem_type, None  # Dynamic shape
        )
        input_tensors.append(input_tensor)

    # Create graph outputs
    output_tensors = []
    for out_name in sorted(final_outputs):
        output_tensor = helper.make_tensor_value_info(
            out_name, elem_type, None  # Dynamic shape
        )
        output_tensors.append(output_tensor)

    # Create the graph
    graph = helper.make_graph(
        onnx_nodes,
        "docc_onnx_graph",
        input_tensors,
        output_tensors,
    )

    # Create the model
    model = helper.make_model(
        graph, opset_imports=[helper.make_opsetid("", 13)]  # ONNX opset 13
    )

    # Set IR version
    model.ir_version = 7

    # Validate the model
    try:
        onnx.checker.check_model(model)
    except onnx.checker.ValidationError as e:
        # Log warning but continue - dynamic shapes may cause validation issues
        print(f"ONNX validation warning: {e}")

    # Save the model
    onnx.save(model, output_path)

    return output_path


def convert_json_to_onnx(build_path: str) -> str:
    """
    Find and convert ONNX JSON files in the build directory.

    Args:
        build_path: Path to the build directory

    Returns:
        Path to the generated ONNX model, or None if no JSON found
    """
    json_paths = Path(build_path).glob("*.onnx.json")
    models = []
    for json_path in json_paths:
        onnx_path = json_path.parent / (json_path.name.replace(".onnx.json", ".onnx"))
        build_onnx_model(str(json_path), str(onnx_path))
        models.append(str(onnx_path))

    return models

1	"""
2	ONNX Model Builder
3
4	Converts JSON graph descriptions emitted by the ONNX tensor dispatchers
5	into valid ONNX model files that can be loaded by ONNX Runtime.
6	"""
7
UNCOV 8	import json	×
UNCOV 9	import os	×
UNCOV 10	from pathlib import Path	×
UNCOV 11	from typing import List, Dict, Any, Set, Tuple	×
12
13
UNCOV 14	def build_onnx_model(json_path: str, output_path: str) -> str:	×
15	"""
16	Build an ONNX model from a JSON graph description.
17
18	Args:
19	json_path: Path to the JSON file containing node definitions
20	output_path: Path where the .onnx model should be written
21
22	Returns:
23	Path to the generated ONNX model file
24	"""
UNCOV 25	try:	×
UNCOV 26	import onnx	×
UNCOV 27	from onnx import helper, TensorProto	×
28	except ImportError:	×
29	raise ImportError(	×
30	"The 'onnx' package is required for ONNX model generation. "
31	"Install it with: pip install onnx"
32	)
33
34	# Read the JSON file
UNCOV 35	with open(json_path, "r") as f:	×
UNCOV 36	content = f.read().strip()	×
37
38	# The JSON file contains comma-separated node definitions
39	# Wrap in array brackets to make it valid JSON
UNCOV 40	if content.endswith(","):	×
UNCOV 41	content = content[:-1] # Remove trailing comma	×
UNCOV 42	json_content = f"[{content}]"	×
43
UNCOV 44	try:	×
UNCOV 45	nodes_data = json.loads(json_content)	×
46	except json.JSONDecodeError as e:	×
47	raise ValueError(	×
48	f"Failed to parse ONNX graph JSON: {e}\nContent: {json_content[:500]}"
49	)
50
51	# Collect all inputs and outputs
UNCOV 52	all_inputs: Set[str] = set()	×
UNCOV 53	all_outputs: Set[str] = set()	×
UNCOV 54	intermediate: Set[str] = set()	×
UNCOV 55	elem_type = TensorProto.FLOAT # Default	×
56
57	# Build ONNX nodes
UNCOV 58	onnx_nodes = []	×
UNCOV 59	for node_data in nodes_data:	×
UNCOV 60	op_type = node_data.get("op_type", "Unknown")	×
UNCOV 61	name = node_data.get("name", "unnamed")	×
UNCOV 62	inputs = node_data.get("inputs", [])	×
UNCOV 63	outputs = node_data.get("outputs", [])	×
UNCOV 64	attributes = node_data.get("attributes", {})	×
65
UNCOV 66	if "elem_type" in node_data:	×
UNCOV 67	elem_type = node_data["elem_type"]	×
68
69	# Track inputs/outputs
UNCOV 70	for inp in inputs:	×
UNCOV 71	if inp not in intermediate:	×
UNCOV 72	all_inputs.add(inp)	×
UNCOV 73	for out in outputs:	×
UNCOV 74	intermediate.add(out)	×
UNCOV 75	all_outputs.add(out)	×
76
77	# Build attribute list
UNCOV 78	onnx_attrs = []	×
UNCOV 79	for attr_name, attr_value in attributes.items():	×
UNCOV 80	if isinstance(attr_value, list):	×
UNCOV 81	if all(isinstance(x, int) for x in attr_value):	×
UNCOV 82	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
83	elif all(isinstance(x, float) for x in attr_value):	×
84	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
UNCOV 85	elif isinstance(attr_value, int):	×
UNCOV 86	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
UNCOV 87	elif isinstance(attr_value, float):	×
UNCOV 88	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
89	elif isinstance(attr_value, str):	×
90	onnx_attrs.append(helper.make_attribute(attr_name, attr_value))	×
91
92	# Create the ONNX node
UNCOV 93	node = helper.make_node(	×
94	op_type,
95	inputs=inputs,
96	outputs=outputs,
97	name=name,
98	)
99	# Add attributes
UNCOV 100	node.attribute.extend(onnx_attrs)	×
UNCOV 101	onnx_nodes.append(node)	×
102
103	# Remove intermediate values from outputs (only keep final outputs)
104	# Final outputs are those not consumed by any other node
UNCOV 105	consumed = set()	×
UNCOV 106	for node_data in nodes_data:	×
UNCOV 107	for inp in node_data.get("inputs", []):	×
UNCOV 108	consumed.add(inp)	×
109
UNCOV 110	final_outputs = all_outputs - consumed	×
UNCOV 111	graph_inputs = all_inputs - intermediate	×
112
113	# If no final outputs, use all outputs
UNCOV 114	if not final_outputs:	×
115	final_outputs = all_outputs	×
116
117	# Create graph inputs (with dynamic shapes using dim_param)
UNCOV 118	input_tensors = []	×
UNCOV 119	for inp_name in sorted(graph_inputs):	×
120	# Use dynamic shape with symbolic dimension
UNCOV 121	input_tensor = helper.make_tensor_value_info(	×
122	inp_name, elem_type, None # Dynamic shape
123	)
UNCOV 124	input_tensors.append(input_tensor)	×
125
126	# Create graph outputs
UNCOV 127	output_tensors = []	×
UNCOV 128	for out_name in sorted(final_outputs):	×
UNCOV 129	output_tensor = helper.make_tensor_value_info(	×
130	out_name, elem_type, None # Dynamic shape
131	)
UNCOV 132	output_tensors.append(output_tensor)	×
133
134	# Create the graph
UNCOV 135	graph = helper.make_graph(	×
136	onnx_nodes,
137	"docc_onnx_graph",
138	input_tensors,
139	output_tensors,
140	)
141
142	# Create the model
UNCOV 143	model = helper.make_model(	×
144	graph, opset_imports=[helper.make_opsetid("", 13)] # ONNX opset 13
145	)
146
147	# Set IR version
UNCOV 148	model.ir_version = 7	×
149
150	# Validate the model
UNCOV 151	try:	×
UNCOV 152	onnx.checker.check_model(model)	×
UNCOV 153	except onnx.checker.ValidationError as e:	×
154	# Log warning but continue - dynamic shapes may cause validation issues
UNCOV 155	print(f"ONNX validation warning: {e}")	×
156
157	# Save the model
UNCOV 158	onnx.save(model, output_path)	×
159
UNCOV 160	return output_path	×
161
162
UNCOV 163	def convert_json_to_onnx(build_path: str) -> str:	×
164	"""
165	Find and convert ONNX JSON files in the build directory.
166
167	Args:
168	build_path: Path to the build directory
169
170	Returns:
171	Path to the generated ONNX model, or None if no JSON found
172	"""
UNCOV 173	json_paths = Path(build_path).glob("*.onnx.json")	×
UNCOV 174	models = []	×
UNCOV 175	for json_path in json_paths:	×
UNCOV 176	onnx_path = json_path.parent / (json_path.name.replace(".onnx.json", ".onnx"))	×
UNCOV 177	build_onnx_model(str(json_path), str(onnx_path))	×
UNCOV 178	models.append(str(onnx_path))	×
179
UNCOV 180	return models	×

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