support loading PTX in RawModule
@@ -205,19 +205,20 @@ cdef class RawModule:
"""User-defined custom module.
This class can be used to either compile raw CUDA sources or load CUDA
- modules (\\*.cubin). This class is useful when a number of CUDA kernels in
- the same source need to be retrieved.
+ modules (\\*.cubin, \\*.ptx). This class is useful when a number of CUDA
+ kernels in the same source need to be retrieved.
For the former case, the CUDA source code is compiled when initializing a
new instance of this class, and the kernels can be retrieved by calling
:meth:`get_function`, which will return an instance of :class:`RawKernel`.
(Same as in :class:`RawKernel`, the generated binary is also cached.)
- For the latter case, an existing CUDA binary (\\*.cubin) can be loaded by
- providing its path, and kernels therein can be retrieved similarly.
+ For the latter case, an existing CUDA binary (\\*.cubin) or a PTX file can
+ be loaded by providing its path, and kernels therein can be retrieved
+ similarly.
Args:
- code_or_path (str): CUDA source code or path to cubin.
+ code_or_path (str): CUDA source code or path to cubin/ptx.
options (tuple of str): Compiler options passed to the backend (NVRTC
or NVCC). For details, see
https://docs.nvidia.com/cuda/nvrtc/index.html#group__options or
@@ -238,7 +239,7 @@ cdef class RawModule:
if isinstance(backend, six.binary_type):
backend = backend.decode('UTF-8')
- if code_or_path.endswith('.cubin'):
+ if code_or_path.endswith('.cubin') or code_or_path.endswith('.ptx'):
path = code_or_path
self.code = None
self.cubin_path = path
