Skip to content

[functorch] introduce an experimental map() op.#88767

Closed
zhxchen17 wants to merge 1 commit intopytorch:masterfrom
zhxchen17:export-D41165796
Closed

[functorch] introduce an experimental map() op.#88767
zhxchen17 wants to merge 1 commit intopytorch:masterfrom
zhxchen17:export-D41165796

Conversation

@zhxchen17
Copy link
Copy Markdown
Contributor

@zhxchen17 zhxchen17 commented Nov 9, 2022

Summary:
We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.

Some calrification on basic requirements we have in mind:

  1. This op can nest cond() and other control flow primitives internally.
  2. We don't necessarily need loop carried dependencies for the models we've seen.
  3. This map() op can handle dynamically shaped tensor as input and return dynamically shaped output based on input shapes.
  4. We should be able to pass through additional arguments to the loop body as extra arguments.

In this diff we introduce a new control flow op map() which has the following semantics:

def map(f: Callable, xs: Tensor, *args):
    # one possible implementation:
    return torch.stack([f(x, *args) for x in xs])

Test Plan:
pytest functorch/test_control_flow.py
CI

Differential Revision: D41165796

@pytorch-bot
Copy link
Copy Markdown

pytorch-bot bot commented Nov 9, 2022

🔗 Helpful Links

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

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

✅ No Failures

As of commit 4017a0b:
💚 Looks good so far! There are no failures yet. 💚

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

@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

@ezyang
Copy link
Copy Markdown
Contributor

ezyang commented Nov 9, 2022

