88from typing import Any , Final
99
1010import torch
11+ import tempfile
12+ import signal
13+ from triton import knobs
14+ from triton .runtime .errors import PTXASError
15+ from functools import partial
1116
1217from triton ._C .libtriton import llvm # pyright: ignore[reportMissingImports]
1318from triton ._C .libtriton .tle .llvm import parse_llvm_ir # pyright: ignore[reportMissingImports]
1722CLANG = os .getenv ("CLANG" , "clang" )
1823CLANG_FLAGS = shlex .split (os .getenv ("CLANG_FLAGS" , "" ))
1924
25+ NVCC = os .getenv ("NVCC" , "nvcc" )
26+ NVCC_FLAGS = shlex .split (os .getenv ("NVCC_FLAGS" , "" ))
27+
28+ PTXAS = os .getenv ("PTXAS" , "ptxas" )
29+
30+ NVLINK = os .getenv ("NVLINK" , "nvlink" )
31+ NVLINK_FLAGS = shlex .split (os .getenv ("NVLINK_FLAGS" , "" ))
32+
33+ OPT = os .getenv ("OPT" , "opt" )
34+
2035
2136def _sanitize_clang_ir (ir : str ) -> str :
2237 # Newer clang emits attributes that this Triton branch's LLVM parser does
@@ -43,23 +58,113 @@ def _get_cuda_gpu_arch() -> str:
4358 if arch :
4459 return f"--cuda-gpu-arch={ arch }
4560 major , minor = torch .cuda .get_device_capability ()
46- return f"--cuda-gpu-arch=sm_{ major } { minor }
61+ suffix = "a" if major >= 9 else ""
62+ return f"--cuda-gpu-arch=sm_{ major } { minor } { suffix }
63+
64+
65+ def make_cubin_inspection_hook (cuda_self , triton_self , stages , options , language , capability ):
66+
67+ def make_cubin (self , src , metadata , opt , capability ):
68+ fsrc_cuda = cuda_self .source_file
69+ arch = _get_cuda_gpu_arch ().split ('=' )[1 ]
70+ fbin_cuda = tempfile .NamedTemporaryFile (delete = False , suffix = '.o' ).name
71+
72+ build = subprocess .run ([NVCC , "-c" , "-rdc=true" , f"-arch={ arch } , * NVCC_FLAGS , "-o" , fbin_cuda , fsrc_cuda ],
73+ capture_output = True )
74+ assert build .returncode == 0 , (f"nvcc failed\n stderr:\n { build .stderr .decode ()} )
75+
76+ with tempfile .NamedTemporaryFile (delete = False , mode = 'w' , suffix = '.ptx' ) as fsrc_triton , \
77+ tempfile .NamedTemporaryFile (delete = False , mode = 'r' , suffix = '.log' ) as flog :
78+ fsrc_triton .write (src )
79+ fsrc_triton .flush ()
80+ fbin_triton = fsrc_triton .name + '.o'
81+ fbin_combined = fbin_triton + '.combined.cubin'
82+
83+ compile_only_cmds = ['-c' ]
84+ line_info = ["-lineinfo" , "-suppress-debug-info" ] if knobs .compilation .disable_line_info else ["-lineinfo" ]
85+ fmad = [] if opt .enable_fp_fusion else ['--fmad=false' ]
86+ disable_opt = ['--opt-level' , '0' ] if knobs .nvidia .disable_ptxas_opt else []
87+ ptx_extra_options = opt .ptx_options .split (" " ) if opt .ptx_options else []
88+
89+ ptxas_cmd = [
90+ PTXAS , * compile_only_cmds , * line_info , * fmad , '-v' , * disable_opt , * ptx_extra_options ,
91+ f'--gpu-name={ arch } , fsrc_triton .name , '-o' , fbin_triton
92+ ]
93+
94+ try :
95+ subprocess .run (ptxas_cmd , check = True , close_fds = False , stderr = flog )
96+ if os .path .exists (fsrc_triton .name ):
97+ os .remove (fsrc_triton .name )
98+ if os .path .exists (flog .name ):
99+ os .remove (flog .name )
100+ except subprocess .CalledProcessError as e :
101+ with open (flog .name ) as log_file :
102+ log = log_file .read ()
103+ if os .path .exists (flog .name ):
104+ os .remove (flog .name )
105+
106+ if e .returncode == 255 :
107+ error = 'Internal Triton PTX codegen error'
108+ elif e .returncode == 128 + signal .SIGSEGV :
109+ error = '`ptxas` raised SIGSEGV'
110+ else :
111+ error = f'`ptxas` failed with error code { e .returncode }
112+ raise PTXASError (f"{ error } \n "
113+ f"`ptxas` stderr:\n { log } \n "
114+ f'Repro command: { " " .join (ptxas_cmd )} \n ' )
115+
116+ nvlink_cmds = [
117+ NVLINK ,
118+ f"-arch={ arch } ,
119+ * NVLINK_FLAGS ,
120+ fbin_triton ,
121+ fbin_cuda ,
122+ "-o" ,
123+ fbin_combined ,
124+ ]
125+
126+ try :
127+ subprocess .run (nvlink_cmds , check = True , close_fds = False , stderr = flog )
128+ except Exception as e :
129+ import logging
130+ logging .error (f"error runing nvlink: { shlex .join (nvlink_cmds )} )
131+ logging .exception (e )
132+
133+ with open (fbin_combined , 'rb' ) as f :
134+ cubin = f .read ()
135+ if os .path .exists (fbin_combined ):
136+ os .remove (fbin_combined )
137+ if os .path .exists (fbin_triton ):
138+ os .remove (fbin_triton )
139+ if os .path .exists (fbin_cuda ):
140+ os .remove (fbin_cuda )
141+ return cubin
142+
143+ stages ["cubin" ] = lambda src , metadata : make_cubin (triton_self , src , metadata , options , triton_self .target .arch )
47144
48145
49146class CUDAJITFunction (object ):
50147
51148 def __init__ (self , fn : Any , file : Path , * args , ** kwargs ) -> None :
52- super ().__init__ (* args , ** {k : v for k , v in kwargs .items () if k not in "extern_func_name" , "deferred" )})
149+ super ().__init__ ()
150+ # super().__init__(*args, **{k: v for k, v in kwargs.items() if k not in ("extern_func_name", "deferred")})
53151 self .fn : Final [Any ] = fn
54152 self .code : Final [str ] = file .read_text ()
55153 self .region_dialect : Final [str ] = "cuda"
56154 self .lowered_region_dialect : Final [str ] = "llvm"
57155 self .arg_dialect : Final [str ] = "llvm"
58156 self .source_file : Final [str ] = str (file )
157+ self .compiler = kwargs .get ("compiler" , None )
158+ self .target = kwargs .get ("target" , None )
159+ self .extern_file = kwargs .get ("extern_file" , None )
59160 self .extern_func_name = kwargs .get ("extern_func_name" , None )
60161 self .deferred : Final [bool ] = kwargs .get ("deferred" , False )
61162 self .__triton_builtin__ : Final [bool ] = True
62163
164+ if self .compiler .lower () == "nvcc" and knobs .runtime .add_stages_inspection_hook is None :
165+ nvcc_cuda_hook = partial (make_cubin_inspection_hook , self )
166+ knobs .runtime .add_stages_inspection_hook = nvcc_cuda_hook
167+
63168 def register_pending_source (self , * , hint : str = "" ) -> str :
64169 if not self .extern_func_name :
65170 raise RuntimeError ("deferred tle_raw CUDA source requires extern_func_name= "
@@ -116,6 +221,38 @@ def make_llvm(self, mlir_context) -> str:
116221 module = parse_llvm_ir (_sanitize_clang_ir (build .stdout .decode ()), llvm_context , mlir_context )
117222 return f"{ module }
118223
224+ def make_bc (self , public_api_names = None ):
225+ fbc_cuda_unopti = Path (self .source_file ).with_suffix ('.bc.unopti' )
226+ fbc_cuda = Path (self .source_file ).with_suffix ('.bc' )
227+
228+ build = subprocess .run ([
229+ CLANG , "-c" , "-x" , "cuda" , "--cuda-device-only" ,
230+ _get_cuda_gpu_arch (), "-emit-llvm" ,
231+ # "-O1",
232+ "-fcuda-flush-denormals-to-zero" , * CLANG_FLAGS , "-o" , fbc_cuda_unopti , self .source_file
233+ ], capture_output = True )
234+ assert build .returncode == 0 , (f"clang failed\n stderr:\n { build .stderr .decode ()} )
235+
236+ if public_api_names is None :
237+ public_api_names = [self .extern_func_name ]
238+ elif isinstance (public_api_names , str ):
239+ public_api_names = [public_api_names ]
240+ else :
241+ public_api_names = list (public_api_names )
242+ if not public_api_names or any (not name for name in public_api_names ):
243+ raise ValueError ("make_bc requires at least one public API name" )
244+ public_api_list = "," .join (dict .fromkeys (public_api_names ))
245+
246+ opt = subprocess .run ([
247+ OPT , "--passes=internalize,inline,globaldce" , f"-internalize-public-api-list={ public_api_list } , "-o" ,
248+ fbc_cuda , fbc_cuda_unopti
249+ ], capture_output = True )
250+ assert opt .returncode == 0 , (f"opt failed\n stderr:\n { opt .stderr .decode ()} )
251+
252+ if os .path .exists (fbc_cuda_unopti ):
253+ os .remove (fbc_cuda_unopti )
254+ return fbc_cuda
255+
119256
120257def compile_deferred_pending_source (entry : dict , * , context ) -> str :
121258 source_text = entry ["source" ]
0 commit comments