Skip to content

CUDA: Replace CachedPTX and CachedCUFunction with CUDACodeLibrary functionality - #6769

Merged
stuartarchibald merged 11 commits into
numba:masterfrom
gmarkall:grm-cuda-codegen-ptx-cubin
Mar 17, 2021
Merged

stuartarchibald merged 11 commits into
numba:masterfrom
gmarkall:grm-cuda-codegen-ptx-cubin

Conversation

@gmarkall

Copy link
Copy Markdown
Member

Note: This PR is based on #6735 - the latest four commits are unique to this PR and can each be examined individually to understand the changes here.

This commit removes CachedPTX and CachedCUFunction, replacing them with additions to the CUDACodeLibrary. This aligns the CUDA target with more of the infrastructure in the core of Numba, simplifies the organization of code generation and caching, and provides some other benefits - generally it lays the foundation for implementations of on-disk caching of compiled kernels, ahead-of-time compilation, and improving debugging support for CUDA kernels. This description is expanded on in the final commit's message.

This PR also includes some other small changes and fixes - see individual commit messages for details.

There will be various places in the CUDA target where it will be helpful
to get the name of a CodeLibrary, so this commit makes it a public
property.
The CUDA Code Library will permit linking for a CC other than the
current device - this change to the linker in driver.py supports
specifying the CC to link for, where the default is the CC of the
current device.
The kernel invoked in the test is sync_group, but then the definitions
of this_grid are checked - the definitions of sync_group should be
checked instead.
This commit removes `CachedPTX` and `CachedCUFunction`, replacing them with
additions to the `CUDACodeLibrary`. This aligns the CUDA target with
more of the infrastructure in the core of Numba, simplifies the
organization of code generation and caching, and provides some other
benefits. Generally it lays the foundation for implementations of
on-disk caching of compiled kernels, ahead-of-time compilation, and
improving debugging support for CUDA kernels. Some specific advantages
of the new implementation:

