Commit Graph

57 Commits

Author SHA1 Message Date
Horace He
39130ccf73 Registered _like metas (#85793)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85793
Approved by: https://github.com/ezyang
2022-09-28 17:23:07 +00:00
PyTorch MergeBot
b44a4a8b51 Revert "Registered _like metas (#85793)"
This reverts commit a4e75ccf85.

Reverted https://github.com/pytorch/pytorch/pull/85793 on behalf of https://github.com/huydhn due to Sorry, reverting as this breaks an aot_autograd mac test on functorch. https://github.com/pytorch/pytorch/pull/85794 was reverted before but it was at the top of the stack so the revert still fail 823dc33b00
2022-09-28 17:18:29 +00:00
Horace He
a4e75ccf85 Registered _like metas (#85793)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85793
Approved by: https://github.com/ezyang
2022-09-28 14:07:57 +00:00
Horace He
2f4a517d67 Ported matmul compositeimplicitautograd impl into core (#85239)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85239
Approved by: https://github.com/ezyang, https://github.com/lezcano
2022-09-21 09:25:24 +00:00
Elias Ellison
a3afb2c2f6 Fake: fix conv_transpose2d striding (#82846)
The output striding channels-last preservation logic differs between cuda and cpu. For the meta kernel, we can peek at the fake tensor device and use that to determine whether to do cpu or cuda.

You could argue there's a leaking of abstraction here but this seems like a pretty minimal leak and I'm not sure there's a much cleaner way forward for device-specific striding tracing logic.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82846
Approved by: https://github.com/ezyang
2022-09-20 18:00:59 +00:00
Sherlock Huang
29eba319b4 Use alias for nop decomp (#84727)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84727
Approved by: https://github.com/Chillee
2022-09-16 18:50:56 +00:00
Horace He
5ea2eb304e Converted batch norm over to use symints (#84113)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84113
Approved by: https://github.com/wconstab, https://github.com/ezyang
2022-09-12 05:36:24 +00:00
Natalia Gimelshein
72f0f24a76 remove unneeded _to_copy meta (#84460)
Fixes #84335

Pull Request resolved: https://github.com/pytorch/pytorch/pull/84460
Approved by: https://github.com/Chillee
2022-09-02 18:08:39 +00:00
Edward Z. Yang
d39490a711 Add meta function for repeat (#84349)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84349
Approved by: https://github.com/Krovatkin
2022-09-01 20:44:52 +00:00
Edward Z. Yang
cd96f3f676 Use register_meta for everything in meta_registrations (#84297)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84297
Approved by: https://github.com/Chillee
2022-08-31 23:58:24 +00:00
PyTorch MergeBot
65f98eb47d Revert "Add meta function for repeat (#84349)"
This reverts commit 44bc6db8f8.

Reverted https://github.com/pytorch/pytorch/pull/84349 on behalf of https://github.com/janeyx99 due to Land race with the revert causing test_fx failures 44bc6db8f8
2022-08-31 18:27:59 +00:00
Edward Z. Yang
44bc6db8f8 Add meta function for repeat (#84349)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84349
Approved by: https://github.com/Krovatkin
2022-08-31 17:20:21 +00:00
PyTorch MergeBot
14093b5979 Revert "Use register_meta for everything in meta_registrations (#84297)"
This reverts commit 8cd296f680.

Reverted https://github.com/pytorch/pytorch/pull/84297 on behalf of https://github.com/suo due to broke test_proxy_tensor on master
2022-08-31 16:32:24 +00:00
Edward Z. Yang
8cd296f680 Use register_meta for everything in meta_registrations (#84297)
Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84297
Approved by: https://github.com/Chillee
2022-08-31 02:13:21 +00:00
Mario Lezcano
f5a3515083 Make linalg.inv composite of linalg.solve (#80074)
The `getri` kernel calls inside `getrs` so we can do so explicitly
ourselves and save ourselves from having to maintain an extra kernel.
This way we just need to optimise `lu_factor` and `lu_solve` and `inv`
will be as efficient as it can be, as it'll be choosing the best backend
to perform the factorisation and the best backend (not necessarily the
same) to perform the solve.

Fixes https://github.com/pytorch/pytorch/issues/77498

The benchmarks: https://github.com/pytorch/pytorch/pull/80074#issuecomment-1164309071
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80074
Approved by: https://github.com/IvanYashchuk, https://github.com/albanD, https://github.com/malfet
2022-08-25 09:28:55 +00:00
PyTorch MergeBot
5321bf52f2 Revert "Make linalg.inv composite of linalg.solve (#80074)"
This reverts commit 4737b33614.

Reverted https://github.com/pytorch/pytorch/pull/80074 on behalf of https://github.com/malfet due to Depends on the changes from https://github.com/pytorch/pytorch/pull/83628
2022-08-25 00:43:00 +00:00
Mario Lezcano
4737b33614 Make linalg.inv composite of linalg.solve (#80074)
The `getri` kernel calls inside `getrs` so we can do so explicitly
ourselves and save ourselves from having to maintain an extra kernel.
This way we just need to optimise `lu_factor` and `lu_solve` and `inv`
will be as efficient as it can be, as it'll be choosing the best backend
to perform the factorisation and the best backend (not necessarily the
same) to perform the solve.

Fixes https://github.com/pytorch/pytorch/issues/77498

The benchmarks: https://github.com/pytorch/pytorch/pull/80074#issuecomment-1164309071
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80074
Approved by: https://github.com/IvanYashchuk, https://github.com/albanD, https://github.com/malfet
2022-08-24 15:18:56 +00:00
Horace He
7ebdb4c72f Refactored ops on size to be dispatcher ops (#83719)
An example of how the graph looks now.
```
def forward(self, x_1):
    size = torch.ops.math.size(x_1, 0)
    size_1 = torch.ops.math.size(x_1, 1);  x_1 = None
    ones = torch.ops.aten.ones.default([1], device = device(type='cpu'), pin_memory = False)
    expand_sym_int = torch.ops.aten.expand.SymInt(ones, [size, size_1]);  ones = size = size_1 = None
    cos_default = torch.ops.aten.cos.default(expand_sym_int);  expand_sym_int = None
    return (cos_default,)
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/83719
Approved by: https://github.com/ezyang
2022-08-23 15:48:00 +00:00
Elias Ellison
8a6b076196 lift numpy tensor, add randperm support (#83191)
Couple changes needed to trace huggingface w fake tensors.

Similar to https://github.com/pytorch/pytorch/pull/81927, need to call liftfresh for tensors created from numpy tensors. Also adds randperm for meta.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/83191
Approved by: https://github.com/bdhirsh
2022-08-10 22:27:51 +00:00
Elias Ellison
1c0f7bd6d2 Enable complex for meta tensors (#79975)
There weren't really any fundamental blockers
- add support for `aten::complex`
- update `angle` for complex
- remove the error in the fallback kernel
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79975
Approved by: https://github.com/ezyang
2022-07-27 22:19:14 +00:00
Horace He
fc389cc0a0 Added new_empty.symint overload and a new_empty ref (#82049)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82049
Approved by: https://github.com/ezyang
2022-07-27 00:31:57 +00:00
Robert
47b4ec8aa7 Conv2d shape calculation for meta tensors (#79834)
Fixes #79512

This PR adds support for convolutional meta modules and computes the output shape correctly for some meta input tensor.
Currently in progress, no tests written so far.

**Feature implementations**:
- [x] `Conv1d`
- [x] `Conv2d`
- [x] `Conv3d`

**Tests**:
- [x] `Conv1d`
- [x] `Conv2d`
- [x] `Conv3d`

cc @albanD @anjali411
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79834
Approved by: https://github.com/ezyang, https://github.com/albanD
2022-07-23 05:58:56 +00:00
Huy Do
12cb26509a Apply ufmt to torch internal (#81643)
This is a big bang PR, merge conflicts are probably expected and will be addressed at merge.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81643
Approved by: https://github.com/ezyang
2022-07-22 02:19:50 +00:00
Horace He
2529ff4bd9 Registered python meta functions to a table (#81092)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81092
Approved by: https://github.com/ezyang, https://github.com/anjali411
2022-07-21 21:45:43 +00:00
Horace He
a5fb41e3d3 Revert "Revert "Refactored prim utils into _prims_utils folder (#81746)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81746
Approved by: https://github.com/anijain2305, https://github.com/Krovatkin
2022-07-20 23:43:57 +00:00
PyTorch MergeBot
e43a02c314 Revert "Refactored prim utils into _prims_utils folder (#81088)"
This reverts commit 80231d0a72.

Reverted https://github.com/pytorch/pytorch/pull/81088 on behalf of https://github.com/jeanschmidt due to breaking internal tests
2022-07-19 19:56:41 +00:00
Horace He
80231d0a72 Refactored prim utils into _prims_utils folder (#81088)
Pull Request resolved: https://github.com/pytorch/pytorch/pull/81088
Approved by: https://github.com/ngimel
2022-07-19 03:55:51 +00:00
lezcano
eb0889cf7d Add support for multiple inputs to out_wrapper and strict dtype checking (#80601)
Reland of https://github.com/pytorch/pytorch/pull/79941
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80601
Approved by: https://github.com/albanD
2022-07-05 12:31:21 +00:00
PyTorch MergeBot
184a065ba7 Revert "Add support for multiple inputs to out_wrapper and strict dtype checking (#79941)"
This reverts commit dc7066a8f0.

Reverted https://github.com/pytorch/pytorch/pull/79941 on behalf of https://github.com/suo due to broke master dc7066a8f0
2022-06-30 03:29:30 +00:00
lezcano
dc7066a8f0 Add support for multiple inputs to out_wrapper and strict dtype checking (#79941)
When a function returns multiple parameters in PyTorch, the `out`
parameter takes a tuple of tensors (see `linalg.svd` for example).
The current implementation in `out_wrapper_multi` modelled this wrong,
as it assumed that it would take a number of different named
parameters.

This PR implements the correct behaviour in `out_wrapper`. As a small
side-effect, we now need to call `@out_wrapper()` when the output is
just one tensor.

This PR also implements an additional optional parameter that checks
whether the dtype of the given `out` is exactly the dtype that the meta
function requires. This is the behaviour that we currently have in
PyTorch, and this check is necessary in eager when we call with these
tensors into external libraries.

We also make the functions with several outputs return a namedtuple,
similar to what we do in PyTorch.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79941
Approved by: https://github.com/mruberry, https://github.com/ezyang
2022-06-30 02:47:16 +00:00
Edward Z. Yang
4331bc436e Ensure torch._refs registrations also get triggered on import torch (#80270)
Fixes https://github.com/pytorch/pytorch/issues/79938

Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80270
Approved by: https://github.com/ngimel
2022-06-26 02:23:03 +00:00
lezcano
ff5a588e6e Port cholesky to structured kernels (#79300)
Yeah.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79300
Approved by: https://github.com/IvanYashchuk, https://github.com/albanD
2022-06-24 02:37:45 +00:00
PyTorch MergeBot
79507d2a9d error when registering meta kernels to composite ops in core
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79741

Approved by: https://github.com/Chillee, https://github.com/albanD
2022-06-21 02:17:13 +00:00
Elias Ellison
9705fb03b3 Add support for a couple ops
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79581

Approved by: https://github.com/Chillee
2022-06-20 22:25:39 +00:00
lezcano
648a6658ec Remove python implementation for eigh meta
Following https://github.com/pytorch/pytorch/pull/79072#discussion_r898210048

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79786

Approved by: https://github.com/ngimel, https://github.com/bdhirsh
2022-06-17 18:52:28 +00:00
kshitij12345
31ada133cb [meta] nansum, nanmedian (and few minor clean-ups) (#79411)
meta support for `nansum` and `nanmedian`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79411
Approved by: https://github.com/anjali411
2022-06-14 16:21:13 +00:00
kshitij12345
a732bbea23 [meta] Add meta support for fft ops (#79311)
As per title
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79311
Approved by: https://github.com/ezyang
2022-06-13 01:56:42 +00:00
kshitij12345
bd1a35dfc8 [meta] diag ops, trace (#79341)
meta registration for `diag.out` and update test skips/expectedFailures
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79341
Approved by: https://github.com/ezyang
2022-06-12 18:45:03 +00:00
kshitij12345
7b307e5fca [meta] angle, angle.out (#79278)
meta registration for `angle, angle.out`.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/79278
Approved by: https://github.com/anjali411
2022-06-10 20:06:31 +00:00
Brian Hirsh
7b3a0ff87a Port index.Tensor to structured kernels.
Tracking issue: #55070

Pull Request resolved: https://github.com/pytorch/pytorch/pull/69607

Approved by: https://github.com/bdhirsh
2022-06-10 17:27:47 +00:00
lezcano
af6321f3d8 Port linalg_qr to structured
This PR simplifies the logic of `linalg.qr` using structured kernels. I
also took this chance and merged a few `copy_` operations with other
ops.

This PR removes a the previous magma implementation as is never faster
than that of cusolver and it's rather buggy. This has the side-effect
that now `qr` is not supported in Rocm. Ivan confirmed that this is
fine, given how incredibly slow was QR on Rocm anyway (we were marking
some tests as slow because of this...).

This PR also corrects the dispatch in geqrf. Before, if we called it
with a matrix for which `input.size(-2) <= 256 && batchCount(input) >= std::max<int64_t>(2, input.size(-2) / 16)` is false, and we have cublas but not cusolver, we would end up calling magma rather than cublas. This is not what the heuristic suggested.
Probaly we should benchmark these heuristics again, but that's beyond the scope of this PR.

Note. It looks like `torch.geqrf` maybe broken in MAGMA as per the
previous comment in `linalg_qr_helper_magma`. IvanYashchuk wdyt?

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79054

Approved by: https://github.com/IvanYashchuk, https://github.com/ezyang
2022-06-09 14:41:30 +00:00
Edward Z. Yang
225bf132ab Black torch._meta_registrations
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/79026

Approved by: https://github.com/Chillee
2022-06-09 03:03:09 +00:00
Edward Z. Yang
50f2af84da Add embedding_bag meta functions
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78997

Approved by: https://github.com/Chillee, https://github.com/Lezcano
2022-06-08 22:03:27 +00:00
Edward Z. Yang
41bd5b85fd cdist meta function
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78993

Approved by: https://github.com/Lezcano, https://github.com/Chillee
2022-06-08 01:57:00 +00:00
Edward Z. Yang
d09e3674d8 addbmm meta function
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78992

Approved by: https://github.com/Lezcano, https://github.com/Chillee
2022-06-07 23:24:57 +00:00
lezcano
c7d6cec078 Add linalg.lu_solve
This PR adds `linalg.lu_solve`. While doing so, I found a bug in MAGMA
when calling the batched MAGMA backend with trans=True. We work around
that by solving the system solving two triangular systems.

We also update the heuristics for this function, as they were fairly
updated. We found that cuSolver is king, so luckily we do not need to
rely on the buggy backend from magma for this function.

We added tests testing this function left and right. We also added tests
for the different backends. We also activated the tests for AMD, as
those should work as well.

Fixes https://github.com/pytorch/pytorch/issues/61657

Pull Request resolved: https://github.com/pytorch/pytorch/pull/77634

Approved by: https://github.com/malfet
2022-06-07 22:28:28 +00:00
Edward Z. Yang
157d478a30 Fix omission of shape in size check in index.
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78897

Approved by: https://github.com/Lezcano, https://github.com/anjali411
2022-06-05 23:10:55 +00:00
Edward Z. Yang
99882fc492 Make check() strongly typed, fix erroneous call sites
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78896

Approved by: https://github.com/Lezcano, https://github.com/anjali411
2022-06-05 23:10:55 +00:00
Edward Z. Yang
83d40a4dba linalg_cholesky_ex meta function
Taken from https://github.com/albanD/subclass_zoo/blob/main/python_meta_tensor.py

Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78604

Approved by: https://github.com/bdhirsh, https://github.com/ngimel, https://github.com/Lezcano
2022-06-03 23:11:02 +00:00
Edward Z. Yang
6120a8e05d Implement meta function for aten::index.Tensor
Signed-off-by: Edward Z. Yang <ezyangfb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/78527

Approved by: https://github.com/bdhirsh, https://github.com/ngimel, https://github.com/Lezcano
2022-06-03 23:11:02 +00:00