[BACKEND] Fix issue of editable mode for sunrise backend

python/setup_tools/setup_helper.py

@@ -293,7 +293,8 @@ def handle_flagtree_backend():
         print(f"\033[1;32m[INFO] FlagtreeBackend is {configs.flagtree_backend}\033[0m")
         configs.extend_backends.append(configs.flagtree_backend)
         if "editable_wheel" in sys.argv and configs.flagtree_backend not in configs.plugin_backends:
-            ext_sourcedir = os.path.abspath(f"./third_party/{configs.flagtree_backend}/python/{ext_sourcedir}") + "/"
+            configs.ext_sourcedir = os.path.abspath(
+                f"./third_party/{configs.flagtree_backend}/python/{configs.ext_sourcedir}") + "/"
 
 
 def handle_plugin_backend(editable):
@@ -310,6 +311,8 @@ def handle_plugin_backend(editable):
                 os.makedirs(dst_build_plugin_dir)
             dst_build_plugin_path = dst_build_plugin_dir / flagtree_plugin_so
             shutil.copy(src_build_plugin_path, dst_build_plugin_path)
+            dst_build_plugin_path = Path(__file__).resolve().parent.parent.parent / "python" / "triton" / "_C"
+            shutil.copy(src_build_plugin_path, dst_build_plugin_path)
         src_install_plugin_path = flagtree_backend_dir / flagtree_plugin_so
         if flagtree_backend in ("mthreads", "sunrise"):
             dst_install_plugin_dir = Path(
@@ -362,8 +365,6 @@ def uninstall_triton():
 download_flagtree_third_party("flir", condition=(flagtree_backend == "aipu"), hock=utils.aipu.precompile_hock,
                               required=True)
 
-handle_plugin_backend(False)
-
 handle_flagtree_backend()
 
 cache = FlagTreeCache()

setup.py

@@ -453,7 +453,8 @@ def build_extension(self, ext):
         thirdparty_cmake_args = get_thirdparty_packages([get_llvm_package_info()])
         thirdparty_cmake_args += self.get_pybind11_cmake_args()
         extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.path)))
-        wheeldir = os.path.dirname(extdir)
+        print(f"extdir {extdir}")
+        wheeldir = os.path.dirname(os.path.dirname(extdir))
 
         # create build directories
         if not os.path.exists(self.build_temp):
@@ -713,6 +714,11 @@ def add_link_to_backends(external_only):
 package_data_tools = ["compile.h", "compile.c"]
 if helper.flagtree_backend == "xpu":
     package_data_tools += ["compile_xpu.h", "compile_xpu.c"]
+
+if helper.flagtree_backend == "sunrise":
+    package_data = {"": ["*TritonPlugin.so"]}
+else:
+    package_data = {}
 #  package_data = {
 #       "triton/tools/extra": sum((b.tools_package_data for b in backends), []),
 #  **{f"triton/backends/{b.name}": b.package_data
@@ -855,8 +861,9 @@ def get_git_version_suffix():
     packages=list(get_packages()),
     package_dir=dict(get_package_dirs()),
     entry_points=get_entry_points(),
+    package_data=package_data,
     include_package_data=True,
-    ext_modules=[CMakeExtension("triton", "triton/_C/")],
+    ext_modules=[CMakeExtension("triton", helper.configs.ext_sourcedir)],
     cmdclass={
         "bdist_wheel": plugin_bdist_wheel,
         "build_ext": CMakeBuild,

third_party/sunrise/CMakeLists.txt

@@ -5,7 +5,7 @@ option(EDITABLE_MODE "Build in developer (editable) mode" OFF)
 if(FLAGTREE_PLUGIN)
   set(SUNRISE_PLUGIN_DIR "${Python3_SITELIB}/triton/_C")
 elseif(EDITABLE_MODE)
-  set(SUNRISE_PLUGIN_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
+  set(SUNRISE_PLUGIN_DIR "${CMAKE_CURRENT_SOURCE_DIR}/python/triton/_C")
 else()
   set(SUNRISE_PLUGIN_DIR "${Python3_SITELIB}/triton/_C")
 endif()

third_party/sunrise/python/triton/backends

@@ -0,0 +1 @@
+../../../../python/triton/backends