[inductor] simplify expr when looking up size hint #123140

ColinPeppler · 2024-04-02T00:38:45Z

Context

Suppose we have two symbols: u0 and s0 where we know that u0 = s0. Now, let's say we tried to look up the size hint for u0 + 1.

Before this PR, we would use a fallback hint if one was provided.

pytorch/torch/_inductor/sizevars.py

Lines 406 to 407 in 3f6acf6

    
           def symbolic_hint(self, expr: Expr) -> Expr: 
        
               # Substitute all hints into expr, but leave unbacked symints alone

With this PR, we would try to replace u0 with s0 via simplify() before using a fallback hint.

pytorch/torch/_inductor/sizevars.py

Lines 46 to 47 in 3f6acf6

def simplify(self, expr: Expr):

return sympy.expand(expr).xreplace(self.replacements)

Concrete Example

A scenario where this is useful is when we're running autotuning benchmarking on bmm with two input nodes: one who has s0 as the batch size and one who has u0 as the batch size. During benchmarking, we'll create two example input tensors where the input with u0 has to use a fallback hint for batch size. This will lead to a mismatch.

pytorch/torch/_inductor/select_algorithm.py

Lines 991 to 997 in e3d80f2

    
           example_inputs_extern = [ 
        
               torch.as_strided( 
        
                   unique_example_inputs[input_node.get_name()], 
        
                   V.graph.sizevars.size_hints( 
        
                       input_node.get_size(), 
        
                       fallback=config.unbacked_symint_fallback, 
        
                   ),

Using the fallback hint (i.e. 8192) leads to a batch size mismatch.

# Note: s0 = 7 and u0 = 7 and fallback hint is 8192.
LoweringException: ErrorFromChoice: Expected size for first two dimensions of batch2 tensor to be: [7, 30] but got: [8192, 30].
From choice ExternKernelCaller(extern_kernels.bmm)

Differential Revision: D55619331

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @amjames @desertfire @chauhang

pytorch-bot · 2024-04-02T00:38:49Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/123140

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 35fcbbb with merge base f15fd65 ():

UNSTABLE - The following job failed but was likely due to flakiness present on trunk and has been marked as unstable:

pull / linux-jammy-py3.8-gcc11 / test (docs_test, 1, 1, linux.2xlarge, unstable) (gh)
Process completed with exit code 2.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

facebook-github-bot · 2024-04-02T00:38:56Z

This pull request was exported from Phabricator. Differential Revision: D55619331

Summary: Pull Request resolved: pytorch#123140 Test Plan: tbd Differential Revision: D55619331

facebook-github-bot · 2024-04-02T23:15:00Z

This pull request was exported from Phabricator. Differential Revision: D55619331

facebook-github-bot · 2024-04-03T00:01:33Z

This pull request was exported from Phabricator. Differential Revision: D55619331

ColinPeppler · 2024-04-03T00:05:59Z

Hi @peterbell10, should this change be okay? or does this look like a fix that should be happening in Dynamo?

aakhundov

Thanks for the fix!

aakhundov · 2024-04-03T00:24:10Z

test/inductor/test_unbacked_symints.py

Seems the bmm's shape is wrong in the comment? Should be [s0, 16, 32]?

woops forgot to update lol

aakhundov · 2024-04-03T00:24:53Z

test/inductor/test_unbacked_symints.py

Could you add a comment above this line mentioning that here s0 and u0 are unified?

aakhundov · 2024-04-03T00:29:47Z

or does this look like a fix that should be happening in Dynamo?

I think, it's fine to do the replacements inside size_hint? cc @ezyang.

facebook-github-bot · 2024-04-03T00:42:41Z

This pull request was exported from Phabricator. Differential Revision: D55619331

Summary: ## Context Suppose we have two symbols: `u0` and `s0` where we know that `u0 = s0`. Now, let's say we tried to look up the size hint for `u0 + 1`. * Before this PR, we would use a fallback hint if one was provided. https://github.com/pytorch/pytorch/blob/3f6acf65fd9b6094513cf28898a42b90dd1169a0/torch/_inductor/sizevars.py#L406-L407 * With this PR, we would try to replace `u0` with `s0` via `simplify()` before using a fallback hint. https://github.com/pytorch/pytorch/blob/3f6acf65fd9b6094513cf28898a42b90dd1169a0/torch/_inductor/sizevars.py#L46-L47 ## Concrete Example A scenario where this is useful is when we're running autotuning benchmarking on bmm with two input nodes: one who has `s0` as the batch size and one who has `u0` as the batch size. During benchmarking, we'll create two example input tensors where the input with `u0` has to use a fallback hint for batch size. This will lead to a mismatch. https://github.com/pytorch/pytorch/blob/e3d80f2fa98d7ab02f88023d381b2e5981dd99ff/torch/_inductor/select_algorithm.py#L991-L997 Using the fallback hint (i.e. 8192) leads to a batch size mismatch. ```python # Note: s0 = 7 and u0 = 7 and fallback hint is 8192. LoweringException: ErrorFromChoice: Expected size for first two dimensions of batch2 tensor to be: [7, 30] but got: [8192, 30]. From choice ExternKernelCaller(extern_kernels.bmm) ``` Test Plan: CI ``` $ CUDA_VISIBLE_DEVICES=0 python test/inductor/test_unbacked_symints.py -k test_equivalent_backed_unbacked_cuda ### Before ### File "torch/_inductor/select_algorithm.py", line 964, in __call__ timings = do_autotuning(precompile_fn) File "torch/_inductor/select_algorithm.py", line 911, in do_autotuning timings = self.lookup( File "torch/_inductor/codecache.py", line 306, in lookup raise e File "torch/_inductor/codecache.py", line 297, in lookup timings = benchmark(choices) File "torch/_inductor/select_algorithm.py", line 897, in autotune return make_benchmark_fn()(choices) File "torch/_inductor/select_algorithm.py", line 1068, in benchmark_in_current_process raise ErrorFromChoice(msg, choice, debug_str()) # noqa: TRY200 torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised: LoweringException: ErrorFromChoice: Expected size for first two dimensions of batch2 tensor to be: [7, 30] but got: [8192, 30]. From choice ExternKernelCaller(extern_kernels.bmm) inputs = [ torch.empty_strided((7, 30, 16), (480, 16, 1), dtype=torch.float32, device='cuda'), torch.empty_strided((30, 32), (32, 1), dtype=torch.float32, device='cuda'), ] ### After ### ---------------------------------------------------------------------- Ran 1 test in 4.627s OK ``` Reviewed By: tissue3, aakhundov Differential Revision: D55619331

facebook-github-bot · 2024-04-03T00:43:42Z

This pull request was exported from Phabricator. Differential Revision: D55619331

ColinPeppler · 2024-04-03T20:06:30Z

CI passes in OSS and internally. Also, we were worried that self.simplify(expr) will fail if expr is an int but sympy.expand(expr) can handle if expr is an int.

ColinPeppler · 2024-04-04T00:00:40Z

@pytorchbot merge

pytorchmergebot · 2024-04-04T00:02:39Z

Merge failed

Reason: This PR needs a release notes: label
If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Details for Dev Infra team

Raised by workflow job

peterbell10 · 2024-04-04T00:27:05Z

torch/_inductor/sizevars.py

        return sympy_subs(expr, self.var_to_val)

    def size_hint(self, expr: Expr, *, fallback: Optional[int] = None) -> int:
+        expr = self.simplify(expr)


This should be in symbolic_hint but otherwise this LGTM

@ColinPeppler this wasn't fixed

clee2000 · 2024-04-04T04:58:11Z

@pytorchbot merge -f "merged internally"

pytorchmergebot · 2024-04-04T04:59:52Z

Merge started

Your change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use -f as last resort and instead consider -i/--ignore-current to continue the merge ignoring current failures. This will allow currently pending tests to finish and report signal before the merge.

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

## Context Suppose we have two symbols: `u0` and `s0` where we know that `u0 = s0`. Now, let's say we tried to look up the size hint for `u0 + 1`. * Before this PR, we would use a fallback hint if one was provided. https://github.com/pytorch/pytorch/blob/3f6acf65fd9b6094513cf28898a42b90dd1169a0/torch/_inductor/sizevars.py#L406-L407 * With this PR, we would try to replace `u0` with `s0` via `simplify()` before using a fallback hint. https://github.com/pytorch/pytorch/blob/3f6acf65fd9b6094513cf28898a42b90dd1169a0/torch/_inductor/sizevars.py#L46-L47 ## Concrete Example A scenario where this is useful is when we're running autotuning benchmarking on bmm with two input nodes: one who has `s0` as the batch size and one who has `u0` as the batch size. During benchmarking, we'll create two example input tensors where the input with `u0` has to use a fallback hint for batch size. This will lead to a mismatch. https://github.com/pytorch/pytorch/blob/e3d80f2fa98d7ab02f88023d381b2e5981dd99ff/torch/_inductor/select_algorithm.py#L991-L997 Using the fallback hint (i.e. 8192) leads to a batch size mismatch. ``` # Note: s0 = 7 and u0 = 7 and fallback hint is 8192. LoweringException: ErrorFromChoice: Expected size for first two dimensions of batch2 tensor to be: [7, 30] but got: [8192, 30]. From choice ExternKernelCaller(extern_kernels.bmm) ``` Differential Revision: D55619331 Pull Request resolved: pytorch#123140 Approved by: https://github.com/aakhundov

pytorch-bot bot added ciflow/inductor module: inductor labels Apr 2, 2024

facebook-github-bot added the fb-exported label Apr 2, 2024

ColinPeppler mentioned this pull request Apr 2, 2024

[inductor] simplify expr before looking up hint #123141

Closed

ColinPeppler force-pushed the export-D55619331 branch from 9da2268 to 3f6acf6 Compare April 2, 2024 23:14

ColinPeppler added a commit to ColinPeppler/pytorch that referenced this pull request Apr 2, 2024

[inductor] simplify expr when looking up size hint (pytorch#123140)

3f6acf6

Summary: Pull Request resolved: pytorch#123140 Test Plan: tbd Differential Revision: D55619331

ColinPeppler changed the title ~~[inductor] simplify expr before looking up hint~~ [inductor] simplify expr when looking up size hint Apr 2, 2024

ColinPeppler force-pushed the export-D55619331 branch from 3f6acf6 to 4b0f72f Compare April 3, 2024 00:01

ColinPeppler requested review from aakhundov and peterbell10 April 3, 2024 00:04

aakhundov approved these changes Apr 3, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 3, 2024

ColinPeppler force-pushed the export-D55619331 branch from 4b0f72f to 0e3d87e Compare April 3, 2024 00:42

ColinPeppler force-pushed the export-D55619331 branch from 0e3d87e to 35fcbbb Compare April 3, 2024 00:43

pytorchmergebot added the merging label Apr 4, 2024

pytorchmergebot removed the merging label Apr 4, 2024

peterbell10 reviewed Apr 4, 2024

View reviewed changes

clee2000 added the release notes: inductor label Apr 4, 2024

pytorchmergebot added the merging label Apr 4, 2024

pytorchmergebot closed this in 2a24b54 Apr 4, 2024

pytorchmergebot added Merged and removed merging labels Apr 4, 2024

ColinPeppler added topic: not user facing topic category and removed release notes: inductor labels Apr 4, 2024

	def symbolic_hint(self, expr: Expr) -> Expr:
	# Substitute all hints into expr, but leave unbacked symints alone

	def simplify(self, expr: Expr):
	return sympy.expand(expr).xreplace(self.replacements)

	example_inputs_extern = [
	torch.as_strided(
	unique_example_inputs[input_node.get_name()],
	V.graph.sizevars.size_hints(
	input_node.get_size(),
	fallback=config.unbacked_symint_fallback,
	),

[inductor] simplify expr when looking up size hint #123140

[inductor] simplify expr when looking up size hint #123140

Uh oh!

Conversation

ColinPeppler commented Apr 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context

Concrete Example

Uh oh!

pytorch-bot bot commented Apr 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/123140

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

facebook-github-bot commented Apr 2, 2024

Uh oh!

facebook-github-bot commented Apr 2, 2024

Uh oh!

facebook-github-bot commented Apr 3, 2024

Uh oh!

ColinPeppler commented Apr 3, 2024

Uh oh!

aakhundov left a comment

Choose a reason for hiding this comment

Uh oh!

aakhundov Apr 3, 2024

Choose a reason for hiding this comment

Uh oh!

ColinPeppler Apr 3, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov Apr 3, 2024

Choose a reason for hiding this comment

Uh oh!

aakhundov commented Apr 3, 2024

Uh oh!

facebook-github-bot commented Apr 3, 2024

Uh oh!

facebook-github-bot commented Apr 3, 2024

Uh oh!

ColinPeppler commented Apr 3, 2024

Uh oh!

ColinPeppler commented Apr 4, 2024

Uh oh!

pytorchmergebot commented Apr 4, 2024

Merge failed

Uh oh!

peterbell10 Apr 4, 2024

Choose a reason for hiding this comment

Uh oh!

peterbell10 Apr 12, 2024

Choose a reason for hiding this comment

Uh oh!

clee2000 commented Apr 4, 2024

Uh oh!

pytorchmergebot commented Apr 4, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ColinPeppler commented Apr 2, 2024 •

edited

Loading

pytorch-bot bot commented Apr 2, 2024 •

edited

Loading