lzhengning
/
JittorMirror
mirror of https://github.com/Jittor/Jittor
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

polish files and add cuda arch hint

This commit is contained in:
Dun Liang 2021-10-01 19:41:04 +08:00
parent a78c3b4f12
commit 0fa0584fd3
2 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/jittor/__init__.py
View File

@ -9,7 +9,7 @@
# file 'LICENSE.txt', which is part of this source code package. # file 'LICENSE.txt', which is part of this source code package.
# *************************************************************** # ***************************************************************
__version__ = '1.3.0.1' __version__ = '1.3.0.2'
from jittor_utils import lock from jittor_utils import lock
with lock.lock_scope(): with lock.lock_scope():
ori_int = int ori_int = int

8
python/jittor/compiler.py
View File

@ -1259,7 +1259,7 @@ if use_data_gz:
.replace("-shared", "") .replace("-shared", "")
vdp = os.path.join(jittor_path, "src", "utils", "vdp") vdp = os.path.join(jittor_path, "src", "utils", "vdp")
run_cmd(fix_cl_flags(f"{cc_path} {dflags} -include {vdp} {data_s_path} -c -o {data_o_path}")) run_cmd(fix_cl_flags(f"{cc_path} {dflags} -include {vdp} {data_s_path} -c -o {data_o_path}"))
# os.remove(data_s_path) os.remove(data_s_path)
with open(data_gz_md5_path, 'w') as f: with open(data_gz_md5_path, 'w') as f:
f.write(md5) f.write(md5)
files.append(data_o_path) files.append(data_o_path)
@ -1280,6 +1280,12 @@ flags = core.flags()
if has_cuda: if has_cuda:
nvcc_flags = convert_nvcc_flags(cc_flags) nvcc_flags = convert_nvcc_flags(cc_flags)
if len(flags.cuda_archs): if len(flags.cuda_archs):
archs = []
for arch in flags.cuda_archs:
if arch<50:
LOG.w(f"CUDA arch({arch})<50 is not supported")
continue
archs.append(arch)
nvcc_flags += f" -arch=compute_{min(flags.cuda_archs)} " nvcc_flags += f" -arch=compute_{min(flags.cuda_archs)} "
nvcc_flags += ''.join(map(lambda x:f' -code=sm_{x} ', flags.cuda_archs)) nvcc_flags += ''.join(map(lambda x:f' -code=sm_{x} ', flags.cuda_archs))