@@ -109,11 +109,11 @@ def get_torch_root_path():
109109 except ImportError :
110110 return None
111111
112- def get_torch_mlu_root_path ():
112+ def get_ixformer_root_path ():
113113 try :
114- import torch_mlu
114+ import ixformer
115115 import os
116- return os .path .dirname (os .path .abspath (torch_mlu .__file__ ))
116+ return os .path .dirname (os .path .abspath (ixformer .__file__ ))
117117 except ImportError :
118118 return None
119119
@@ -229,7 +229,15 @@ def set_cuda_envs():
229229 os .environ ["CUDA_TOOLKIT_ROOT_DIR" ] = "/usr/local/cuda"
230230 os .environ ["NCCL_ROOT" ] = get_nccl_root_path ()
231231 os .environ ["NCCL_VERSION" ] = "2"
232-
232+
233+ def set_ilu_envs ():
234+ os .environ ["PYTHON_INCLUDE_PATH" ] = get_python_include_path ()
235+ os .environ ["PYTHON_LIB_PATH" ] = get_torch_root_path ()
236+ os .environ ["LIBTORCH_ROOT" ] = get_torch_root_path ()
237+ os .environ ["PYTORCH_INSTALL_PATH" ] = get_torch_root_path ()
238+ os .environ ["CUDA_TOOLKIT_ROOT_DIR" ] = "/usr/local/corex"
239+ os .environ ["IXFORMER_INSTALL_PATH" ] = get_ixformer_root_path ()
240+
233241class CMakeExtension (Extension ):
234242 def __init__ (self , name : str , path : str , sourcedir : str = "" ) -> None :
235243 super ().__init__ (name , sources = [])
@@ -280,8 +288,7 @@ def run(self):
280288 for ext in self .extensions :
281289 self .build_extension (ext )
282290 except Exception as e :
283- print ("ERROR: Build failed." )
284- print (f"Details: { e } " )
291+ print ("Build failed." )
285292 exit (1 )
286293
287294 def build_extension (self , ext : CMakeExtension ):
@@ -313,7 +320,7 @@ def build_extension(self, ext: CMakeExtension):
313320 f"-DDEVICE_ARCH={ self .arch .upper ()} " ,
314321 f"-DINSTALL_XLLM_KERNELS={ 'ON' if self .install_xllm_kernels else 'OFF' } " ,
315322 ]
316-
323+
317324 if self .device == "a2" or self .device == "a3" :
318325 cmake_args += ["-DUSE_NPU=ON" ]
319326 # set npu environment variables
@@ -328,6 +335,9 @@ def build_extension(self, ext: CMakeExtension):
328335 f"-DCMAKE_CUDA_ARCHITECTURES={ cuda_architectures } " ]
329336 # set cuda environment variables
330337 set_cuda_envs ()
338+ elif self .device == "ilu" :
339+ cmake_args += ["-DUSE_ILU=ON" ]
340+ set_ilu_envs ()
331341 else :
332342 raise ValueError ("Please set --device to a2 or a3 or mlu or cuda." )
333343
@@ -351,6 +361,7 @@ def build_extension(self, ext: CMakeExtension):
351361
352362 build_args = ["--config" , build_type ]
353363 max_jobs = os .getenv ("MAX_JOBS" , str (os .cpu_count ()))
364+ # max_jobs="2"
354365 build_args += ["-j" + max_jobs ]
355366
356367 env = os .environ .copy ()
@@ -572,7 +583,7 @@ def pre_build():
572583 idx = sys .argv .index ('--device' )
573584 if idx + 1 < len (sys .argv ):
574585 device = sys .argv [idx + 1 ].lower ()
575- if device not in ('a2' , 'a3' , 'mlu' , 'cuda' ):
586+ if device not in ('a2' , 'a3' , 'mlu' , 'cuda' , 'ilu' ):
576587 print ("Error: --device must be a2 or a3 or mlu (case-insensitive)" )
577588 sys .exit (1 )
578589 # Remove the arguments so setup() doesn't see them
0 commit comments