-
Notifications
You must be signed in to change notification settings - Fork 21.3k
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
fix soundness bug with unsupported constraints #102897
Closed
Closed
Changes from 7 commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
fe20907
fix soundness bug with unsupported constraints
avikchaudhuri ebb86f5
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 19a998c
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 10dde76
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 11ddeb3
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 3b528f1
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri d422007
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 1bbf112
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 0ff2384
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 3376754
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri 947ca5a
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri cc80e79
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri d8a72c8
Update on "fix soundness bug with unsupported constraints"
avikchaudhuri File filter
Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
There are no files selected for viewing
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
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
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
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
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -1499,7 +1499,7 @@ class DimConstraints: | |
Solutions are "static" values or simplified "dynamic" constraints. | ||
""" | ||
|
||
def __init__(self, symbol_to_source, var_to_val): | ||
def __init__(self, symbol_to_source, var_to_val, marked_dynamic): | ||
# We try to solve systems of inequalities with 1 free variable. | ||
self._univariate_inequalities: Dict[sympy.Symbol, Set[sympy.Expr]] = defaultdict(set) | ||
# Among them, we prioritize solving for a free variable that has equalities. | ||
|
@@ -1538,6 +1538,9 @@ def __init__(self, symbol_to_source, var_to_val): | |
# inconsistencies found on substituting with concrete values / static solutions | ||
self._inconsistencies: List[str] = [] | ||
|
||
# symbols that are marked dynamic | ||
self._marked_dynamic = marked_dynamic | ||
|
||
def rewrite_with_congruences(self, s, expr): | ||
""" | ||
Eliminate expressions of the form b // d and b % d while adding congruences of the form b % d == k. | ||
|
@@ -1686,7 +1689,35 @@ def raise_inconsistencies(self): | |
self._inconsistencies.clear() | ||
raise ValueError(f"The following inconsistencies were found:\n{msg}") | ||
|
||
def solve(self): | ||
def _force_specialization(self, s): | ||
val = self._var_to_val[s] | ||
self._static_results.add(f"{self._dcp.symbol_to_source[s][0].name()} == {val}") | ||
self._substitutions[s] = val | ||
|
||
def specialize_divisor_symbols(self): | ||
for expr in self._multivariate_inequalities: | ||
for atom in expr.atoms(FloorDiv, sympy.Mod): | ||
_, divisor = atom.args | ||
for s in divisor.free_symbols: | ||
self._force_specialization(s) | ||
|
||
multivariate_inequalities = self._multivariate_inequalities | ||
self._multivariate_inequalities = set() | ||
for expr in multivariate_inequalities: | ||
self.add(expr.subs(self._substitutions)) | ||
self.raise_inconsistencies() | ||
self._univariate_inequalities = { | ||
s: exprs | ||
for s, exprs in self._univariate_inequalities.items() | ||
if s not in self._substitutions | ||
} | ||
self._congruences = { | ||
s: congruences | ||
for s, congruences in self._congruences.items() | ||
if s not in self._substitutions | ||
} | ||
|
||
def solve(self, disable_congruences=True): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If and when we support specifying |
||
self.raise_inconsistencies() | ||
# as long as there are symbols with equalities, solve for them | ||
# NOTE(avik): this is guaranteed to terminate (#iterations <= #symbols) | ||
|
@@ -1711,11 +1742,20 @@ def solve(self): | |
self.add(expr.subs(s, self._substitutions[s])) | ||
self.raise_inconsistencies() | ||
|
||
# simplify symbolic equivalences: some of them will now become specializations! | ||
symbolic_equivalences = self._symbolic_equivalences | ||
self._symbolic_equivalences = [] | ||
for source, expr in symbolic_equivalences: | ||
self.add_equality(source, expr.subs(s, self._substitutions[s])) | ||
self.specialize_divisor_symbols() | ||
|
||
# solve linear congruences | ||
# NOTE(avik): We do not need to solve them for symbols that have already been specialized. | ||
reduced_congruences = self.reduce_congruences() | ||
for s, congruences in reduced_congruences.items(): | ||
for congruence in congruences: | ||
# any congruence that cannot be checked becomes a dynamic constraint as well | ||
if s not in self._substitutions or not sympy.checksol(congruence, {s: self._substitutions[s]}): | ||
if disable_congruences: | ||
self._force_specialization(s) | ||
self._univariate_inequalities.pop(s, None) | ||
else: | ||
self._dynamic_results.add(self._dcp.doprint(sympy.Eq(congruence, 0))) | ||
|
||
# remaining symbols have only pure inequalities (no equalities) | ||
for s, exprs in self._univariate_inequalities.items(): | ||
|
@@ -1732,18 +1772,25 @@ def solve(self): | |
for expr in exprs: | ||
self._dynamic_results.add(self._dcp.doprint(expr)) | ||
|
||
# simplify symbolic equivalences: some of them will now become specializations! | ||
symbolic_equivalences = self._symbolic_equivalences | ||
self._symbolic_equivalences = [] | ||
for source, expr in symbolic_equivalences: | ||
self.add_equality(source, expr.subs(self._substitutions)) | ||
|
||
# remaining symbolic equivalences become dynamic equality constraints | ||
for source, expr in self._symbolic_equivalences: | ||
self._dynamic_results.add(f"{self._dcp.print_source(source)} == {self._dcp.doprint(expr)}") | ||
|
||
# solve linear congruences | ||
# NOTE(avik): We do not need to solve them for symbols that have already been specialized. | ||
reduced_congruences = self.reduce_congruences() | ||
for s, congruences in reduced_congruences.items(): | ||
for congruence in congruences: | ||
# any congruence that cannot be checked becomes a dynamic constraint as well | ||
if s not in self._substitutions or not sympy.checksol(congruence, {s: self._substitutions[s]}): | ||
self._dynamic_results.add(self._dcp.doprint(sympy.Eq(congruence, 0))) | ||
def forced_specializations(self): | ||
return "\n".join([ | ||
( | ||
f"\t{self._dcp.symbol_to_source[s][0].name()}, which was marked dynamic, " | ||
f"must be specialized to {val}." | ||
) | ||
for s, val in self._substitutions.items() | ||
if s in self._marked_dynamic | ||
]) | ||
|
||
def prettify_results(self, original_signature: inspect.Signature): | ||
# Note: Model inputs are wrapped as LocalSource in dynamo. | ||
|
@@ -2366,7 +2413,11 @@ def hint(s): | |
# if we have an input (2, 3), we must show s0*2 == 2 and s1 == 3. | ||
# This does a lot of work: it covers duck sizing and equality guards. | ||
exprs = [] | ||
self.dim_constraints = DimConstraints(symbol_to_source, self.var_to_val) | ||
self.dim_constraints = DimConstraints( | ||
symbol_to_source, | ||
self.var_to_val, | ||
set(symbol_to_constraints.keys()), | ||
) | ||
|
||
if not _simplified: | ||
for source, expr in input_guards: | ||
|
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think you need to rebase.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I recently landed a change that makes this private method.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, I'll wait for https://www.internalfb.com/diff/D46471787 to land internally before rebasing.