Skip to content

Commit b5b6b98

Browse files
committed
update the nvjitlink binding test to use multiple architectures and nvrtc to generate the ltoir
1 parent 88d3503 commit b5b6b98

File tree

1 file changed

+86
-58
lines changed

1 file changed

+86
-58
lines changed

cuda_bindings/tests/test_nvjitlink.py

Lines changed: 86 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,21 @@
55
import pytest
66

77
from cuda.bindings import nvjitlink
8+
from cuda import nvrtc
89

10+
# Establish a handful of compatible architectures and PTX versions to test with
11+
ARCHITECTURES = ["sm_60", "sm_75", "sm_80", "sm_90"]
12+
PTX_VERSIONS = ["5.0", "6.4", "7.0", "8.5"]
913

10-
ptx_kernel = """
11-
.version 8.5
12-
.target sm_90
14+
def ptx_header(version, arch):
15+
return f"""
16+
.version {version}
17+
.target {arch}
1318
.address_size 64
19+
"""
1420

21+
22+
ptx_kernel = """
1523
.visible .entry _Z6kernelPi(
1624
.param .u64 _Z6kernelPi_param_0
1725
)
@@ -29,20 +37,33 @@
2937
"""
3038

3139
minimal_ptx_kernel = """
32-
.version 8.5
33-
.target sm_90
34-
.address_size 64
35-
3640
.func _MinimalKernel()
3741
{
3842
ret;
3943
}
4044
"""
4145

42-
pre_compiled_empty_kernel_ltoir = b'\xedCN\x7f\x01>\x02T\x03\x01\x07\x00\x18\x00\x01\x00(\x00\x00\x00,\x00\x00\x00\x01\x00\x08\x02\x02\x00\x02\x00c\x00\x00\x00\x00\x00\x00\x00\xc8}O$\xd6\x03U\xee\xacI\x85\xd8\x00\x05PST\x13c]\xe8\x1a&\xa5\xb7\x81\xdb\xb2\x08:\x97\x12\xb0\xf0\x02\x8f\x1d\xda\xf3\x98\xd2L\xd0\xfcCW\x9e\xb1R\xffB%\xc7\xbeS\xe9\x13G_\xa1;\x04\xa8\xb4\x8f\xab\x01\xadTk\x1e\xa5\xb8\xef\xb5\x01\x0b\xfaM,\xaf\x7f\xaa\xa6/\xfc\x13\x80b\xb7\x08>&\xc0I}\xe72|\xf5Q\x7f|\x17\xd5\xf3\xefY\x90\xce\x90z\xae\xf5\x12\x94!P@=\xe3\x1e\xb8\x19\xd4.*\xb5I\xeaA)\xf4\x11\t\xedg$d\x03Bf\x90\xd7\xfc\x81%\xb8V\x8bS<N\xddao\x1c\t\xd3\xbf8\xd4d\x01[Z\x19\xb5\x99\xfc\x87\x81\xcf\xb2\xe0ZH.\x8a\xf2\xafq:\x9e\xdb<\xa1\xe1R\x0c\xf9X\xbe_\r\xb9F\xa8\x9dg\x18\xac\xeb\x1c\xc1/\x03\x84\xc2tVp\x9b5\xd9V\xb7\xfcGE)\xc1\xafF\xfe$&&\x94\xe7DM-\x87\x88n\xa7\xb7\x93|\xd5@\x84\xcfl\x04{\x13(\xaf\xf5\xdc\xf0,\xd3\x1b\x80\xc3\xcbc\x97\xd3UB\xa2\xebm\x8c\x9f`\xe7\x17R\x191y\xbd\xf4\x08\xc0\xf7\x97\x91P\xbbzzv\xe5K\x82\x89}\xe3\xc7\xab\x98G\xd8\xe5\xba\xbe\xe8\xe7R,\xa9\xd3\xc8Gl\xe5\xc9\xb7$\xff\xe7\xefB\xa5\x81\xcc\xc5<\xce/\x94fg\x18\xf5\xa9\\\xe3\xef\xae\xd2\xbc\xd7#\x1d\xbb6\xea\xd2cH\xdf\xf5\x9b\xb2\xde\xb5\xea\xd7\x0c\x8c7S\x07\x86:\x94W\xcaz\x1f<H!B\x08s\x83rqo<G\x8b\xa0\xfc\xe5\x12mi\xf2\x01X<\x15\x93L\xa1s\x00\xa2\xe5\\\xa9\x1b\x9c\xdf\xd5&_\xc0\xf2l1_Z*r\xf3\x91\xb3\xc2+\xa9\xfcFrB@>e\x10\x16\xd2\xbdC*\xcb-%Gl}\xe0\x1fNv\x7fmx%E~\x1d!h\xd8\xb68\xc1|\xb1\xcf\x0e\x01.^\x9f7\xffF\x9f\xd8Fz\x91\x99>\xfb\x8aN\xa2\x9f\xd4S\xdc|\xeb\x9c\x90\xca\x1c\x1f\xac\xee`\r\xbaRm\x99\xe3=\x17\xa9wB3s_\xa1\x13\xe0@,\x9a\x94\x8a$\xfc\x83\xfb\xae\xde\xa0A\xbc\xf8\r@\xe7\xe0JS\xdeN\xe4\xbc\xce\x91\xa5E\xf322y\x13\\\xc5\x7f\xec\xb0\xb3\x9a\xac\xfd\xe3\r\xaf\x86\x9f\x95\x1c{$\x95\xce\xd1^\x94\xa6-\xf8\xea\xaa\x0f\xfe\x17@\x19<e\x07\x0f\xb9\xfa\xd8\x15\x10\x89^\xad=\x11\xb2\xbb\x8b\x14\xc9\xe3v$\xac0\xe8Y\x8a\x9e\xb3\xed\xf8\x99\xb6\xd5\x93\x00\xd4\x13\x03\xeb~\xff\xcb0\xd8\x1a:\x19[\x14\xcd\x95\xeb\xed\xce\x98\x03J\x16\x99\x0f\x1f\x99h\xbfE\xcc\xcf\xd9q\x1b\xab\x19\xce^\x1f\xa1\xbc,\\\xe1\xa0D\xc9\xa2\xfd\x96\xfa\xbc\x9ej\xb5\xd05\x85\x92.\xe1u\xa5\xfc\x83\xa7\xe6q\xfc1\xce\xa9\x7f@\xd3eq\x9d\x93\xa4\x9a\xef\xe4\xb7\xf2\xa4\x8c\x13\x8e\xdf\x8dT\x9b\xda\x02$\x93\'Ek\x16\x9d\xcc+\xa8\x98E\xab\xadf\xa1[\xbe\x91\xfb\xec\xaa7\xe1\xa5\xc1\xb8\xfd\xc9\xbc\xbc\x1fr\xc9q\xd7\xcf\x99\x0e\x93\xc3\xdfp\xed\xcf\xe3\x0f\xbd\xf3\x95\x04\xf3\x7f\xbe?\x0b~\x15\xe3\xfe\xba\x80k`\xd4\xe5\xd2\x1e["D6?<\xe7\xb1\x16\x96\x7f\x9cXFB\tK\xf7\x96\xad\xb0cg\x1f\xc4q\xeb})\x8d\x94\xa3\xd0\x0fo\xcb\xc7\x0ewq\xd0\xe0\x10\xe3\xd1F\x80B\xec\x8c\x87AI/\xd0\xf9\xc9\r\x81\xf4\xd4\xb38\x01\x06Y\xb3P\xa7Q\xf0\xe5\x17\x93\xa3\xc880\xc9z1\xf9R\x94\x80\x08\x1c\xfd\xe3\xc1\xac_L\xd5\xb4s\x1b^t8F\x13A\xb0A\xe86q{\xa1\xc3\x11\xf8\x9c\x16\x87\xe1c\xba\x8a\xe0\xa8Nn\x9d\xa05\x1fKG/"\xa6\xff\x9a\x85)\x04r\xb7\xe4n\x91\xc4\x00I\xa7l\x1fv\xb8\x82j\xa7\xfb\x89\xd3y;\xe2\xa2(\xd6\xb3F\xaaD\x0b\xdfA\x05\x07\x02Z\xe13\xc79\x8e\xb1\xd4:L@\xf0\x0bT\x96\xd0\xddh%b\x9a`l\xebt\x96ht"\x8e\xe0_\xd1\xc5\x96F\xe5}Z\x02\xdby\xa2O\x99\xe8m\x00\x9ag\n1\xae\xea7\xc2L>\xee\x9c\xf0nt\xf3$\xa5s\xc2Dn\xa8\xd6\x85\xd8\xbc`7\x92$w\x9b\x85r\x84\xb0\x0b \x01\xd0C\xdc\xe4Y\x90o\x018o+p:\x1e\x9b(\x15\x8b\xa81\x8al5\xb4\x04\x0e\xcbt1\xae\xfe\x1f\x1cDJ\xb7\'\x14\xa2NJ\x0cn(e\xfc\xa5\xb0\x8d\xc8PlI\x17$\r\xaa,$\x82\x0b\x07\x82\xe7\xf0\x85\x9b\x03\x85\xa4\x9fZ\r\x1f\x18\x95\xb3\xe5/I\xbdoov\xc2l\x8d\x0f\x05If\x94\xb7y*\x1b\r\x80\xfa\xce\xe08]\x98\xb2{\xcb\xa6\xd7\x83I\x84\x8f \r\x93\\8$\x80\xed\xf4}\x1c\xd2\x86\xcd~S\x10b!\xd4\x05\xd5\x80I\x97\x8c\xeaC\xbf\x88)\xcd\xb9R\xc6A\xc2Y\x03j\x89\x87\x16\xfb\xfb\x1f\xdf\x02\xe7"0\xfc\xe7\n\x01x\xa25-\xc7\xbd\xd06\x18\x8c\xd2\x9dy\x19\xd9\xee\x1a\x15\xc0F\x7f\xc2/\xb0\x01\x0b\x14?U%\n\x0bdQj\xcf@\xbc\x17\xd7\x8eH#o\x1b>\xc9\x97\x9f6\x0e\x9b]\xdeF\x00:\xc40@@9\x9b\xac8\x17Z\xa29\xcd\xc2g\xb4(1\xe4\xd2Gk\x7f>IP\xdb(\x12\x13&W\xb2\xef\x96\x95\x98\x1d\xba\x9d\xc3\x8f5\x94\xdc\xf2\x02\xf5\xdbfc\x00\xca\xca \x18\xcb#\x05@\x95~\x0b\xca\xea\xd6\xe0\x04\x99N\xd0\xee?Br\x81\x96c\x9d\xa0Y=$\x10\xb6B\xf7>\xa2\xfc|\xa4\xac\x0b\xdd\xd0S\x1d\xc7\x82r~\x1b\xbcE\x89\xf2\xa2\xfeQ\t}sr\xd8V^$Ba\r\x1f\x9b\xf9\xf5\xa2\x0b\x8d\xc0p\xec\xe9\x87\xcf\xfc\xbc\xc9)\x95P\x88bl7\xcf4e\xfd#c/(4\xf2F\xb1\x81M\xb3CSPN\x12*\x06Ci<\xf9\xcd\xdc\xb2\xebp7\xd3\xae\xea\x99\xf6\xba(\xf8\x90\xb5\x0cV\xfbg[VLee\xff\xdb\xa2\x10C\x10\xe0\x97\x99\xf1\x0f\xa4e\x14\xb8\xdb\xd8\x81\xb5\xef\x1c\x80\x98\xe1c\xd8\xe3\xd9\x10o\xe1k!a\xf9.H\xaco\x88N\xb7([\\3A\xc2`\xce;\xa4\x1c\x05\t\x0e6\xf0\xc0\xd2 m\xbco\x0b\xc0\xe3}\xe2\x97\xd0h\x0b\xc4\xa08S\xe9\xfb(\xe04nAV(\xbb!\xa0Og\xc5"\xa0y\x00\x82\xa8\x81\x1e\xdd\x17\xc3\x9e\xa6S6"0\x1d\x83\xbbvF\x03\xbbH\x96\xde\x08\xba\x9f\x93\xe2\xabbr\xc9\xcc\x0f!I\x9b\xffS\x17ErH\xc7\n\x02:}\x92\xf7\xd5\xbe\xb2\x13\xa2\xa2\xd5\x06J\xc4\xddA\x9b\xb4\xbb\xe6\xea\xa9\xc8IuS11e\x80\xbd\x85\xeaP\x11}\xc8A0\x8eL\x86\xf7\xbd\xd3\xb3\xbd^\xff\xc7\xf6\x852\xff\xa7n\xb1\x82\xef\x17\x99\xf4\xe6\xa7bj+\xf8\x90\xe8\x95q\xdc\n\x9eI\x04\x8f\xb6\xa9\x0b\x84\x98n\xc2fNYu4\x83\tk\xf9F\x83\xda\xc8@9\xee\x01P\x16\xb1A\xa0\x05\x9d\xb5\xd8K\x95M\xeb\xd6\xeb\x0b\xb0\x8b\x82\xbd\xdf\xf4\xb6\tR\xfaD\xa5\xddW8E=e\x880^\x89\xcf\xcf\xa1\xbd0;\x86\x85\xa5\xc5\xab|[ \x99\xaeW\xbd\x0cF\xa9\x9e\xd2\x85MW\xa0H\x0f\xda\xccvZ\xb7hh\x022>\x19"\xb3I\xa2.K\x94G\xbf\xe0I\xe7U\xf35\xe9iP\x13I\x8e\xad\xef\xa0\x89\x997{}6\xd7\x07]\xab\x07t\x9c\xcb\x8c\x87\xe7\x8fh J\xa0L\xd0\xde\xa9MCyK0:\xbdH\x90d$\xad\xd3\xe0\xc5\x7f\xc7\xb3\x94WV^\xfa\x82E\xe1\xa14\xec&U-~\x0e\x7fJJ\xde\x0c\xa4\x19\x9f\xbf,\n\xa0\xd77,\n!\xb6\xb9\x90%\xa4\xd6\x93AI\x0ew*\xd2 M-,\x00\x1e\tbl\xb3E\x08\xd4\x18\x01\xf5W\xd6\x15\xb7\x1dh\xb3\xec\xa9\xf3\xd1\xe6g\xa9\'\xef\xa3%Y\x16\x17m\xf3Q\x18X\xca\x1d\xda8\x96\xad\xcf\xf4\x82JQ\x9e@\x844V\xb5\xfd\xa5\xe8\xbd\x00\xb4\x98\x10!\xec\xd5`\xb8\xcex\xb8\x9e\x96\xb7\xed\x06;\xcb\xce\x11\xea\x8d|]to-{\x9d\x8f\x91n\xd9\x01c%C\xc05\xac\xd9$wf\xbe<\x10\xd9\x12\xc0\x8c\xd4\xf3\x17\xf1\xa9d<\xa04\xa9c\xf2@\x1d\x9du:\xf0\xb4\n\xd1\x9c*\x9ea*\x82\x1b\x1e*\xcc\xbb\xf8\x81\xbf>\xcb\xabA\x8c\xd8\xd8\x95F\xacWG]\xd1\xf8\xd5\xa4{\xdd\xca@\xde\xac\xff0\xf5\xea\xfaE|\xc9\xe8\xf6C\x96=\x16\xb4\xec\xf5\xf2\xc1\xc6\xa6\xd8\xc5\xb8#!\x10[<\xfdB~\x02\x93T\xc6\xba\x8aT\xf7\x10\xb9\xe5 \xd6\xe0\xb8-`f1@\xecB\x87\xd6&p\x922\xc0\x13\xa9u:<\rr\xbd\xd9;K\xc6n\xaaRt#\xfe\x08bb\x16\x9c^w\x13\xdbpg\'\x16\x7f>~\xa2n(\xefI\x1b\xfegK\x86\x8en\x1a\x08\xab\xdfk2\x00s\x12\x1b\xb5\xf6/h\x8f\xeb\x90\xc80 \x01\xf6\xc5Z\xbd\x0f\x9f\xef\xb8\xedt\x06i,q\\G5\xb0\xe1\x8f\x17\xa1\xa9\xb21ql\xc5\x84\xd0\xdc\xc7\xcb|\xe2\x15\xc4\xc4;\xd5\xc8\xb9\x95\xae-e\xc8\x19\xcfy\xeb\x9c^\x0ei\xd8\xea\x9d\x9eg\xb9\xb5g{|?\x16\xbf\xdaU\x12\xf4aI\x0b\xb3Q({\xdd\xeb\x90u\xae28\xcd\xa9\x11-\x12p\x8dW_\x94Z\x96\xbe\x9c\\\xecB\xe7~w\x18?\xfd\xe5q\xa7\xa7\xdc\xc1<S F\x05\xcd\x0c\x81K \x83R\xea\xb7\xfb\x1b\xe5/T\xe3Ix\x06^`\xb22\xeb\x8b*\xad\xd3\x84L\xc8f\xee9\xa8E\xfcjDV^p\xb7\xcb}k\xca$6\x03\xa5_\xb9e\xe9l`\xaa\x86\xe8n\n\xb7\xac\xb6\xbf\xfel\x96\x84\xa3\xcb>\xb4=\x96\x86]\x8d\x15\xf7~e\x159\xa9\xb6\x94Fs\x1b'
46+
ptx_kernel_bytes = [(ptx_header(version, arch) + ptx_kernel).encode('utf-8') for version, arch in zip(PTX_VERSIONS, ARCHITECTURES)]
47+
minimal_ptx_kernel_bytes = [(ptx_header(version, arch) + minimal_ptx_kernel).encode('utf-8') for version, arch in zip(PTX_VERSIONS, ARCHITECTURES)]
48+
49+
50+
#create a valid LTOIR input for testing
51+
def CHECK_NVRTC(err):
52+
if err != nvrtc.nvrtcResult.NVRTC_SUCCESS:
53+
raise RuntimeError('Nvrtc Error: {}'.format(err))
54+
55+
empty_cplusplus_kernel = "__global__ void A() {}"
56+
err, program_handle = nvrtc.nvrtcCreateProgram(empty_cplusplus_kernel.encode(), b"", 0, [], [])
57+
CHECK_NVRTC(err)
58+
nvrtc.nvrtcCompileProgram(program_handle, 1, ["-dlto".encode()])
59+
err, size = nvrtc.nvrtcGetLTOIRSize(program_handle)
60+
CHECK_NVRTC(err)
61+
empty_kernel_ltoir = b" " * size
62+
err, = nvrtc.nvrtcGetLTOIR(program_handle, empty_kernel_ltoir)
63+
CHECK_NVRTC(err)
64+
err, = nvrtc.nvrtcDestroyProgram(program_handle)
65+
CHECK_NVRTC(err)
4366

44-
ptx_kernel_bytes = ptx_kernel.encode('utf-8')
45-
minimal_ptx_kernel_bytes = minimal_ptx_kernel.encode('utf-8')
4667

4768
def test_unrecognized_option_error():
4869
with pytest.raises(nvjitlink.nvJitLinkError, match="ERROR_UNRECOGNIZED_OPTION"):
@@ -55,75 +76,82 @@ def test_invalid_arch_error():
5576

5677

5778
def test_create_and_destroy():
58-
handle = nvjitlink.create(1, ["-arch=sm_90"])
59-
assert handle != 0
60-
nvjitlink.destroy(handle)
79+
for option in ARCHITECTURES:
80+
handle = nvjitlink.create(1, [f"-arch={option}"])
81+
assert handle != 0
82+
nvjitlink.destroy(handle)
6183

6284

6385
def test_complete_empty():
64-
handle = nvjitlink.create(1, ["-arch=sm_90"])
65-
nvjitlink.complete(handle)
66-
nvjitlink.destroy(handle)
86+
for option in ARCHITECTURES:
87+
handle = nvjitlink.create(1, [f"-arch={option}"])
88+
nvjitlink.complete(handle)
89+
nvjitlink.destroy(handle)
6790

6891

6992
def test_add_data():
70-
handle = nvjitlink.create(1, ["-arch=sm_90"])
71-
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_kernel_bytes, len(ptx_kernel_bytes), "test_data")
72-
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, minimal_ptx_kernel_bytes, len(minimal_ptx_kernel_bytes), "minimal_test_data")
73-
nvjitlink.complete(handle)
74-
nvjitlink.destroy(handle)
93+
for option, ptx_bytes in zip(ARCHITECTURES, ptx_kernel_bytes):
94+
handle = nvjitlink.create(1, [f"-arch={option}"])
95+
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_bytes, len(ptx_bytes), "test_data")
96+
nvjitlink.complete(handle)
97+
nvjitlink.destroy(handle)
7598

7699

77100
def test_add_file(tmp_path):
78-
handle = nvjitlink.create(1, ["-arch=sm_90"])
79-
file_path = tmp_path / "test_file.cubin"
80-
file_path.write_bytes(ptx_kernel_bytes)
81-
nvjitlink.add_file(handle, nvjitlink.InputType.ANY, str(file_path))
82-
nvjitlink.complete(handle)
83-
nvjitlink.destroy(handle)
101+
for option, ptx_bytes in zip(ARCHITECTURES, ptx_kernel_bytes):
102+
handle = nvjitlink.create(1, [f"-arch={option}"])
103+
file_path = tmp_path / "test_file.cubin"
104+
file_path.write_bytes(ptx_bytes)
105+
nvjitlink.add_file(handle, nvjitlink.InputType.ANY, str(file_path))
106+
nvjitlink.complete(handle)
107+
nvjitlink.destroy(handle)
84108

85109

86110
def test_get_error_log():
87-
handle = nvjitlink.create(1, ["-arch=sm_90"])
88-
nvjitlink.complete(handle)
89-
log_size = nvjitlink.get_error_log_size(handle)
90-
log = bytearray(log_size)
91-
nvjitlink.get_error_log(handle, log)
92-
assert len(log) == log_size
93-
nvjitlink.destroy(handle)
111+
for option in ARCHITECTURES:
112+
handle = nvjitlink.create(1, [f"-arch={option}"])
113+
nvjitlink.complete(handle)
114+
log_size = nvjitlink.get_error_log_size(handle)
115+
log = bytearray(log_size)
116+
nvjitlink.get_error_log(handle, log)
117+
assert len(log) == log_size
118+
nvjitlink.destroy(handle)
94119

95120

96121
def test_get_info_log():
97-
handle = nvjitlink.create(1, ["-arch=sm_90"])
98-
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_kernel_bytes, len(ptx_kernel_bytes), "test_data")
99-
nvjitlink.complete(handle)
100-
log_size = nvjitlink.get_info_log_size(handle)
101-
log = bytearray(log_size)
102-
nvjitlink.get_info_log(handle, log)
103-
assert len(log) == log_size
104-
nvjitlink.destroy(handle)
122+
for option, ptx_bytes in zip(ARCHITECTURES, ptx_kernel_bytes):
123+
handle = nvjitlink.create(1, [f"-arch={option}"])
124+
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_bytes, len(ptx_bytes), "test_data")
125+
nvjitlink.complete(handle)
126+
log_size = nvjitlink.get_info_log_size(handle)
127+
log = bytearray(log_size)
128+
nvjitlink.get_info_log(handle, log)
129+
assert len(log) == log_size
130+
nvjitlink.destroy(handle)
105131

106132

107133
def test_get_linked_cubin():
108-
handle = nvjitlink.create(1, ["-arch=sm_90"])
109-
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_kernel_bytes, len(ptx_kernel_bytes), "test_data")
110-
nvjitlink.complete(handle)
111-
cubin_size = nvjitlink.get_linked_cubin_size(handle)
112-
cubin = bytearray(cubin_size)
113-
nvjitlink.get_linked_cubin(handle, cubin)
114-
assert len(cubin) == cubin_size
115-
nvjitlink.destroy(handle)
134+
for option, ptx_bytes in zip(ARCHITECTURES, ptx_kernel_bytes):
135+
handle = nvjitlink.create(1, [f"-arch={option}"])
136+
nvjitlink.add_data(handle, nvjitlink.InputType.ANY, ptx_bytes, len(ptx_bytes), "test_data")
137+
nvjitlink.complete(handle)
138+
cubin_size = nvjitlink.get_linked_cubin_size(handle)
139+
cubin = bytearray(cubin_size)
140+
nvjitlink.get_linked_cubin(handle, cubin)
141+
assert len(cubin) == cubin_size
142+
nvjitlink.destroy(handle)
116143

117144

118145
def test_get_linked_ptx():
119-
handle = nvjitlink.create(3, ["-arch=sm_90", "-lto", "-ptx"])
120-
nvjitlink.add_data(handle, nvjitlink.InputType.LTOIR, pre_compiled_empty_kernel_ltoir, len(pre_compiled_empty_kernel_ltoir), "test_data")
121-
nvjitlink.complete(handle)
122-
ptx_size = nvjitlink.get_linked_ptx_size(handle)
123-
ptx = bytearray(ptx_size)
124-
nvjitlink.get_linked_ptx(handle, ptx)
125-
assert len(ptx) == ptx_size
126-
nvjitlink.destroy(handle)
146+
for option in ARCHITECTURES:
147+
handle = nvjitlink.create(3, [f"-arch={option}", "-lto", "-ptx"])
148+
nvjitlink.add_data(handle, nvjitlink.InputType.LTOIR, empty_kernel_ltoir, len(empty_kernel_ltoir), "test_data")
149+
nvjitlink.complete(handle)
150+
ptx_size = nvjitlink.get_linked_ptx_size(handle)
151+
ptx = bytearray(ptx_size)
152+
nvjitlink.get_linked_ptx(handle, ptx)
153+
assert len(ptx) == ptx_size
154+
nvjitlink.destroy(handle)
127155

128156

129157
def test_package_version():

0 commit comments

Comments
 (0)