What exactly is being looped over here? Are you permitting lists of varying sizes in the IR format? (If so, that makes me very negative on this feature. cond is OK because it doesn't introduce new data types to the FX IR. This would.)

@zhxchen17
Copy link
Copy Markdown
Contributor Author

What exactly is being looped over here? Are you permitting lists of varying sizes in the IR format? (If so, that makes me very negative on this feature. cond is OK because it doesn't introduce new data types to the FX IR. This would.)

@ezyang

What exactly is being looped over here?

I think the semantics is very similar to how jax.lax.map works, basically we loop over the second argument:

def map(f, xs, *args):
    return torch.stack([f(x, *args) for x in xs])

Are you permitting lists of varying sizes in the IR format?

No we don't permit that, varying size list is not useful to us and we don't want to add it. xs argument must be a tensor, and we have a check during tracing:

# map.py, line 29
assert isinstance(xs, torch.Tensor), "map() must operate over a tensor"

@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

1 similar comment
@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

@ezyang
Copy link
Copy Markdown
Contributor

ezyang commented Nov 10, 2022

OK, makes more sense.

@ezyang
Copy link
Copy Markdown
Contributor

ezyang commented Nov 11, 2022

deferring mostly to @voznesenskym and @zou3519 for review on this. Holler if you want me to look at something specific.

Copy link
Copy Markdown
Contributor

@zou3519 zou3519 left a comment

Choose a reason for hiding this comment

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

I have some questions, but this seems reasonable to me, it's in line with what we did with cond

@facebook-github-bot
Copy link
Copy Markdown
Contributor

@zhxchen17 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@zhxchen17 zhxchen17 force-pushed the export-D41165796 branch 2 times, most recently from 1fc1088 to e4c39a0 Compare November 16, 2022 07:48
@zhxchen17
Copy link
Copy Markdown
Contributor Author

addressed all comments.

@zhxchen17 zhxchen17 requested a review from zou3519 November 16, 2022 07:54
@zhxchen17 zhxchen17 force-pushed the export-D41165796 branch 2 times, most recently from b81d463 to 631790b Compare November 18, 2022 02:24
zhxchen17 added a commit that referenced this pull request Nov 18, 2022
Summary:
I saw the following issue only on Windows build in PR #88767:
```
RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge'
```
The reason it's only on Windows is because we get the attributes of SymNode in C++ with
`__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior.
In this case, MSVC will include a function's namespace and class name, which is not intended here.

Instead we should use `__func__`. see: https://en.cppreference.com/w/c/language/function_definition

godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 57cbfab
Pull Request resolved: #89264
zhxchen17 added a commit that referenced this pull request Nov 18, 2022
Summary:
I saw the following issue only on Windows build in PR #88767:
```
RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge'
```
The reason it's only on Windows is because we get the attributes of SymNode in C++ with
`__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior.
In this case, MSVC will include a function's namespace and class name, which is not intended here.

Instead we should use `__func__`. see: https://en.cppreference.com/w/c/language/function_definition

godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
pytorchmergebot pushed a commit that referenced this pull request Nov 18, 2022
Summary:
I saw the following issue only on Windows build in PR #88767:
```
RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge'
```
It's only on Windows because we get the attributes of SymNode in C++ with
`__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior.
In this case, MSVC will include a function's namespace and class name, which is not intended here.

Instead we should use `__func__`. see: https://en.cppreference.com/w/cpp/language/function#Function_definition

godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:
Pull Request resolved: #89264
Approved by: https://github.com/ezyang
Summary:

We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.

Some calrification on basic requirements we have in mind:
1. This op can nest cond() and other control flow primitives internally.
2. We don't necessarily need loop carried dependencies for the models we've seen.
3. This map() op can handle dynamically shaped tensor as input and return dynamically shaped output based on input shapes.
4. We should be able to pass through additional arguments to the loop body as extra arguments.

In this diff we introduce a new control flow op `map()` which has the following semantics:
```
def map(f: Callable, xs: Tensor, *args):
    # one possible implementation:
    return torch.stack([f(x, *args) for x in xs])
```

Test Plan:
pytest functorch/test_control_flow.py
CI

Reviewers:

Subscribers:

Tasks:

Tags:
@facebook-github-bot
Copy link
Copy Markdown
Contributor

@zhxchen17 has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@zhxchen17
Copy link
Copy Markdown
Contributor Author

@pytorchbot merge

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Nov 18, 2022
@pytorchmergebot
Copy link
Copy Markdown
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

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

zhxchen17 added a commit that referenced this pull request Nov 29, 2022
…nctorch.experimental.control_flow.

Summary:
Similar to #88767 we want to reduce the chance that users
accidentally import private functions from `functorch.experimental.cond` as if they were public
interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op.

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
zhxchen17 added a commit that referenced this pull request Nov 29, 2022
…nctorch.experimental.control_flow.

Summary:
Similar to #88767 we want to reduce the chance that users
accidentally import private functions from `functorch.experimental.cond` as if they were public
interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op.

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: bcce75d
Pull Request resolved: #89819
pytorchmergebot pushed a commit that referenced this pull request Nov 30, 2022
…nctorch.experimental.control_flow. (#89819)

Summary:
Similar to #88767 we want to reduce the chance that users
accidentally import private functions from `functorch.experimental.cond` as if they were public
interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op.

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

Pull Request resolved: #89819
Approved by: https://github.com/zou3519
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Dec 10, 2022
Summary:
I saw the following issue only on Windows build in PR pytorch#88767:
```
RuntimeError: AttributeError: 'SymNode' object has no attribute 'torch::impl::PythonSymNodeImpl::ge'
```
It's only on Windows because we get the attributes of SymNode in C++ with
`__FUNCTION__` macro, which is not in C++ standard, therefore has platform specific behavior.
In this case, MSVC will include a function's namespace and class name, which is not intended here.

Instead we should use `__func__`. see: https://en.cppreference.com/w/cpp/language/function#Function_definition

godbolt example to show the difference: https://godbolt.org/z/PGfvecxPx

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:
Pull Request resolved: pytorch#89264
Approved by: https://github.com/ezyang
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Dec 10, 2022
Summary:
We want to introduce an experimental control flow op: map() to export some models as FX graphs correctly.

Some calrification on basic requirements we have in mind:
1. This op can nest cond() and other control flow primitives internally.
2. We don't necessarily need loop carried dependencies for the models we've seen.
3. This map() op can handle dynamically shaped tensor as input and return dynamically shaped output based on input shapes.
4. We should be able to pass through additional arguments to the loop body as extra arguments.

In this diff we introduce a new control flow op `map()` which has the following semantics:
```
def map(f: Callable, xs: Tensor, *args):
    # one possible implementation:
    return torch.stack([f(x, *args) for x in xs])
```

Test Plan:
pytest functorch/test_control_flow.py
CI

Differential Revision: D41165796

Pull Request resolved: pytorch#88767
Approved by: https://github.com/zou3519
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Dec 10, 2022
…nctorch.experimental.control_flow. (pytorch#89819)

Summary:
Similar to pytorch#88767 we want to reduce the chance that users
accidentally import private functions from `functorch.experimental.cond` as if they were public
interfaces. We also move `cond()` under `control_flow.py` to stay consistent with `map()` op.

Test Plan:
CI

Reviewers:

Subscribers:

Tasks:

Tags:

Pull Request resolved: pytorch#89819
Approved by: https://github.com/zou3519
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk Trigger trunk jobs on your pull request fb-exported Merged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants