Skip to content

[ET-VK][Ops] aten.embedding#3762

Closed
junpi3 wants to merge 4 commits intogh/jorgep31415/64/basefrom
gh/jorgep31415/64/head
Closed

[ET-VK][Ops] aten.embedding#3762
junpi3 wants to merge 4 commits intogh/jorgep31415/64/basefrom
gh/jorgep31415/64/head

Conversation

@junpi3
Copy link
Copy Markdown
Contributor

@junpi3 junpi3 commented May 29, 2024

Stack from ghstack (oldest at bottom):

The Operator

nn.Module invocations on the embedding returned by torch.nn.Embedding get compiled to aten.embedding.default in the Edge Dialect, which carries the following signature.

- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor

Implementation

This is a C-packing-only implementation.

Interestingly, the 1D-indices case is equivalent to the dim=0 case of the preceding aten.index_select: #3744

- func: index_select(Tensor self, int dim, Tensor index) -> Tensor

I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-indices cases are more involved to the extent that we require a standalone cpp/glsl file.

Codegen

We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the case_name string to recursively handle list pylists.

// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),

Differential Revision: D57880520

## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: #3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)

[ghstack-poisoned]
@pytorch-bot
Copy link
Copy Markdown

pytorch-bot Bot commented May 29, 2024

🔗 Helpful Links

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

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 7b4b2ae with merge base 55d11e1 (image):
💚 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 facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 29, 2024
@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

junpi3 pushed a commit that referenced this pull request May 29, 2024
## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: #3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)

ghstack-source-id: 228038402
Pull Request resolved: #3762
## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: #3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)

[ghstack-poisoned]
@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: #3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)

[ghstack-poisoned]
@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: #3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)

[ghstack-poisoned]
@facebook-github-bot
Copy link
Copy Markdown
Contributor

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

@facebook-github-bot
Copy link
Copy Markdown
Contributor

This pull request has been merged in a36ace7.

kedarnath03 pushed a commit to kedarnath03/executorch that referenced this pull request Jun 25, 2025
Pull Request resolved: pytorch/executorch#3762

## The Operator
`nn.Module` invocations on the embedding returned by [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) get compiled to `aten.embedding.default` in the Edge Dialect, which carries the following signature.
```
- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor
```

## Implementation
This is a C-packing-only implementation.

Interestingly, the 1D-`indices` case is equivalent to the `dim=0` case of the preceding `aten.index_select`: pytorch/executorch#3744
```
- func: index_select(Tensor self, int dim, Tensor index) -> Tensor
```
I naïvely thought the rest of the operator would be similarly easy but it wasn't. The 2D and 3D-`indices` cases are more involved to the extent that we require a standalone `cpp`/`glsl` file.

## Codegen
We add support for making 2D and 3D index tensors. This requires new generation functions as well as renaming of the `case_name` string to recursively handle list `pylist`s.
```
// 1D
Test(weight=[10, 9], indices=[0, 2]),
// 2D
Test(weight=[10, 9], indices=[[0, 2], [1, 4], [7, 7]]),
// 3D
Test(weight=[10, 9], indices=[[[3, 1, 4], [1, 5, 9]], [[2, 6, 5], [3, 5, 8]]]),
```
ghstack-source-id: 228201965

Differential Revision: [D57880520](https://our.internmc.facebook.com/intern/diff/D57880520/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. fb-exported Merged

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants