Add compile_all API

[isVoid]

This PR adds compile_all API. Instead of only returning the compiled kernel, it also compiles cuda source files for its external linking files. In the case that external files are not compilable (such as cubin, ptx, LTOIR etc.), they will be passed through into the return list.

Testing of compile_all completely is tough due to combinatorial complexity from the number of arguments. Since it's mainly a passthrough of compile with additional nvrtc pass on external linking files, I reused all existing tests in TestCompile to make sure the API is consistent with previous APIs.

The default arguments of compile_all is the same as compile.

When lineinfo or debug is set to true, the nvrtc flag for either options are set accordingly for external files.

[copy-pr-bot]

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[isVoid]

/ok to test

[isVoid]

/ok to test

[isVoid]

/ok to test

[rparolin]

Does this need to be guarded by TEST_BIN_DIR ?

[isVoid]

It shouldn't be run if TEST_BIN_DIR doesn't exist.

[rparolin]

Does this need to be a hard error? Is forceinline a guarantee or can the compiler still ignore it?

[isVoid]

I imagine this is because there's no way to designate the inline information when you specify output as PTX. The information can only persist via LTOIR.

[gmarkall]

That's correct - you can't inline PTX. This looks like new code on the diff, but this PR is only really moving it from the original compile() implementation.

[rparolin]

Nit: I'd consider renaming lto to do_lto or is_lto to more clearly indicate that its a flag controlling behaviour.

[isVoid]

The name here means whether the output should be ltoir, so is_ltoir is the best fit here IMO.

[rparolin]

The variable is called path_or_obj but the code appears to handle cu and obj use cases. Should the variable be renamed or are you missing the path handling use case?

[isVoid]

The objects in lib._linking_files can be arbitrary linkable code object or paths, among which only cuda source files require compilation before it's feedable to the linker. (That's the implied assumption of return objects for compile_all, everything returning from the API should be passable to linker without additional processing). Therefore we are special casing below, compile them with nvrtc before returning it to user.

[rparolin]

I left a couple comments for you to consider. Looks good overall.

[isVoid]

/ok to test

[isVoid]

/ok to test

[isVoid]

/ok to test

[isVoid]

/ok to test

[isVoid]

/ok to test

[gmarkall]

It's a bit of an oversight that we never passed these flags to NVRTC before - have you noticed if this now enables debugging / profiling of NVRTC-compiled code used in Numba kernels with these changes?

[isVoid]

There was a discussion offline on this. I think this will serve as part of the solution.

[gmarkall]

Because this docstring will appear in the docs, I think it needs some modification to fit into the context of the documentation on the compilation APIs:

Suggested change

	"""Return a list of PTX/LTO-IR for kernel as well as external functions depended by the kernel.
	If external functions are cuda c/c++ source, they will be compiled with NVRTC. The output code
	kind is the same as the `output` parameter.
	Otherwise, they will be passed through to the return list.
	"""
	"""Similar to ``compile()``, but returns a list of PTX codes/LTO-IRs for
	the compiled function and the external functions it depends on.
	If external functions are CUDA C++ source, they will be compiled with
	NVRTC. Other kinds of external function code (e.g. cubins, fatbins, etc.)
	will be added directly to the return list. The output code kind is
	determined by the ``output`` parameter, and defaults to ``"ltoir"``.
	"""

This change is aimed at pointing out the differences between compile() and compile_all().

The docs change to include this is also required - this could be done by adding

.. autofunction:: numba.cuda.compile_all

around line 126 of docs/source/reference/host.rst.

[gmarkall]

It seems a bit strange that the property isn't just called linking_files, but maybe I'm not understanding some nuance here.

[gmarkall]

I see this isn't used - perhaps it was meant for use in https://github.com/NVIDIA/numba-cuda/pull/484/files#diff-f40f9ce54f6cd734dff0d6436111d1e6f3a0be43a8d98f489bdd6bf726ba813aR1036 ?

[isVoid]

Right, I believed at one point I wanted to use this property to normalize object and paths into LinkableCode. But now this logic lives in LinkableCode factory functions.

[gmarkall]

The changes here work - I think I would have been inclined to use a pattern like:

    def _test_device_function(self, compile_function):
        def add(x, y):
            return x + y

        args = (float32, float32)

        ptx, resty = compile_function(add, args, device=True)

        # Device functions take a func_retval parameter for storing the
        # returned value in by reference
        self.assertIn("func_retval", ptx)
        # .visible .func is used to denote a device function
        self.assertIn(".visible .func", ptx)
        # .visible .entry would denote the presence of a global function
        self.assertNotIn(".visible .entry", ptx)
        # Inferred return type as expected?
        self.assertEqual(resty, float32)

        # Check that function's output matches signature
        sig_int32 = int32(int32, int32)
        ptx, resty = compile_ptx(add, sig_int32, device=True)
        self.assertEqual(resty, int32)

        sig_int16 = int16(int16, int16)
        ptx, resty = compile_ptx(add, sig_int16, device=True)
        self.assertEqual(resty, int16)
        # Using string as signature
        sig_string = "uint32(uint32, uint32)"
        ptx, resty = compile_ptx(add, sig_string, device=True)
        self.assertEqual(resty, uint32)

    def test_device_function(self):
        self._test_device_function(compile_ptx)

    def test_device_function_all(self):
        def compile_all_wrapper(*args, **kwargs):
            kwargs["abi"] = "c"
            kwargs["output"] = "ptx"
            ptx_list, resty = compile_all(*args, **kwargs)
            self.assertEqual(len(ptx_list), 1)
            return ptx_list[0], resty

        self._test_device_function(compile_all_wrapper)

to test compile_all() without duplicating all the test code, but I think this needn't necessarily be changed.

(To make the changes from the original main code upstream a bit clearer, they are:

• Rename test_device_function() to _test_device_function() and add the compile_function parameter
• Replace the compile_ptx() call with a call to compile_function()
• Add the test_device_function() and test_device_function_all() methods, where test_device_function_all() emulates the API of compile_ptx() with a wrapper around compile_all().

)

[gmarkall]

Thanks for the PR. I think this looks good in general. We need some changes for documentation (which are noted on the diff):

• To add the compile_all() function to the reference documentation
• To modify compile_all()'s docstring so it reads well in the context of the rest of the documentation.

I made a comment about the duplications introduced in the tests for compile*() and compile_all(), but I think that needn't necessarily be changed.

[gmarkall]

/ok to test