- Although the `CUDACodeLibrary` can assume that code is to be generated
  for the current device (as is required to re-implement the behaviour
  of `CachedPTX` and `CachedCUFunction`), it is also possible to specify
  the compute capability (for PTX and cubins) for which the result
  should be returned. This aims towards avoiding interacting with a CUDA
  context at all for PTX and cubin generation - however, the latter will
  require further work to avoid using the Driver API for linking PTX,
  such as using the PTX Compiler API
  (https://docs.nvidia.com/cuda/ptx-compiler-api/index.html) instead. I
  did think about creating a `JITCUDACodeLibrary` class that contains
  all the context interactions, so that there are none in the base
  class, but until there's an actual implementation of a putative
  `AOTCUDACodeLibrary` this just feels like premature abstraction.
- The setting of NVVM options is easier to keep consistent - the options
  are initialized and placed into the flags in the compiler state, and
  are set in the code library when it is initialized. This avoids code
  being generated by a code library with a potentially different set of
  flags than the one it used earlier, as was possible with the old
  `Cached*` classes.
- Serializing a `CUDACodeLibrary` also serializes the cubins (which are
  just ELF files), which the `CachedCUFunction` class did not.

Some other notes on the changes:

- The `_func` member and property has been removed, and has not been
  replaced by anything. It could have been replaced by providing access
  to the code library of a dispatcher or kernel object, but instead I'd
  like to discourage poking about at the internals of these objects.
  This requires a slight change in `ForAll.compute_thread_per_block`
  which is a little less than perfect, but it seems no worse than it was
  before.
- Because of the initialization of a CUDA code library requiring the
  NVVM options, it was necessary to create a CUDA flags class and a pass
  to set up the code library with the NVVM options from the flags which
  runs just prior to the `NativeLowering` pass.
@gmarkall gmarkall added 2 - In Progress CUDA CUDA related issue/PR labels Feb 26, 2021

@stuartarchibald stuartarchibald left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for doing this refactoring, it's made the CUDA target much clearer and helped tidy up the CPU side. On inspection it looks largely good, just a few minor comments to resolve. I think this should have both a CPU and CUDA test run in the build farm just to make sure there's no surprises.

@sklam you might want to take a look at this too?

Thanks again!

Comment thread numba/core/codegen.py Outdated
Comment thread numba/core/codegen.py Outdated
Comment thread numba/cuda/codegen.py
Comment on lines +18 to +44
def disassemble_cubin(cubin):
# nvdisasm only accepts input from a file, so we need to write out to a
# temp file and clean up afterwards.
fd = None
fname = None
try:
fd, fname = tempfile.mkstemp()
with open(fname, 'wb') as f:
f.write(cubin)

def _optimize_final_module(self):
pass
try:
cp = subprocess.run(['nvdisasm', fname], check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
except FileNotFoundError as e:
if e.filename == 'nvdisasm':
msg = ("nvdisasm is required for SASS inspection, and has not "
"been found.\n\nYou may need to install the CUDA "
"toolkit and ensure that it is available on your "
"PATH.\n")
raise RuntimeError(msg)
return cp.stdout.decode('utf-8')
finally:
if fd is not None:
os.close(fd)
if fname is not None:
os.unlink(fname)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

NOTE: this is a direct move of the function from numba.cuda.compiler.

Comment thread numba/cuda/codegen.py Outdated
Comment thread numba/cuda/compiler.py Outdated
Comment on lines +299 to +301
if nvvm_options:
msg = 'nvvm_options kwarg for inspect_ptx is deprecated'
warn(msg, category=NumbaDeprecationWarning)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Want to add this to the deprecation notice doc?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I decided that deprecating it was a mistake (from the commit message un-deprecating it):

The idea behind deprecating this was to prevent a mismatch between
compiled code that is executed and code produced by the `inspect_ptx()`
function - however, as device functions aren't materialized standalone
(they are always part of a kernel), there is no risk of this mismatch.
Instead, `inspect_ptx()` stands as a useful mechanism for obtaining PTX
code for a device function if required - noting that `compile_ptx()`
does provide a better interface, so this whole function may be
deprecated / removed in future.

Comment thread numba/cuda/compiler.py Outdated
Returns the LLVM IR for this kernel.
'''
return str(self._func.ptx.llvmir)
return self._codelibrary.get_asm_str()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Is this right? Should it be get_llvm_str?

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

It's wrong - it should be as you suggest - now fixed.

@stuartarchibald

Copy link
Copy Markdown
Contributor

Also, conflicts need resolving, thanks!

@stuartarchibald stuartarchibald added 4 - Waiting on author Waiting for author to respond to review and removed 3 - Ready for Review labels Mar 12, 2021
@stuartarchibald stuartarchibald added this to the Numba 0.54 RC milestone Mar 12, 2021
The idea behind deprecating this was to prevent a mismatch between
compiled code that is executed and code produced by the `inspect_ptx()`
function - however, as device functions aren't materialized standalone
(they are always part of a kernel), there is no risk of this mismatch.
Instead, `inspect_ptx()` stands as a useful mechanism for obtaining PTX
code for a device function if required - noting that `compile_ptx()`
does provide a better interface, so this whole function may be
deprecated / removed in future.
@gmarkall

Copy link
Copy Markdown
Member Author

@stuartarchibald Many thanks for the review - conflicts resolved and comments now addressed.

@gmarkall gmarkall added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 4 - Waiting on author Waiting for author to respond to review labels Mar 12, 2021

@stuartarchibald stuartarchibald left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Thanks for the patch and fixes!

@stuartarchibald stuartarchibald added 4 - Waiting on CI Review etc done, waiting for CI to finish Pending BuildFarm For PRs that have been reviewed but pending a push through our buildfarm labels Mar 12, 2021
@stuartarchibald

Copy link
Copy Markdown
Contributor

Leaving the waiting-on-reviewer label present as @sklam might want to take a look too.

@stuartarchibald

Copy link
Copy Markdown
Contributor

Buildfarm ID: numba_smoketest_cuda_yaml_26.

@sklam sklam left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Just one nitpick on the baseclass ordering for mixins

Comment thread numba/cuda/codegen.py Outdated


def __init__(self, codegen, name):
class CUDACodeLibrary(CodeLibrary, serialize.ReduceMixin):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Although it won't affect the code here, mixins should go to the left. The MRO goes left to right. If CodeLibrary defines __reduce__, it will take priority over the one in the ReduceMixin.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Many thanks - I've swapped these round now.

@sklam sklam left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Thanks for the patch!

@sklam sklam removed the 4 - Waiting on reviewer Waiting for reviewer to respond to author label Mar 17, 2021
@stuartarchibald

Copy link
Copy Markdown
Contributor

Think this is just waiting for a build farm run now.

@stuartarchibald

Copy link
Copy Markdown
Contributor

Buildfarm ID: numba_smoketest_cuda_yaml_30.

@stuartarchibald

Copy link
Copy Markdown
Contributor

Buildfarm ID: numba_smoketest_cuda_yaml_30.

Passed.

@stuartarchibald stuartarchibald added 5 - Ready to merge Review and testing done, is ready to merge BuildFarm Passed For PRs that have been through the buildfarm and passed and removed 4 - Waiting on CI Review etc done, waiting for CI to finish Pending BuildFarm For PRs that have been reviewed but pending a push through our buildfarm labels Mar 17, 2021
@stuartarchibald
stuartarchibald merged commit 81f7082 into numba:master Mar 17, 2021
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

5 - Ready to merge Review and testing done, is ready to merge BuildFarm Passed For PRs that have been through the buildfarm and passed CUDA CUDA related issue/PR Effort - medium Medium size effort needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants