change akg package path to mindspore

akg

@@ -1 +1 @@
-Subproject commit 94cb709ecaf5d1d869883dfe80cee7497dd0692c
+Subproject commit e2a0a264a0be549b51a035e5f783927052f8ead8

cmake/package.cmake

@@ -297,10 +297,14 @@ install(
 
 if((ENABLE_D OR ENABLE_GPU) AND ENABLE_AKG)
     set (AKG_PATH ${CMAKE_SOURCE_DIR}/build/mindspore/akg)
+    file(REMOVE_RECURSE ${AKG_PATH}/_akg)
+    file(MAKE_DIRECTORY ${AKG_PATH}/_akg)
+    file(TOUCH ${AKG_PATH}/_akg/__init__.py)
+    install(DIRECTORY "${AKG_PATH}/akg" DESTINATION "${AKG_PATH}/_akg")
     install(
         DIRECTORY
-            ${AKG_PATH}/akg
-        DESTINATION ${INSTALL_PY_DIR}/..
+            ${AKG_PATH}/_akg
+        DESTINATION ${INSTALL_PY_DIR}/
         COMPONENT mindspore
     )
 endif()

mindspore/_extends/parallel_compile/akg_compiler/akg_process.py

@@ -18,6 +18,22 @@ import shutil
 import subprocess
 import sys
 from multiprocessing import Pool, cpu_count
+import importlib
+
+def get_akg_path():
+    """get akg directory base path"""
+    search_res = importlib.util.find_spec("mindspore")
+    if search_res is None:
+        raise RuntimeError("Cannot find mindspore module!")
+
+    res_path = search_res.origin
+    find_pos = res_path.find("__init__.py")
+    if find_pos == -1:
+        raise RuntimeError("Find module mindspore origin file failed!")
+    akg_path = "{}_akg".format(res_path[:find_pos])
+    if not os.path.isdir(akg_path):
+        raise RuntimeError("Cannot find akg from mindspore module!")
+    return akg_path
 
 def copy_json(pid_path, ppid_path):
     """
@@ -36,7 +52,7 @@ def _compile_akg_task_gpu(*json_strs):
     Parameters:
         json_strs: list. List contains multiple kernel infos, suitable for json compile api.
     """
-
+    sys.path.insert(0, get_akg_path())
     p = __import__("akg", globals(), locals(), ['ms'], 0)
     func = getattr(p.ms, "compilewithjson")
 

mindspore/_extends/parallel_compile/akg_compiler/compiler.py

@@ -13,7 +13,25 @@
 # limitations under the License.
 # ============================================================================
 """Providing akg compile with json"""
+import importlib
+import os
 import sys
+
+def get_akg_path():
+    """get akg directory base path"""
+    search_res = importlib.util.find_spec("mindspore")
+    if search_res is None:
+        raise RuntimeError("Cannot find mindspore module!")
+
+    res_path = search_res.origin
+    find_pos = res_path.find("__init__.py")
+    if find_pos == -1:
+        raise RuntimeError("Find module mindspore origin file failed!")
+    akg_path = "{}_akg".format(res_path[:find_pos])
+    if not os.path.isdir(akg_path):
+        raise RuntimeError("Cannot find akg from mindspore module!")
+    return akg_path
+
 def run_compiler(op_json):
     """
     Run AKG compiler to compile op with subprocess, if this process of
@@ -25,6 +43,7 @@ def run_compiler(op_json):
     Returns:
         None
     """
+    sys.path.insert(0, get_akg_path())
     p = __import__("akg", globals(), locals(), ['ms'], 0)
     func = getattr(p.ms, "compilewithjson")
     res = func(op_json)

setup.py

@@ -179,7 +179,7 @@ class BuildPy(build_py):
         super().run()
         mindspore_dir = os.path.join(pkg_dir, 'build', 'lib', 'mindspore')
         update_permissions(mindspore_dir)
-        mindspore_dir = os.path.join(pkg_dir, 'build', 'lib', 'akg')
+        mindspore_dir = os.path.join(pkg_dir, 'build', 'lib', 'mindspore', '_akg')
         update_permissions(mindspore_dir)