[FRONTEND] clean build directly when executing python setup.py clean (#2238)

Current setup.py could not clean the build directly because the default
build directly has been changed in `CMakeBuild`. This PR is to clean
build directly in this regard.

python/setup.py

@@ -8,6 +8,7 @@ import sysconfig
 import tarfile
 import tempfile
 import urllib.request
+from distutils.command.clean import clean
 from pathlib import Path
 from typing import NamedTuple
 
@@ -158,6 +159,25 @@ def download_and_copy_ptxas():
 
 # ---- cmake extension ----
 
+def get_base_dir():
+    return os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir))
+
+
+def get_cmake_dir():
+    plat_name = sysconfig.get_platform()
+    python_version = sysconfig.get_python_version()
+    dir_name = f"cmake.{plat_name}-{sys.implementation.name}-{python_version}"
+    cmake_dir = Path(get_base_dir()) / "python" / "build" / dir_name
+    cmake_dir.mkdir(parents=True, exist_ok=True)
+    return cmake_dir
+
+
+class CMakeClean(clean):
+    def initialize_options(self):
+        clean.initialize_options(self)
+        self.build_temp = get_cmake_dir()
+
+
 class CMakeBuildPy(build_py):
     def run(self) -> None:
         self.run_command('build_ext')
@@ -178,10 +198,7 @@ class CMakeBuild(build_ext):
 
     def initialize_options(self):
         build_ext.initialize_options(self)
-        self.base_dir = os.path.abspath(
-            os.path.join(
-                os.path.dirname(__file__),
-                os.pardir))
+        self.base_dir = get_base_dir()
 
     def finalize_options(self):
         build_ext.finalize_options(self)
@@ -200,14 +217,6 @@ class CMakeBuild(build_ext):
         for ext in self.extensions:
             self.build_extension(ext)
 
-    def get_cmake_dir(self):
-        plat_name = sysconfig.get_platform()
-        python_version = sysconfig.get_python_version()
-        dir_name = f"cmake.{plat_name}-{sys.implementation.name}-{python_version}"
-        cmake_dir = Path(self.base_dir) / "python" / "build" / dir_name
-        cmake_dir.mkdir(parents=True, exist_ok=True)
-        return cmake_dir
-
     def build_extension(self, ext):
         lit_dir = shutil.which('lit')
         user_home = os.getenv("HOME") or os.getenv("USERPROFILE") or \
@@ -265,7 +274,7 @@ class CMakeBuild(build_ext):
                            "-DCMAKE_SHARED_LINKER_FLAGS=-fuse-ld=lld"]
 
         env = os.environ.copy()
-        cmake_dir = self.get_cmake_dir()
+        cmake_dir = get_cmake_dir()
         subprocess.check_call(["cmake", self.base_dir] + cmake_args, cwd=cmake_dir, env=env)
         subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=cmake_dir)
 
@@ -300,7 +309,7 @@ setup(
     ],
     include_package_data=True,
     ext_modules=[CMakeExtension("triton", "triton/_C/")],
-    cmdclass={"build_ext": CMakeBuild, "build_py": CMakeBuildPy},
+    cmdclass={"build_ext": CMakeBuild, "build_py": CMakeBuildPy, "clean": CMakeClean},
     zip_safe=False,
     # for PyPI
     keywords=["Compiler", "Deep Learning"],