Skip to content

Add batched linear solver to torch.gesv() #6100

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.

Sign up for GitHub

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

Jump to bottom
Merged
soumith merged 11 commits into from
May 8, 2018
Merged

Add batched linear solver to torch.gesv() #6100

soumith merged 11 commits into from
May 8, 2018

Conversation

@zou3519
Copy link
Contributor

@zou3519 zou3519 commented Mar 29, 2018

Fixes #3164
Picks up from #4502

I moved gesv to ATen.
Adds bindings for MAGMA's gesv_batched function for CUDA.
For CPU, runs THLapack(gesv) in a for loop.

The new function supports arbitrary batch dimensions (and broadcasting
of those dimensions). For example, the 4-d tensor A x B x M x M should
be treated as having batch-size (A x B).

The overhead of creating the magma_queue_t is: ~350000 microseconds
the first time it's called and ~6 microseconds every time after that.

cc @colesbury @apaszke

@onnxbot-worker-3 onnxbot-worker-3 mentioned this pull request Mar 29, 2018
Closed
ssnl
ssnl reviewed Mar 29, 2018
View reviewed changes
test/test_torch.py Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed Mar 29, 2018
View reviewed changes
tools/autograd/templates/Functions.cpp Outdated

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed Mar 29, 2018
View reviewed changes
aten/src/ATen/native/LinearAlgebraUtils.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed Mar 29, 2018
View reviewed changes
aten/src/ATen/native/LinearAlgebraUtils.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed Mar 30, 2018
View reviewed changes
aten/src/ATen/Declarations.cwrap Outdated

This comment was marked as off-topic.

Sign in to view

@fmassa fmassa mentioned this pull request Apr 25, 2018
Closed
@zou3519
Copy link
Contributor Author

zou3519 commented Apr 27, 2018

@ssnl could I get another review on this please? I've addressed all the previous comments. IIRC, your last pending comment was that we should delete the TH gesv and move it to ATen as well (because the batched gesv does the TH gesv in a loop). Could we leave that for a future PR? (I'll open another issue for it). It's much more difficult to implement because of the following:

  • the single version of gesv takes an _out keyword. Implementing that efficiently (or even just as efficiently as in TH) is going to take a lot of refactoring, especially because many of the ATen functions are designed without an out keyword in mind.

I don't think it's worth the marginal benefit to deduplicate our gesv code right now, especially since users have been requesting batched gesv for a while.

ssnl
ssnl reviewed Apr 28, 2018
View reviewed changes
aten/src/ATen/native/LinearAlgebraUtils.cpp Outdated

This comment was marked as off-topic.

Sign in to view

aten/src/ATen/native/LinearAlgebraUtils.cpp Outdated

This comment was marked as off-topic.

Sign in to view

aten/src/ATen/native/Gesv.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

aten/src/ATen/native/cuda/Gesv.cu Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

tools/autograd/templates/python_variable_methods.cpp Outdated

This comment was marked as off-topic.

Sign in to view

torch/_torch_docs.py Outdated

This comment was marked as off-topic.

Sign in to view

ssnl
ssnl reviewed May 1, 2018
View reviewed changes
aten/src/ATen/native/Gesv.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

aten/src/ATen/native/Gesv.cpp Outdated

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

This comment was marked as off-topic.

Sign in to view

@onnxbot-worker-1 onnxbot-worker-1 mentioned this pull request May 1, 2018
Closed
ssnl
ssnl reviewed May 4, 2018
View reviewed changes
Copy link
Collaborator

@ssnl ssnl left a comment

Choose a reason for hiding this comment

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

Two minor comments

torch/_torch_docs.py Outdated

This comment was marked as off-topic.

Sign in to view

test/test_torch.py Outdated

This comment was marked as off-topic.

Sign in to view

zou3519 added 8 commits May 7, 2018 07:20
@zou3519
Add batched linear solver to torch.gesv()
d75c86e 
Fixes pytorch#3164
Picks up from pytorch#4502

I moved `gesv` to ATen.
Adds bindings for MAGMA's `gesv_batched` function for CUDA.
For CPU, runs `THLapack(gesv)` in a for loop.

The new function supports arbitrary batch dimensions (and broadcasting
of those dimensions). For example, the 4-d tensor `A x B x M x M` should
be treated as having batch-size `(A x B)`.

The overhead of creating the magma_queue_t is: ~350000 microseconds
the first time it's called and ~6 microseconds every time after that.
@zou3519
Tests and docs
eadcc4d
@zou3519
Address comments
67baea9
@zou3519
Address comments
319f1d8
@zou3519
Rebase
5acacf2
@zou3519
Address comments
1ab57fe
@zou3519
Fix rebase
a88046f
@zou3519
Addressed comments
62bd344
@zou3519
Copy link
Contributor Author

zou3519 commented May 7, 2018

@ssnl done :) sorry for the wait!

@ssnl
Copy link
Collaborator

ssnl commented May 7, 2018

@pytorchbot retest this please

1 similar comment
@ssnl
Copy link
Collaborator

ssnl commented May 8, 2018

@pytorchbot retest this please

ssnl
ssnl approved these changes May 8, 2018
View reviewed changes
Copy link
Collaborator

@ssnl ssnl left a comment

Choose a reason for hiding this comment

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

LGTM

@soumith soumith merged commit 7162649 into pytorch:master May 8, 2018
onnxbot added a commit to onnxbot/onnx-fb-universe that referenced this pull request May 8, 2018
@onnxbot
[auto] Update pytorch to 7162649 - Add batched linear solver to torch…
1b5f21c 
….gesv() (pytorch/pytorch#6100)

pytorch/pytorch@7162649
weiyangfb pushed a commit to weiyangfb/pytorch that referenced this pull request Jun 11, 2018
@zou3519 @weiyangfb
Add batched linear solver to torch.gesv() (pytorch#6100)
207b5d6 
* Add batched linear solver to torch.gesv()

Fixes pytorch#3164
Picks up from pytorch#4502

I moved `gesv` to ATen.
Adds bindings for MAGMA's `gesv_batched` function for CUDA.
For CPU, runs `THLapack(gesv)` in a for loop.

The new function supports arbitrary batch dimensions (and broadcasting
of those dimensions). For example, the 4-d tensor `A x B x M x M` should
be treated as having batch-size `(A x B)`.

The overhead of creating the magma_queue_t is: ~350000 microseconds
the first time it's called and ~6 microseconds every time after that.

* Tests and docs

* Address comments

* Address comments

* Rebase

* Address comments

* Fix rebase

* Addressed comments

* Address comments

* Address comments

* Addressed comments
@zou3519 zou3519 mentioned this pull request Jul 24, 2018
Closed
@fmassa fmassa mentioned this pull request Nov 6, 2018
Closed
@gpleiss gpleiss mentioned this pull request May 28, 2019
Closed
14 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ssnl ssnl ssnl approved these changes

@apaszke apaszke Awaiting requested review from apaszke

@colesbury colesbury Awaiting requested review from colesbury

@ezyang ezyang Awaiting requested review from ezyang

@gchanan gchanan Awaiting requested review from gchanan

@zdevito zdevito Awaiting requested review from zdevito

+1 more reviewer

@soumith soumith soumith approved these changes

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@zou3519 @ssnl @soumith