已合并
修复mie_ops_lib.so安全编译选项问题 #925
hw-zhoutianyang创建于 4月20日
修复mie_ops_lib.so安全编译选项问题 #925
已合并
共 1 个文件变更+17-18
| @@ -17,22 +17,21 @@ | |||
| 17 | import os | 17 | import os |
| 18 | import sys | 18 | import sys |
| 19 | import glob | 19 | import glob |
| 20 | -import torch | 20 | +from setuptools import setup |
| 21 | -from setuptools import setup, find_packages | ||
| 22 | from torch.utils.cpp_extension import BuildExtension | 21 | from torch.utils.cpp_extension import BuildExtension |
| 23 | 22 | ||
| 24 | import torch_npu | 23 | import torch_npu |
| 25 | from torch_npu.utils.cpp_extension import NpuExtension | 24 | from torch_npu.utils.cpp_extension import NpuExtension |
| 26 | 25 | ||
| 27 | PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.abspath(torch_npu.__file__)) | 26 | PYTORCH_NPU_INSTALL_PATH = os.path.dirname(os.path.abspath(torch_npu.__file__)) |
| 28 | -USE_NINJA = os.getenv('USE_NINJA') == '1' | 27 | +USE_NINJA = os.getenv("USE_NINJA") == "1" |
| 29 | BASE_DIR = os.path.dirname(os.path.realpath(__file__)) | 28 | BASE_DIR = os.path.dirname(os.path.realpath(__file__)) |
| 30 | 29 | ||
| 31 | if "--ops" in sys.argv: | 30 | if "--ops" in sys.argv: |
| 32 | idx = sys.argv.index("--ops") | 31 | idx = sys.argv.index("--ops") |
| 33 | if idx + 1 < len(sys.argv): | 32 | if idx + 1 < len(sys.argv): |
| 34 | ops = sys.argv[idx + 1] | 33 | ops = sys.argv[idx + 1] |
| 35 | - del sys.argv[idx:idx + 2] # 从 argv 中移除 --ops 及其值，避免 setuptools 报错 | 34 | + del sys.argv[idx : idx + 2] # 从 argv 中移除 --ops 及其值，避免 setuptools 报错 |
| 36 | else: | 35 | else: |
| 37 | raise ValueError("Missing the value of input parameter '--ops', mie_ops setup failed!") | 36 | raise ValueError("Missing the value of input parameter '--ops', mie_ops setup failed!") |
| 38 | else: | 37 | else: |
| @@ -46,33 +45,33 @@ ext = NpuExtension( | |||
| 46 | name=mie_ops_version + ".mie_ops_lib", | 45 | name=mie_ops_version + ".mie_ops_lib", |
| 47 | sources=source_files, | 46 | sources=source_files, |
| 48 | extra_compile_args=[ | 47 | extra_compile_args=[ |
| 49 | - '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include"), | 48 | + "-I" + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include"), |
| 50 | - '-I' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/acl/inc"), | 49 | + "-I" + os.path.join(PYTORCH_NPU_INSTALL_PATH, "include/third_party/acl/inc"), |
| 51 | - '-Wno-unused-variable', | 50 | + "-Wno-unused-variable", |
| 52 | - '-Wno-unused-parameter', | 51 | + "-Wno-unused-parameter", |
| 53 | - '-Wno-unused-function', | 52 | + "-Wno-unused-function", |
| 54 | - '-Wno-narrowing', | 53 | + "-Wno-narrowing", |
| 55 | ], | 54 | ], |
| 56 | extra_link_args=[ | 55 | extra_link_args=[ |
| 57 | - '-L' + os.path.join(PYTORCH_NPU_INSTALL_PATH, "lib"), # torch_npu 的 so 所在目录 | 56 | + "-L" + os.path.join(PYTORCH_NPU_INSTALL_PATH, "lib"), # torch_npu 的 so 所在目录 |
| 58 | - '-ltorch_npu', # 链接 libtorch_npu.so | 57 | + "-ltorch_npu", # 链接 libtorch_npu.so |
| 59 | - '-ltorch', # 可能还需要链接 libtorch.so | 58 | + "-ltorch", # 可能还需要链接 libtorch.so |
| 60 | - '-lc10', # PyTorch 依赖的库 | 59 | + "-lc10", # PyTorch 依赖的库 |
| 61 | - '-Wl,-rpath,/path/to/torch_npu/lib', # 运行时 rpath | 60 | + "-Wl,-z,now", |
| 62 | ], | 61 | ], |
| 63 | ) | 62 | ) |
| 64 | exts.append(ext) | 63 | exts.append(ext) |
| 65 | 64 | ||
| 66 | setup( | 65 | setup( |
| 67 | name=mie_ops_version, | 66 | name=mie_ops_version, |
| 68 | - version='1.0', | 67 | + version="1.0", |
| 69 | keywords=mie_ops_version, | 68 | keywords=mie_ops_version, |
| 70 | packages=[mie_ops_version], | 69 | packages=[mie_ops_version], |
| 71 | package_dir={mie_ops_version: "mie_ops"}, # 将实际版本的mie_ops映射到mie_ops目录 | 70 | package_dir={mie_ops_version: "mie_ops"}, # 将实际版本的mie_ops映射到mie_ops目录 |
| 72 | ext_modules=exts, | 71 | ext_modules=exts, |
| 73 | package_data={ | 72 | package_data={ |
| 74 | - mie_ops_version: ['*.py', '*.so', 'opp/**/*'], | 73 | + mie_ops_version: ["*.py", "*.so", "opp/**/*"], |
| 75 | - mie_ops_version + '.torch_ops_extension': ['*.py', '*.so'], | 74 | + mie_ops_version + ".torch_ops_extension": ["*.py", "*.so"], |
| 76 | }, | 75 | }, |
| 77 | cmdclass={"build_ext": BuildExtension.with_options(use_ninja=USE_NINJA)}, | 76 | cmdclass={"build_ext": BuildExtension.with_options(use_ninja=USE_NINJA)}, |
| 78 | ) | 77 | ) |