From cb05334f081740ea7c2ec06c656b2418e1de26ca Mon Sep 17 00:00:00 2001
From: tronzhang <zhangzhaochuang@huawei.com>
Date: Mon, 15 Mar 2021 11:43:54 +0800
Subject: [PATCH] change akg package path to mindspore

---
 akg                                           |  2 +-
 cmake/package.cmake                           |  8 ++++++--
 .../akg_compiler/akg_process.py               | 18 +++++++++++++++++-
 .../parallel_compile/akg_compiler/compiler.py | 19 +++++++++++++++++++
 setup.py                                      |  2 +-
 5 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/akg b/akg
index 94cb709ecaf..e2a0a264a0b 160000
--- a/akg
+++ b/akg
@@ -1 +1 @@
-Subproject commit 94cb709ecaf5d1d869883dfe80cee7497dd0692c
+Subproject commit e2a0a264a0be549b51a035e5f783927052f8ead8
diff --git a/cmake/package.cmake b/cmake/package.cmake
index 55cdac65b63..3f26e32d8b5 100644
--- a/cmake/package.cmake
+++ b/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()
diff --git a/mindspore/_extends/parallel_compile/akg_compiler/akg_process.py b/mindspore/_extends/parallel_compile/akg_compiler/akg_process.py
index bab42d8b161..c5ccc13a4cd 100644
--- a/mindspore/_extends/parallel_compile/akg_compiler/akg_process.py
+++ b/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")
 
diff --git a/mindspore/_extends/parallel_compile/akg_compiler/compiler.py b/mindspore/_extends/parallel_compile/akg_compiler/compiler.py
index de78aad7e49..7aae1db2a92 100644
--- a/mindspore/_extends/parallel_compile/akg_compiler/compiler.py
+++ b/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)
diff --git a/setup.py b/setup.py
index d7db3c9dbb0..b5ebc21bd67 100644
--- a/setup.py
+++ b/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)