Skip to content

Add mixed dtype check for XNNPACK partitioner #9612

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mcr229 merged 5 commits into from
Mar 27, 2025
Merged

Add mixed dtype check for XNNPACK partitioner #9612

Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
46b708f
Add mixed dtype check for XNNPACK partitioner
keyprocedure Mar 25, 2025
6bd2196
Merge branch 'pytorch:main' into 9023_mixed-dtype-partitioner
keyprocedure Mar 25, 2025
2c10b33
Add why log for mixed input dtypes; remove mixed dtype check from output
keyprocedure Mar 26, 2025
2ddc16a
Improve mixed dtype logging with op name
keyprocedure Mar 27, 2025
bc7dba2
Merge branch 'main' into 9023_mixed-dtype-partitioner
keyprocedure Mar 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions backends/xnnpack/partition/config/xnnpack_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,9 +144,10 @@ def check_common_constraints(
return True

def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):
# Check inputs are valid dtypes
# Check inputs are valid and have the same dtypes
# Gather all args which are nodes
args_to_check = []
reference_dtype = None
for arg in node.args:
if isinstance(arg, list) or isinstance(arg, tuple):
for item in arg:
Expand Down Expand Up @@ -174,11 +175,19 @@ def _check_inputs_are_valid_dtypes(self, node, valid_dtypes):
if arg_val.dtype not in valid_dtypes:
return False

# Check for mixed dtypes
if reference_dtype is None:
reference_dtype = arg_val.dtype
elif arg_val.dtype != reference_dtype:
return False

return True

def _check_outputs_are_valid_dtypes(self, node, valid_dtypes):
# Check outputs are valid dtype
# Check outputs are valid and have the same dtypes
node_val = node.meta.get("val", None)
reference_dtype = None

if node_val is None:
return True

Expand All @@ -192,6 +201,12 @@ def _check_outputs_are_valid_dtypes(self, node, valid_dtypes):
if val.dtype not in valid_dtypes:
return False

# Check for mixed dtypes
if reference_dtype is None:
reference_dtype = val.dtype
elif val.dtype != reference_dtype:
return False

return True

def _check_node_has_valid_dtype(self, node):
Expand Down
Loading