address some review comments

Author: ksimpson-work

cuda_core/cuda/core/experimental/_linker.py

@@ -86,7 +86,7 @@ class Linker:
 
     __slots__ = ("_handle")
 
-    def __init__(self, options: LinkerOptions, object_codes = None):
+    def __init__(self, *object_codes : ObjectCode, options: LinkerOptions = None):
         self._handle = None
         options = check_or_create_options(LinkerOptions, options, "Linker options")
         self._handle = nvjitlink.create(len(options.formatted_options), options.formatted_options)

cuda_core/tests/test_linker.py

@@ -28,101 +28,88 @@ def compile_ltoir_functions(init_cuda):
     return object_code_a_ltoir, object_code_b_ltoir, object_code_c_ltoir
 
 
-def test_linker_init_valid_options():
-    options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    assert linker.handle is not None
+@pytest.mark.parametrize("options", [
+    LinkerOptions(arch=ARCH),
+    LinkerOptions(arch=ARCH, max_register_count=32),
+    LinkerOptions(arch=ARCH, time=True),
+    LinkerOptions(arch=ARCH, verbose=True),
+    LinkerOptions(arch=ARCH, optimization_level=3),
+    LinkerOptions(arch=ARCH, debug=True),
+    LinkerOptions(arch=ARCH, lineinfo=True),
+    LinkerOptions(arch=ARCH, ftz=True),
+    LinkerOptions(arch=ARCH, prec_div=True),
+    LinkerOptions(arch=ARCH, prec_sqrt=True),
+    LinkerOptions(arch=ARCH, fma=True),
+    LinkerOptions(arch=ARCH, kernels_used=["kernel1"]),
+    LinkerOptions(arch=ARCH, variables_used=["var1"]),
+    LinkerOptions(arch=ARCH, optimize_unused_variables=True),
+    LinkerOptions(arch=ARCH, xptxas=["-v"]),
+    LinkerOptions(arch=ARCH, split_compile=0),
+    LinkerOptions(arch=ARCH, split_compile_extended=1),
+    LinkerOptions(arch=ARCH, jump_table_density=100),
+    LinkerOptions(arch=ARCH, no_cache=True)
+])
+def test_linker_init(compile_ptx_functions, options):
+    linker = Linker(options, object_codes=compile_ptx_functions)
+    object_code = linker.link("cubin")
+    assert isinstance(object_code, ObjectCode)
+
 
 def test_linker_init_invalid_arch():
     options = LinkerOptions(arch=None)
     with pytest.raises(ValueError):
         Linker(options)
 
-def test_linker_init(compile_ptx_functions):
-    combinations = [
-        LinkerOptions(arch=ARCH),
-        LinkerOptions(arch=ARCH, max_register_count=32),
-        LinkerOptions(arch=ARCH, time=True),
-        LinkerOptions(arch=ARCH, verbose=True),
-        LinkerOptions(arch=ARCH, optimization_level=3),
-        LinkerOptions(arch=ARCH, debug=True),
-        LinkerOptions(arch=ARCH, lineinfo=True),
-        LinkerOptions(arch=ARCH, ftz=True),
-        LinkerOptions(arch=ARCH, prec_div=True),
-        LinkerOptions(arch=ARCH, prec_sqrt=True),
-        LinkerOptions(arch=ARCH, fma=True),
-        LinkerOptions(arch=ARCH, kernels_used=["kernel1"]),
-        LinkerOptions(arch=ARCH, variables_used=["var1"]),
-        LinkerOptions(arch=ARCH, optimize_unused_variables=True),
-        LinkerOptions(arch=ARCH, xptxas=["-v"]),
-        LinkerOptions(arch=ARCH, split_compile=0),
-        LinkerOptions(arch=ARCH, split_compile_extended=1),
-        LinkerOptions(arch=ARCH, jump_table_density=100),
-        LinkerOptions(arch=ARCH, no_cache=True)
-    ]
-
-    # Try the combinations, with and without providing object codes to the constructor
-    for i, options in enumerate(combinations):
-        linker = Linker(options, object_codes=compile_ptx_functions)
-        object_code = linker.link("cubin")
-        assert isinstance(object_code, ObjectCode)
 
 def test_linker_add_code_object(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(options)
-    functions = compile_ptx_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ptx_functions:
+        linker.add_code_object(functions)
+
 
 def test_linker_link_ptx(compile_ltoir_functions):
     options = LinkerOptions(arch=ARCH, link_time_optimization=True, ptx=True)
     linker = Linker(options)
-    functions = compile_ltoir_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ltoir_functions:
+        linker.add_code_object(functions)
     linked_code = linker.link("ptx")
     assert isinstance(linked_code, ObjectCode)
 
+
 def test_linker_link_cubin(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(options)
-    functions = compile_ptx_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ptx_functions:
+        linker.add_code_object(functions)
     linked_code = linker.link("cubin")
     assert isinstance(linked_code, ObjectCode)
 
+
 def test_linker_link_invalid_target_type(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(options)
-    functions = compile_ptx_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ptx_functions:
+        linker.add_code_object(functions)
     with pytest.raises(ValueError):
         linker.link("invalid_target")
 
+
 def test_linker_get_error_log(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(options)
-    functions = compile_ptx_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ptx_functions:
+        linker.add_code_object(functions)
     linker.link("cubin")
     log = linker.get_error_log()
     assert isinstance(log, str)
 
+
 def test_linker_get_info_log(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
     linker = Linker(options)
-    functions = compile_ptx_functions
-    linker.add_code_object(functions[0])
-    linker.add_code_object(functions[1])
-    linker.add_code_object(functions[2])
+    for functions in compile_ptx_functions:
+        linker.add_code_object(functions)
     linker.link("cubin")
     log = linker.get_info_log()
     assert isinstance(log, str)