Import symbolic shape information in MLIR while importing stateless graph #3655
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
patel-vimal
force-pushed
the
vimal/upstream/import_symbolic_shape_expr_for_stateless_graph_import
branch
2 times, most recently
from
db18918
to
5000283
Compare
New sympy type is introduced to represent integer infinity in upstream PyTorch repo. Subsequently, sympy.oo is no longer used to represent infinity upper bound for dynamic dimensions where the upper bound is unknown. Instead `int_oo` is used to represent integer infinity. This commit updates the `_sympy_int_to_int` utility in light of this change.
For the `stateless_fx_import` entry-point to import FX graph to MLIR, importing symbolic shape expressions was not enabled. This commit enabled this. Unlike the case of importing `ExportedProgram` here we need to get range constraints from the shape environment that was used during tracing. We rely on example inputs to get hold of ShapeEnv object which we later use to get the variable to range mapping.
This commit adds information about which symbol a given function argument corresponds to if any. This makes it possible to find the corresponding `torch.symbolic_int` operation that provides bounds on the range of a given function argument.
patel-vimal
force-pushed
the
vimal/upstream/import_symbolic_shape_expr_for_stateless_graph_import
branch
from
5000283
to
a10215c
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR contains changes required to enable the import of symbolic shape information on the stateless graph import entry point of FxImporter.
I note that the
FxImporter.import_stateless_graphAPI is deprecated. But, as per my knowledge, it's the only API that can be used withtorch.compile(). The other API, namelyimport_programandimportantly accepts torch's ExportedProgram format. However, thetorch.export.exportthat generatesExportedProgram` does not support graph breaks, which is a limiting factor for us. We are ready to maintain this API.