5
5
import pytest
6
6
7
7
from cuda .bindings import nvjitlink
8
+ from cuda import nvrtc
8
9
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" ]
9
13
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 }
13
18
.address_size 64
19
+ """
14
20
21
+
22
+ ptx_kernel = """
15
23
.visible .entry _Z6kernelPi(
16
24
.param .u64 _Z6kernelPi_param_0
17
25
)
29
37
"""
30
38
31
39
minimal_ptx_kernel = """
32
- .version 8.5
33
- .target sm_90
34
- .address_size 64
35
-
36
40
.func _MinimalKernel()
37
41
{
38
42
ret;
39
43
}
40
44
"""
41
45
42
- pre_compiled_empty_kernel_ltoir = b'\xed CN\x7f \x01 >\x02 T\x03 \x01 \x07 \x00 \x18 \x00 \x01 \x00 (\x00 \x00 \x00 ,\x00 \x00 \x00 \x01 \x00 \x08 \x02 \x02 \x00 \x02 \x00 c\x00 \x00 \x00 \x00 \x00 \x00 \x00 \xc8 }O$\xd6 \x03 U\xee \xac I\x85 \xd8 \x00 \x05 PST\x13 c]\xe8 \x1a &\xa5 \xb7 \x81 \xdb \xb2 \x08 :\x97 \x12 \xb0 \xf0 \x02 \x8f \x1d \xda \xf3 \x98 \xd2 L\xd0 \xfc CW\x9e \xb1 R\xff B%\xc7 \xbe S\xe9 \x13 G_\xa1 ;\x04 \xa8 \xb4 \x8f \xab \x01 \xad Tk\x1e \xa5 \xb8 \xef \xb5 \x01 \x0b \xfa M,\xaf \x7f \xaa \xa6 /\xfc \x13 \x80 b\xb7 \x08 >&\xc0 I}\xe7 2|\xf5 Q\x7f |\x17 \xd5 \xf3 \xef Y\x90 \xce \x90 z\xae \xf5 \x12 \x94 !P@=\xe3 \x1e \xb8 \x19 \xd4 .*\xb5 I\xea A)\xf4 \x11 \t \xed g$d\x03 Bf\x90 \xd7 \xfc \x81 %\xb8 V\x8b S<N\xdd ao\x1c \t \xd3 \xbf 8\xd4 d\x01 [Z\x19 \xb5 \x99 \xfc \x87 \x81 \xcf \xb2 \xe0 ZH.\x8a \xf2 \xaf q:\x9e \xdb <\xa1 \xe1 R\x0c \xf9 X\xbe _\r \xb9 F\xa8 \x9d g\x18 \xac \xeb \x1c \xc1 /\x03 \x84 \xc2 tVp\x9b 5\xd9 V\xb7 \xfc GE)\xc1 \xaf F\xfe $&&\x94 \xe7 DM-\x87 \x88 n\xa7 \xb7 \x93 |\xd5 @\x84 \xcf l\x04 {\x13 (\xaf \xf5 \xdc \xf0 ,\xd3 \x1b \x80 \xc3 \xcb c\x97 \xd3 UB\xa2 \xeb m\x8c \x9f `\xe7 \x17 R\x19 1y\xbd \xf4 \x08 \xc0 \xf7 \x97 \x91 P\xbb zzv\xe5 K\x82 \x89 }\xe3 \xc7 \xab \x98 G\xd8 \xe5 \xba \xbe \xe8 \xe7 R,\xa9 \xd3 \xc8 Gl\xe5 \xc9 \xb7 $\xff \xe7 \xef B\xa5 \x81 \xcc \xc5 <\xce /\x94 fg\x18 \xf5 \xa9 \\ \xe3 \xef \xae \xd2 \xbc \xd7 #\x1d \xbb 6\xea \xd2 cH\xdf \xf5 \x9b \xb2 \xde \xb5 \xea \xd7 \x0c \x8c 7S\x07 \x86 :\x94 W\xca z\x1f <H!B\x08 s\x83 rqo<G\x8b \xa0 \xfc \xe5 \x12 mi\xf2 \x01 X<\x15 \x93 L\xa1 s\x00 \xa2 \xe5 \\ \xa9 \x1b \x9c \xdf \xd5 &_\xc0 \xf2 l1_Z*r\xf3 \x91 \xb3 \xc2 +\xa9 \xfc FrB@>e\x10 \x16 \xd2 \xbd C*\xcb -%Gl}\xe0 \x1f Nv\x7f mx%E~\x1d !h\xd8 \xb6 8\xc1 |\xb1 \xcf \x0e \x01 .^\x9f 7\xff F\x9f \xd8 Fz\x91 \x99 >\xfb \x8a N\xa2 \x9f \xd4 S\xdc |\xeb \x9c \x90 \xca \x1c \x1f \xac \xee `\r \xba Rm\x99 \xe3 =\x17 \xa9 wB3s_\xa1 \x13 \xe0 @,\x9a \x94 \x8a $\xfc \x83 \xfb \xae \xde \xa0 A\xbc \xf8 \r @\xe7 \xe0 JS\xde N\xe4 \xbc \xce \x91 \xa5 E\xf3 22y\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 \xe3 v$\xac 0\xe8 Y\x8a \x9e \xb3 \xed \xf8 \x99 \xb6 \xd5 \x93 \x00 \xd4 \x13 \x03 \xeb ~\xff \xcb 0\xd8 \x1a :\x19 [\x14 \xcd \x95 \xeb \xed \xce \x98 \x03 J\x16 \x99 \x0f \x1f \x99 h\xbf E\xcc \xcf \xd9 q\x1b \xab \x19 \xce ^\x1f \xa1 \xbc ,\\ \xe1 \xa0 D\xc9 \xa2 \xfd \x96 \xfa \xbc \x9e j\xb5 \xd0 5\x85 \x92 .\xe1 u\xa5 \xfc \x83 \xa7 \xe6 q\xfc 1\xce \xa9 \x7f @\xd3 eq\x9d \x93 \xa4 \x9a \xef \xe4 \xb7 \xf2 \xa4 \x8c \x13 \x8e \xdf \x8d T\x9b \xda \x02 $\x93 \' Ek\x16 \x9d \xcc +\xa8 \x98 E\xab \xad f\xa1 [\xbe \x91 \xfb \xec \xaa 7\xe1 \xa5 \xc1 \xb8 \xfd \xc9 \xbc \xbc \x1f r\xc9 q\xd7 \xcf \x99 \x0e \x93 \xc3 \xdf p\xed \xcf \xe3 \x0f \xbd \xf3 \x95 \x04 \xf3 \x7f \xbe ?\x0b ~\x15 \xe3 \xfe \xba \x80 k`\xd4 \xe5 \xd2 \x1e ["D6?<\xe7 \xb1 \x16 \x96 \x7f \x9c XFB\t K\xf7 \x96 \xad \xb0 cg\x1f \xc4 q\xeb })\x8d \x94 \xa3 \xd0 \x0f o\xcb \xc7 \x0e wq\xd0 \xe0 \x10 \xe3 \xd1 F\x80 B\xec \x8c \x87 AI/\xd0 \xf9 \xc9 \r \x81 \xf4 \xd4 \xb3 8\x01 \x06 Y\xb3 P\xa7 Q\xf0 \xe5 \x17 \x93 \xa3 \xc8 80\xc9 z1\xf9 R\x94 \x80 \x08 \x1c \xfd \xe3 \xc1 \xac _L\xd5 \xb4 s\x1b ^t8F\x13 A\xb0 A\xe8 6q{\xa1 \xc3 \x11 \xf8 \x9c \x16 \x87 \xe1 c\xba \x8a \xe0 \xa8 Nn\x9d \xa0 5\x1f KG/"\xa6 \xff \x9a \x85 )\x04 r\xb7 \xe4 n\x91 \xc4 \x00 I\xa7 l\x1f v\xb8 \x82 j\xa7 \xfb \x89 \xd3 y;\xe2 \xa2 (\xd6 \xb3 F\xaa D\x0b \xdf A\x05 \x07 \x02 Z\xe1 3\xc7 9\x8e \xb1 \xd4 :L@\xf0 \x0b T\x96 \xd0 \xdd h%b\x9a `l\xeb t\x96 ht"\x8e \xe0 _\xd1 \xc5 \x96 F\xe5 }Z\x02 \xdb y\xa2 O\x99 \xe8 m\x00 \x9a g\n 1\xae \xea 7\xc2 L>\xee \x9c \xf0 nt\xf3 $\xa5 s\xc2 Dn\xa8 \xd6 \x85 \xd8 \xbc `7\x92 $w\x9b \x85 r\x84 \xb0 \x0b \x01 \xd0 C\xdc \xe4 Y\x90 o\x01 8o+p:\x1e \x9b (\x15 \x8b \xa8 1\x8a l5\xb4 \x04 \x0e \xcb t1\xae \xfe \x1f \x1c DJ\xb7 \' \x14 \xa2 NJ\x0c n(e\xfc \xa5 \xb0 \x8d \xc8 PlI\x17 $\r \xaa ,$\x82 \x0b \x07 \x82 \xe7 \xf0 \x85 \x9b \x03 \x85 \xa4 \x9f Z\r \x1f \x18 \x95 \xb3 \xe5 /I\xbd oov\xc2 l\x8d \x0f \x05 If\x94 \xb7 y*\x1b \r \x80 \xfa \xce \xe0 8]\x98 \xb2 {\xcb \xa6 \xd7 \x83 I\x84 \x8f \r \x93 \\ 8$\x80 \xed \xf4 }\x1c \xd2 \x86 \xcd ~S\x10 b!\xd4 \x05 \xd5 \x80 I\x97 \x8c \xea C\xbf \x88 )\xcd \xb9 R\xc6 A\xc2 Y\x03 j\x89 \x87 \x16 \xfb \xfb \x1f \xdf \x02 \xe7 "0\xfc \xe7 \n \x01 x\xa2 5-\xc7 \xbd \xd0 6\x18 \x8c \xd2 \x9d y\x19 \xd9 \xee \x1a \x15 \xc0 F\x7f \xc2 /\xb0 \x01 \x0b \x14 ?U%\n \x0b dQj\xcf @\xbc \x17 \xd7 \x8e H#o\x1b >\xc9 \x97 \x9f 6\x0e \x9b ]\xde F\x00 :\xc4 0@@9\x9b \xac 8\x17 Z\xa2 9\xcd \xc2 g\xb4 (1\xe4 \xd2 Gk\x7f >IP\xdb (\x12 \x13 &W\xb2 \xef \x96 \x95 \x98 \x1d \xba \x9d \xc3 \x8f 5\x94 \xdc \xf2 \x02 \xf5 \xdb fc\x00 \xca \xca \x18 \xcb #\x05 @\x95 ~\x0b \xca \xea \xd6 \xe0 \x04 \x99 N\xd0 \xee ?Br\x81 \x96 c\x9d \xa0 Y=$\x10 \xb6 B\xf7 >\xa2 \xfc |\xa4 \xac \x0b \xdd \xd0 S\x1d \xc7 \x82 r~\x1b \xbc E\x89 \xf2 \xa2 \xfe Q\t }sr\xd8 V^$Ba\r \x1f \x9b \xf9 \xf5 \xa2 \x0b \x8d \xc0 p\xec \xe9 \x87 \xcf \xfc \xbc \xc9 )\x95 P\x88 bl7\xcf 4e\xfd #c/(4\xf2 F\xb1 \x81 M\xb3 CSPN\x12 *\x06 Ci<\xf9 \xcd \xdc \xb2 \xeb p7\xd3 \xae \xea \x99 \xf6 \xba (\xf8 \x90 \xb5 \x0c V\xfb g[VLee\xff \xdb \xa2 \x10 C\x10 \xe0 \x97 \x99 \xf1 \x0f \xa4 e\x14 \xb8 \xdb \xd8 \x81 \xb5 \xef \x1c \x80 \x98 \xe1 c\xd8 \xe3 \xd9 \x10 o\xe1 k!a\xf9 .H\xac o\x88 N\xb7 ([\\ 3A\xc2 `\xce ;\xa4 \x1c \x05 \t \x0e 6\xf0 \xc0 \xd2 m\xbc o\x0b \xc0 \xe3 }\xe2 \x97 \xd0 h\x0b \xc4 \xa0 8S\xe9 \xfb (\xe0 4nAV(\xbb !\xa0 Og\xc5 "\xa0 y\x00 \x82 \xa8 \x81 \x1e \xdd \x17 \xc3 \x9e \xa6 S6"0\x1d \x83 \xbb vF\x03 \xbb H\x96 \xde \x08 \xba \x9f \x93 \xe2 \xab br\xc9 \xcc \x0f !I\x9b \xff S\x17 ErH\xc7 \n \x02 :}\x92 \xf7 \xd5 \xbe \xb2 \x13 \xa2 \xa2 \xd5 \x06 J\xc4 \xdd A\x9b \xb4 \xbb \xe6 \xea \xa9 \xc8 IuS11e\x80 \xbd \x85 \xea P\x11 }\xc8 A0\x8e L\x86 \xf7 \xbd \xd3 \xb3 \xbd ^\xff \xc7 \xf6 \x85 2\xff \xa7 n\xb1 \x82 \xef \x17 \x99 \xf4 \xe6 \xa7 bj+\xf8 \x90 \xe8 \x95 q\xdc \n \x9e I\x04 \x8f \xb6 \xa9 \x0b \x84 \x98 n\xc2 fNYu4\x83 \t k\xf9 F\x83 \xda \xc8 @9\xee \x01 P\x16 \xb1 A\xa0 \x05 \x9d \xb5 \xd8 K\x95 M\xeb \xd6 \xeb \x0b \xb0 \x8b \x82 \xbd \xdf \xf4 \xb6 \t R\xfa D\xa5 \xdd W8E=e\x88 0^\x89 \xcf \xcf \xa1 \xbd 0;\x86 \x85 \xa5 \xc5 \xab |[ \x99 \xae W\xbd \x0c F\xa9 \x9e \xd2 \x85 MW\xa0 H\x0f \xda \xcc vZ\xb7 hh\x02 2>\x19 "\xb3 I\xa2 .K\x94 G\xbf \xe0 I\xe7 U\xf3 5\xe9 iP\x13 I\x8e \xad \xef \xa0 \x89 \x99 7{}6\xd7 \x07 ]\xab \x07 t\x9c \xcb \x8c \x87 \xe7 \x8f h J\xa0 L\xd0 \xde \xa9 MCyK0:\xbd H\x90 d$\xad \xd3 \xe0 \xc5 \x7f \xc7 \xb3 \x94 WV^\xfa \x82 E\xe1 \xa1 4\xec &U-~\x0e \x7f JJ\xde \x0c \xa4 \x19 \x9f \xbf ,\n \xa0 \xd7 7,\n !\xb6 \xb9 \x90 %\xa4 \xd6 \x93 AI\x0e w*\xd2 M-,\x00 \x1e \t bl\xb3 E\x08 \xd4 \x18 \x01 \xf5 W\xd6 \x15 \xb7 \x1d h\xb3 \xec \xa9 \xf3 \xd1 \xe6 g\xa9 \' \xef \xa3 %Y\x16 \x17 m\xf3 Q\x18 X\xca \x1d \xda 8\x96 \xad \xcf \xf4 \x82 JQ\x9e @\x84 4V\xb5 \xfd \xa5 \xe8 \xbd \x00 \xb4 \x98 \x10 !\xec \xd5 `\xb8 \xce x\xb8 \x9e \x96 \xb7 \xed \x06 ;\xcb \xce \x11 \xea \x8d |]to-{\x9d \x8f \x91 n\xd9 \x01 c%C\xc0 5\xac \xd9 $wf\xbe <\x10 \xd9 \x12 \xc0 \x8c \xd4 \xf3 \x17 \xf1 \xa9 d<\xa0 4\xa9 c\xf2 @\x1d \x9d u:\xf0 \xb4 \n \xd1 \x9c *\x9e a*\x82 \x1b \x1e *\xcc \xbb \xf8 \x81 \xbf >\xcb \xab A\x8c \xd8 \xd8 \x95 F\xac WG]\xd1 \xf8 \xd5 \xa4 {\xdd \xca @\xde \xac \xff 0\xf5 \xea \xfa E|\xc9 \xe8 \xf6 C\x96 =\x16 \xb4 \xec \xf5 \xf2 \xc1 \xc6 \xa6 \xd8 \xc5 \xb8 #!\x10 [<\xfd B~\x02 \x93 T\xc6 \xba \x8a T\xf7 \x10 \xb9 \xe5 \xd6 \xe0 \xb8 -`f1@\xec B\x87 \xd6 &p\x92 2\xc0 \x13 \xa9 u:<\r r\xbd \xd9 ;K\xc6 n\xaa Rt#\xfe \x08 bb\x16 \x9c ^w\x13 \xdb pg\' \x16 \x7f >~\xa2 n(\xef I\x1b \xfe gK\x86 \x8e n\x1a \x08 \xab \xdf k2\x00 s\x12 \x1b \xb5 \xf6 /h\x8f \xeb \x90 \xc8 0 \x01 \xf6 \xc5 Z\xbd \x0f \x9f \xef \xb8 \xed t\x06 i,q\\ G5\xb0 \xe1 \x8f \x17 \xa1 \xa9 \xb2 1ql\xc5 \x84 \xd0 \xdc \xc7 \xcb |\xe2 \x15 \xc4 \xc4 ;\xd5 \xc8 \xb9 \x95 \xae -e\xc8 \x19 \xcf y\xeb \x9c ^\x0e i\xd8 \xea \x9d \x9e g\xb9 \xb5 g{|?\x16 \xbf \xda U\x12 \xf4 aI\x0b \xb3 Q({\xdd \xeb \x90 u\xae 28\xcd \xa9 \x11 -\x12 p\x8d W_\x94 Z\x96 \xbe \x9c \\ \xec B\xe7 ~w\x18 ?\xfd \xe5 q\xa7 \xa7 \xdc \xc1 <S F\x05 \xcd \x0c \x81 K \x83 R\xea \xb7 \xfb \x1b \xe5 /T\xe3 Ix\x06 ^`\xb2 2\xeb \x8b *\xad \xd3 \x84 L\xc8 f\xee 9\xa8 E\xfc jDV^p\xb7 \xcb }k\xca $6\x03 \xa5 _\xb9 e\xe9 l`\xaa \x86 \xe8 n\n \xb7 \xac \xb6 \xbf \xfe l\x96 \x84 \xa3 \xcb >\xb4 =\x96 \x86 ]\x8d \x15 \xf7 ~e\x15 9\xa9 \xb6 \x94 Fs\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 )
43
66
44
- ptx_kernel_bytes = ptx_kernel .encode ('utf-8' )
45
- minimal_ptx_kernel_bytes = minimal_ptx_kernel .encode ('utf-8' )
46
67
47
68
def test_unrecognized_option_error ():
48
69
with pytest .raises (nvjitlink .nvJitLinkError , match = "ERROR_UNRECOGNIZED_OPTION" ):
@@ -55,75 +76,82 @@ def test_invalid_arch_error():
55
76
56
77
57
78
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 )
61
83
62
84
63
85
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 )
67
90
68
91
69
92
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 )
75
98
76
99
77
100
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 )
84
108
85
109
86
110
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 )
94
119
95
120
96
121
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 )
105
131
106
132
107
133
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 )
116
143
117
144
118
145
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 )
127
155
128
156
129
157
def test_package_version ():
0 commit comments