Skip to content

Commit

Permalink
Set the url to pytorch/ao and clean up some duplication (#1446)
Browse files Browse the repository at this point in the history
  • Loading branch information
janeyx99 authored Dec 19, 2024
1 parent 1beb6d0 commit aea2356
Showing 1 changed file with 18 additions and 24 deletions.
42 changes: 18 additions & 24 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -64,36 +64,29 @@ def get_extensions():
use_cuda = torch.cuda.is_available() and CUDA_HOME is not None
extension = CUDAExtension if use_cuda else CppExtension

extra_link_args = []
extra_compile_args = {
"nvcc": [
"-O3" if not debug_mode else "-O0",
"-t=0",
]
}

if not IS_WINDOWS:
extra_link_args = []
extra_compile_args = {
"cxx": [
"-O3" if not debug_mode else "-O0",
"-fdiagnostics-color=always",
],
"nvcc": [
"-O3" if not debug_mode else "-O0",
"-t=0",
]
}
extra_compile_args["cxx"] = [
"-O3" if not debug_mode else "-O0",
"-fdiagnostics-color=always",
]

if debug_mode:
extra_compile_args["cxx"].append("-g")
extra_compile_args["nvcc"].append("-g")
extra_link_args.extend(["-O0", "-g"])

else:
extra_link_args = []
extra_compile_args = {
"cxx": [
"/O2" if not debug_mode else "/Od",
"/permissive-"
],
"nvcc": [
"-O3" if not debug_mode else "-O0",
"-t=0",
]
}
extra_compile_args["cxx"] = [
"/O2" if not debug_mode else "/Od",
"/permissive-"
]

if debug_mode:
extra_compile_args["cxx"].append("/ZI")
Expand Down Expand Up @@ -125,6 +118,7 @@ def get_extensions():

return ext_modules


setup(
name="torchao",
version=version+version_suffix,
Expand All @@ -138,7 +132,7 @@ def get_extensions():
description="Package for applying ao techniques to GPU models",
long_description=open("README.md").read(),
long_description_content_type="text/markdown",
url="https://github.com/pytorch-labs/ao",
url="https://github.com/pytorch/ao",
cmdclass={"build_ext": BuildExtension},
options={"bdist_wheel": {
"py_limited_api": "cp39"
Expand Down

0 comments on commit aea2356

Please sign in to comment.