Skip to content

Allow cuFFT plans to be used as a context manager; Set stream before executing cuFFT plans#2362

Merged
asi1024 merged 9 commits intocupy:masterfrom
leofang:fft_ctx_manager
Aug 14, 2019
Merged

Allow cuFFT plans to be used as a context manager; Set stream before executing cuFFT plans#2362
asi1024 merged 9 commits intocupy:masterfrom
leofang:fft_ctx_manager

Conversation

@leofang
Copy link
Copy Markdown
Member

@leofang leofang commented Aug 5, 2019 •
edited
Loading

Resolves #2356, resolves #2357. This PR enables the following (perhaps more pythonic) usage:

import cupy
from cupyx.scipy.fftpack import get_fft_plan

x = cupy.arange(10)
fftplan = get_fft_plan(x)

with fftplan:
    y = cupy.fft.fft(x)

Note that the context manager pattern will work for functions in both cupy.fft and cupyx.scipy.fftpack.

Furthermore, this PR makes it easy to work with the new module cupyx.scipy.fft (to be added in #2355 following upstream changes):

import cupy
import cupyx.scipy.fft as cu_fft
import scipy.fft
from cupyx.scipy.fftpack import get_fft_plan

x = cupy.arange(10)
fftplan = get_fft_plan(x)

with scipy.fft.set_backend(cu_fft), fftplan:
    scipy.fft.fft(x)  # Calls into cupyx, use fftplan, and returns a cupy array

While the plan argument might be added in upstream, see #2356 (comment) and scipy/scipy#10302, we should not count on that.

Finally, no backward compatibility is broken in this PR.

Attn: @grlee77, @peterbell10

grlee77
grlee77 reviewed Aug 6, 2019
View reviewed changes
Comment thread cupy/fft/fft.py Outdated
out_size = a.shape[-1]

batch = a.size // a.shape[-1]
plan = cufft.get_current_plan()
Copy link
Copy Markdown
Member

@grlee77 grlee77 Aug 6, 2019

Choose a reason for hiding this comment

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

Doesn't this cause _exec_fft to ignore any plan the user passed in as a kwarg? i.e. if plan is not None, and we are not in a plan context then the provided plan would not be used and Plan1d would be called?

I thought the logic in this section would be something like:

if plan is None:
    plan = cufft.get_current_plan()
    if plan is None:
        plan = cufft.Plan1d(out_size, fft_type, batch)
else:

Copy link
Copy Markdown
Member

@grlee77 grlee77 Aug 6, 2019

Choose a reason for hiding this comment

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

or alternatively, like in the _default_fft_func below:

curr_plan = cufft.get_current_plan()
if curr_plan is not None:
    plan = curr_plan

Copy link
Copy Markdown
Member Author

@leofang leofang Aug 6, 2019

Choose a reason for hiding this comment

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

Yeah good catch! Thank you, @grlee77. I probably got distracted when testing and forgot to revisit this part...

For the record, I was thinking the precedence is:

  1. plan from context
  2. plan from argument
  3. auto planning (for cupy.fft).

So, if the user uses both context and argument, by design I should have raised an error. I'll add this check when fixing the overridden plan.

grlee77
grlee77 reviewed Aug 6, 2019
View reviewed changes
Comment thread cupy/fft/fft.py Outdated
else:
raise ValueError('a must be contiguous')

plan = cufft.get_current_plan()
Copy link
Copy Markdown
Member

@grlee77 grlee77 Aug 6, 2019

Choose a reason for hiding this comment

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

same comment here as made for _exec_fft

@grlee77
Copy link
Copy Markdown
Member

grlee77 commented Aug 6, 2019

Allowing a plan context manager like this seems pretty nice. I am in favor of it.

@leofang
Copy link
Copy Markdown
Member Author

leofang commented Aug 6, 2019

@grlee77 thanks for reviewing and supporting this PR. The bug is fixed.

asi1024
asi1024 reviewed Aug 7, 2019
View reviewed changes
Comment thread cupy/cuda/cufft.pyx Outdated
asi1024
asi1024 reviewed Aug 7, 2019
View reviewed changes
Comment thread tests/cupyx_tests/scipy_tests/fftpack_tests/test_fftpack.py Outdated
asi1024
asi1024 reviewed Aug 7, 2019
View reviewed changes
Comment thread docs/source/reference/fft.rst Outdated
@leofang leofang mentioned this pull request Aug 7, 2019
Merged
@asi1024 asi1024 added the cat:enhancement Improvements to existing features label Aug 8, 2019
@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 8, 2019

pfnCI, test this please.

@pfn-ci-bot
Copy link
Copy Markdown
Collaborator

pfn-ci-bot commented Aug 8, 2019

Successfully created a job for commit 2517edb:

@asi1024 asi1024 added this to the v7.0.0b3 milestone Aug 8, 2019
@chainer-ci
Copy link
Copy Markdown
Member

chainer-ci commented Aug 8, 2019

Jenkins CI test (for commit 2517edb, target branch master) failed with status FAILURE.

@leofang
Copy link
Copy Markdown
Member Author

leofang commented Aug 8, 2019

Jenkins error on FFT tests is due to #2361.

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 8, 2019

We will merge #2363 first.

@asi1024 asi1024 added the st:blocked-by-another-pr Blocked by another pull-request label Aug 8, 2019
@leofang
Copy link
Copy Markdown
Member Author

leofang commented Aug 9, 2019

It's merged. Would you like to kick off Jenkins again?

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 9, 2019

pfnCI, test this please.

@pfn-ci-bot
Copy link
Copy Markdown
Collaborator

pfn-ci-bot commented Aug 9, 2019

Successfully created a job for commit 2517edb:

@asi1024 asi1024 removed the st:blocked-by-another-pr Blocked by another pull-request label Aug 9, 2019
@chainer-ci
Copy link
Copy Markdown
Member

chainer-ci commented Aug 9, 2019

Jenkins CI test (for commit 2517edb, target branch master) failed with status FAILURE.

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 9, 2019

pfnCI, test this please.

@pfn-ci-bot
Copy link
Copy Markdown
Collaborator

pfn-ci-bot commented Aug 9, 2019

Successfully created a job for commit 2517edb:

@chainer-ci
Copy link
Copy Markdown
Member

chainer-ci commented Aug 9, 2019

Jenkins CI test (for commit 2517edb, target branch master) failed with status FAILURE.

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 9, 2019

These test failures seem unrelated to this PR.

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 14, 2019

pfnCI, test this please.

@pfn-ci-bot
Copy link
Copy Markdown
Collaborator

pfn-ci-bot commented Aug 14, 2019

Successfully created a job for commit 2517edb:

@chainer-ci
Copy link
Copy Markdown
Member

chainer-ci commented Aug 14, 2019

Jenkins CI test (for commit 2517edb, target branch master) succeeded!

@asi1024
Copy link
Copy Markdown
Member

asi1024 commented Aug 14, 2019

LGTM. Thank you for the PR!

@asi1024 asi1024 merged commit 610e213 into cupy:master Aug 14, 2019
@leofang leofang deleted the fft_ctx_manager branch August 14, 2019 17:34
@leofang
Copy link
Copy Markdown
Member Author

leofang commented Aug 15, 2019

Thanks @asi1024!

@leofang leofang mentioned this pull request Aug 28, 2019
Merged
@leofang leofang mentioned this pull request Oct 3, 2019
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@asi1024 asi1024 asi1024 left review comments

+1 more reviewer

@grlee77 grlee77 grlee77 left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

@asi1024 asi1024

Labels

cat:enhancement Improvements to existing features

Projects

None yet

Milestone

v7.0.0b3

Development

Successfully merging this pull request may close these issues.

Set the stream for cuFFT plans as late as possible Possibility of making cufft.Plan1d and cufft.PlanNd be a context manager?

5 participants

@leofang @grlee77 @asi1024 @pfn-ci-bot @chainer-ci