Repository: numba/numba
Pull request: CUDA: Fix and deprecate `inspect_ptx()`, fix NVVM option setup for device functions
This fixes issues with device functions:

- `inspect_ptx` was broken by PR #6735 (commit 11a7397). The device  function should use the code library to obtain the PTX instead of attempting to compile directly itself.
- Optimization and debug options were not passed to NVVM when compiling device functions. These options are now set up and passed to `compile_cuda` when a device function is compiled.

This commit also deprecates the use of `inspect_ptx` - the preferred API for compiling Python to PTX is to use the `compile_ptx` function instead. The `inspect_ptx()` method had a couple of issues:

- Allowing it to receive a dict of options for NVVM enabled it to produce different PTX to what might actually have been compiled so far, as highlighted by the fact that optimization and debug flags failed to be passed to NVVM for the `compile()` method, but not in the `inspect_ptx()` method.
- The NVVM options dict kwarg was unsafely initialized anyway (`nvvm_options={}`) - See issue #5811.
- It returned encoded bytes, which differs from similar APIs (e.g. `inspect_llvm()`, and the Dispatcher's `inspect_asm()`, which return `str` instead.

As there is no easy way to correctly and consistently pass NVVM options through `inspect_ptx()`, a warning is emitted stating that these are ignored if a user passes any options in.

Some tests are also added, to ensure that it works correctly, and warns the user appropriately.

Fixes #6950.

Problem statement:
Numba master breaks cuDF
Failures with cuDF branch-0.20 with Numba master:

```
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[float32] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[float64] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[int16] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[int32] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[int64] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[uint16] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[uint32] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[uint64] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
FAILED cudf/tests/test_udf_binops.py::test_generic_ptx[uint8] - numba.cuda.cudadrv.error.NvvmError: Failed to compile
```

because of e.g.:

```
__________________________________________________________________________________________________________________________________________________ test_generic_ptx[uint8] __________________________________________________________________________________________________________________________________________________

dtype = 'uint8'

    @pytest.mark.parametrize(
        "dtype", sorted(list(dtypeutils.NUMERIC_TYPES - {"int8"}))
    )
    def test_generic_ptx(dtype):
    
        size = 500
    
        lhs_arr = np.random.random(size).astype(dtype)
        lhs_col = Series(lhs_arr)._column
    
        rhs_arr = np.random.random(size).astype(dtype)
        rhs_col = Series(rhs_arr)._column
    
        @numba.cuda.jit(device=True)
        def generic_function(a, b):
            return a ** 3 + b
    
        nb_type = numpy_support.from_dtype(np.dtype(dtype))
        type_signature = (nb_type, nb_type)
    
        result = generic_function.compile(type_signature)
>       ptx = generic_function.inspect_ptx(type_signature)

cudf/tests/test_udf_binops.py:41: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../../numba/numba/cuda/compiler.py:319: in inspect_ptx
    ptx = nvvm.llvm_to_ptx(llvmir, opt=opt, arch=arch, **nvvm_options)
../../../numba/numba/cuda/cudadrv/nvvm.py:708: in llvm_to_ptx
    ptx = cu.compile(**opts)
../../../numba/numba/cuda/cudadrv/nvvm.py:294: in compile
    self._try_error(err, 'Failed to compile\n')
../../../numba/numba/cuda/cudadrv/nvvm.py:312: in _try_error
    self.driver.check_error(err, "%s\n%s" % (msg, self.get_log()))
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <numba.cuda.cudadrv.nvvm.NVVM object at 0x7f31684ae910>, error = 9, msg = "Failed to compile\n\n<unnamed> (85, 24): parse invalid redefinition of function '_ZN5numba7cpython7numbers14int_power_impl12$3clocals$3e14int_power$2416Eyx'", exit = False

    def check_error(self, error, msg, exit=False):
        if error:
            exc = NvvmError(msg, RESULT_CODE_NAMES[error])
            if exit:
                print(exc)
                sys.exit(1)
            else:
>               raise exc
E               numba.cuda.cudadrv.error.NvvmError: Failed to compile
E               
E               <unnamed> (85, 24): parse invalid redefinition of function '_ZN5numba7cpython7numbers14int_power_impl12$3clocals$3e14int_power$2416Eyx'
E               NVVM_ERROR_COMPILATION

../../../numba/numba/cuda/cudadrv/nvvm.py:192: NvvmError
```

Minimal reproducer:

```python
from numba import cuda, types

@cuda.jit(device=True)
def generic_function(a, b):
    return a ** 3 + b

type_signature = (types.int8, types.int8)
generic_function.compile(type_signature)
generic_function.inspect_ptx(type_signature)
```

produces:

```
$ python repro.py 
Traceback (most recent call last):
  File "repro.py", line 11, in <module>
    generic_function.inspect_ptx(type_signature)
  File "/home/gmarkall/numbadev/numba/numba/cuda/compiler.py", line 319, in inspect_ptx
    ptx = nvvm.llvm_to_ptx(llvmir, opt=opt, arch=arch, **nvvm_options)
  File "/home/gmarkall/numbadev/numba/numba/cuda/cudadrv/nvvm.py", line 708, in llvm_to_ptx
    ptx = cu.compile(**opts)
  File "/home/gmarkall/numbadev/numba/numba/cuda/cudadrv/nvvm.py", line 294, in compile
    self._try_error(err, 'Failed to compile\n')
  File "/home/gmarkall/numbadev/numba/numba/cuda/cudadrv/nvvm.py", line 312, in _try_error
    self.driver.check_error(err, "%s\n%s" % (msg, self.get_log()))
  File "/home/gmarkall/numbadev/numba/numba/cuda/cudadrv/nvvm.py", line 192, in check_error
    raise exc
numba.cuda.cudadrv.error.NvvmError: Failed to compile

<unnamed> (85, 24): parse invalid redefinition of function '_ZN5numba7cpython7numbers14int_power_impl12$3clocals$3e13int_power$242Exx'
NVVM_ERROR_COMPILATION
```

Hints from issue discussion:
11a739762600a1d6f1ab75988876a48c015c391f is responsible for this.
This is from PR #6735.
Actually compilation is not needed to trigger this, it's just a bug in `inspect_ptx`:

```python
from numba import cuda, int8

@cuda.jit(device=True)
def generic_function(a):
    return a ** 3

generic_function.inspect_ptx((int8,))
```

will produce the same error.
The issue is that device functions generate PTX by calling `llvm_to_ptx` on their IR returned by their own `inspect_llvm()` method: https://github.com/numba/numba/blob/39271156a52c58ca18b15aebcb1c85e4a07e49ed/numba/cuda/compiler.py#L312-L320

The `inspect_llvm()` method concatenates the IR for all modules so that the user gets to see it all - however, this isn't in the right form for compilation any more - instead the `inspect_ptx()` function should use the `CompileResult`'s code library to get the PTX.
Starter patch:

```diff
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
index 41f3058a6..07f1424d2 100644
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -290,12 +290,8 @@ class DeviceFunctionTemplate(serialize.ReduceMixin):
         -------
         llvmir : str
         """
-        # Force a compilation to occur if none has yet - this can be needed if
-        # the user attempts to inspect LLVM IR or PTX before the function has
-        # been called for the given arguments from a jitted kernel.
-        self.compile(args)
-        cres = self._compileinfos[args]
-        return "\n\n".join([str(mod) for mod in cres.library.modules])
+        modules = self.compile(args).library.modules
+        return "\n\n".join([str(mod) for mod in modules])
 
     def inspect_ptx(self, args, nvvm_options={}):
         """Returns the PTX compiled for *args* for the currently active GPU
@@ -309,15 +305,7 @@ class DeviceFunctionTemplate(serialize.ReduceMixin):
         -------
         ptx : bytes
         """
-        llvmir = self.inspect_llvm(args)
-        # Make PTX
-        cuctx = get_context()
-        device = cuctx.device
-        cc = device.compute_capability
-        arch = nvvm.get_arch_option(*cc)
-        opt = 3 if self.opt else 0
-        ptx = nvvm.llvm_to_ptx(llvmir, opt=opt, arch=arch, **nvvm_options)
-        return ptx
+        return self.compile(args).library.get_asm_str()
```

Needs a test case adding and also fixing up tests that decode the PTX (or alternatively encoding the result first).

Commit message: Add tests for inspect_ptx

Patch to review:
diff --git a/numba/cuda/compiler.py b/numba/cuda/compiler.py
index 41f3058a697..222bf247228 100644
--- a/numba/cuda/compiler.py
+++ b/numba/cuda/compiler.py
@@ -25,7 +25,7 @@
 import numba
 from .cudadrv.devices import get_context
 from .cudadrv.libs import get_cudalib
-from .cudadrv import nvvm, driver
+from .cudadrv import driver
 from .errors import missing_launch_config_msg, normalize_kernel_dimensions
 from .api import get_current_device
 from .args import wrap_arg
@@ -260,8 +260,13 @@ def compile(self, args):
         Returns the `CompileResult`.
         """
         if args not in self._compileinfos:
+            nvvm_options = {
+                'opt': 3 if self.opt else 0,
+                'debug': self.debug,
+            }
+
             cres = compile_cuda(self.py_func, None, args, debug=self.debug,
-                                inline=self.inline)
+                                inline=self.inline, nvvm_options=nvvm_options)
             first_definition = not self._compileinfos
             self._compileinfos[args] = cres
             libs = [cres.library]
@@ -290,12 +295,8 @@ def inspect_llvm(self, args):
         -------
         llvmir : str
         """
-        # Force a compilation to occur if none has yet - this can be needed if
-        # the user attempts to inspect LLVM IR or PTX before the function has
-        # been called for the given arguments from a jitted kernel.
-        self.compile(args)
-        cres = self._compileinfos[args]
-        return "\n\n".join([str(mod) for mod in cres.library.modules])
+        modules = self.compile(args).library.modules
+        return "\n\n".join([str(mod) for mod in modules])
 
     def inspect_ptx(self, args, nvvm_options={}):
         """Returns the PTX compiled for *args* for the currently active GPU
@@ -309,15 +310,15 @@ def inspect_ptx(self, args, nvvm_options={}):
         -------
         ptx : bytes
         """
-        llvmir = self.inspect_llvm(args)
-        # Make PTX
-        cuctx = get_context()
-        device = cuctx.device
-        cc = device.compute_capability
-        arch = nvvm.get_arch_option(*cc)
-        opt = 3 if self.opt else 0
-        ptx = nvvm.llvm_to_ptx(llvmir, opt=opt, arch=arch, **nvvm_options)
-        return ptx
+        msg = ('inspect_ptx for device functions is deprecated. Use '
+               'compile_ptx instead.')
+        warn(msg, category=NumbaDeprecationWarning)
+
+        if nvvm_options:
+            msg = ('nvvm_options are ignored. Use compile_ptx if you want to '
+                   'set NVVM options.')
+            warn(msg, category=NumbaDeprecationWarning)
+        return self.compile(args).library.get_asm_str().encode()
 
 
 def compile_device_template(pyfunc, debug=False, inline=False, opt=True):
diff --git a/numba/cuda/tests/cudapy/test_inspect.py b/numba/cuda/tests/cudapy/test_inspect.py
index d5412a523e3..b7fb5bdc4d5 100644
--- a/numba/cuda/tests/cudapy/test_inspect.py
+++ b/numba/cuda/tests/cudapy/test_inspect.py
@@ -1,4 +1,5 @@
 import numpy as np
+import warnings
 
 from io import StringIO
 from numba import cuda, float32, float64, int32, intp
@@ -204,6 +205,49 @@ def f(x):
         msg = 'ptx will always return a dict in future'
         self.assertIn(msg, str(warns.warnings[0]))
 
+    def test_inspect_ptx(self):
+        # Ensure that inspect_ptx works - see issue #6950
+        @cuda.jit(device=True)
+        def use_power(x):
+            return x ** 3
+
+        ptx = use_power.inspect_ptx((int32,)).decode()
+
+        # A simple sanity check - look for multiplications used to implement
+        # `x ** 3` as `x * x * x`.
+        self.assertIn('mul.wide.s32', ptx)
+        self.assertIn('mul.lo.s64', ptx)
+
+    def test_inspect_ptx_deprecation(self):
+        # Make sure we warn the user that inspect_ptx is deprecated and tell
+        # them to use compile_ptx instead.
+        @cuda.jit(device=True)
+        def use_power(x):
+            return x ** 3
+
+        with warnings.catch_warnings(record=True) as w:
+            use_power.inspect_ptx((int32,))
+
+        self.assertEqual(w[0].category, NumbaDeprecationWarning)
+        self.assertIn('Use compile_ptx instead', str(w[0].message))
+
+    def test_inspect_ptx_nvvm_options_deprecation(self):
+        # Ensure that we warn the user of the deprecation of inspect_ptx and
+        # that nvvm_options are ignored if they pass nvvm_options to
+        # inspect_ptx.
+        @cuda.jit(device=True)
+        def use_power(x):
+            return x ** 3
+
+        with warnings.catch_warnings(record=True) as w:
+            nvvm_options = {'opt': 0}
+            use_power.inspect_ptx((int32,), nvvm_options=nvvm_options)
+
+        self.assertEqual(w[0].category, NumbaDeprecationWarning)
+        self.assertIn('Use compile_ptx instead', str(w[0].message))
+        self.assertEqual(w[1].category, NumbaDeprecationWarning)
+        self.assertIn('nvvm_options are ignored', str(w[1].message))
+
 
 if __name__ == '__main__':
     unittest.main()


Code review comments:
diff_hunk: @@ -260,8 +260,13 @@ def compile(self, args):
         Returns the `CompileResult`.
         """
         if args not in self._compileinfos:
+            nvvm_options = {
+                'opt': 3 if self.opt else 0,
+                'debug': self.debug,
+            }
+
             cres = compile_cuda(self.py_func, None, args, debug=self.debug,
-                                inline=self.inline)
+                                inline=self.inline, nvvm_options=nvvm_options)
path: numba/cuda/compiler.py
text: @user1:
Whilst this replicates the `opt` behaviour in `inspect_ptx` is this now changing the default optimisation behaviour for callers of this method? Was it just being ignored before/something else would override it (direct inlining)?

@author:
It's fixing the optimization behaviour - even if you passed `opt=False`, it would still optimize because we passed nothing to NVVM and it optimizes by default. Worse, the test case that checked this (`numba.cuda.tests.cudapy.test_optimization.TestOptimization.test_device_noopt`) was erroneously passing because of the incorrect implementation of `inspect_ptx`! Once I fixed the behaviour of `inspect_ptx()`, this test started failing, highlighting the need for the change on the line these comments pertain to.

@user1:
Thanks for the explanation. Seems like this is doing the correct behaviour now!