Skip to content
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.

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

Add type inference #252

Merged
merged 2 commits into from
Nov 6, 2024
Merged

Add type inference #252

merged 2 commits into from
Nov 6, 2024

Conversation

harsh-nod
Copy link
Contributor

This PR adds a type inference pass to wave. Previously, the types were infered by looking up types from neighbors resulting in inefficient type inference.

Instead, we now introduce a pass that infers the types for all operators in the graph and the inferred type is then stores in the node. New nodes that are constructed in downstream passes are responsible for annotating types for the new operators.

@harsh-nod harsh-nod force-pushed the type_inference branch 2 times, most recently from e5846a5 to 1e98db2 Compare November 5, 2024 19:57
This PR adds a type inference pass to wave. Previously,
the types were infered by looking up types from neighbors
resulting in inefficient type inference.

Instead, we now introduce a pass that infers the types for
all operators in the graph and the inferred type is then
stored in the node.

Signed-off-by: Harsh Menon <[email protected]>
Infer the type of this operator using the types
of its arguments.
"""
pass
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe we can raise NotImplemented here?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question. So the idea is that most operators will implement infer_types except for Placeholders, etc. For those operators, since they already have a type, we want this to just pass through instead of throwing an error.

@@ -336,7 +336,7 @@ def _expand_reduction(
# Add GetResult nodes for the corresponding dimensions
reduction.graph.inserting_after(reduction.fx_node)
new_node = GetResult(reduction.fx_node, len(new_output_args))
new_node.add_to_graph(reduction.graph)
new_node.add_to_graph(reduction.graph, arg.type)
Copy link
Contributor

@raikonenfnu raikonenfnu Nov 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just curious what/why is the regular GetResult get_infer_type not working for this case? Can we add a comment to explain a bit.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure I can add that. The short summary is that it would require us to find the output node to get the right type since this is during expansion and the init_args have not been set yet. So since we have access to arg, we just use that instead.

Copy link
Contributor

@raikonenfnu raikonenfnu left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks awesome! thanks for the refactoring, this is really clean! just couple Qs

Signed-off-by: Harsh Menon <[email protected]>
@harsh-nod harsh-nod merged commit db1ec57 into iree-org:main Nov 6, 2024
6 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants