提交 60f105c9 编写于 作者: M mindspore-ci-bot 提交者: Gitee

!108 use akg build insteand of tvm build to dump ir in kernel named folder.

Merge pull request !108 from lvwenyuan/gpu-env
...@@ -15,7 +15,9 @@ ...@@ -15,7 +15,9 @@
# limitations under the License. # limitations under the License.
"""build module""" """build module"""
import os
import json import json
import akg
from akg import tvm from akg import tvm
from akg.tvm import _api_internal from akg.tvm import _api_internal
from .repository import __all__ as repository from .repository import __all__ as repository
...@@ -178,6 +180,6 @@ def build_cuda(outputs, args, sch_name, kernel_name): ...@@ -178,6 +180,6 @@ def build_cuda(outputs, args, sch_name, kernel_name):
s = scheduler[sch_name](outputs) s = scheduler[sch_name](outputs)
dump_ir = os.getenv('MS_AKG_DUMP_IR') == "on" dump_ir = os.getenv('MS_AKG_DUMP_IR') == "on"
with tvm.build_config(dump_pass_ir = dump_ir): with tvm.build_config(dump_pass_ir = dump_ir):
mod = tvm.build(s, args, cuda, name = kernel_name) mod = akg.build(s, list(args), "cuda", name = kernel_name)
dump_cuda_meta.dump(mod, kernel_name, s, list(args)) dump_cuda_meta.dump(mod, kernel_name, s, list(args))
return mod return mod
...@@ -20,7 +20,7 @@ import json ...@@ -20,7 +20,7 @@ import json
import json.decoder as jd import json.decoder as jd
import logging import logging
import traceback import traceback
import os.path import os
from pathlib import Path from pathlib import Path
import akg.tvm import akg.tvm
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册