diff --git a/src/codegen/build_module.cc b/src/codegen/build_module.cc index e1c258f046c2fd859510cc690ef275e2ef0d8b1d..d4fbc6bf67c27393a1dc1fc14a7648f8ac0e75a3 100644 --- a/src/codegen/build_module.cc +++ b/src/codegen/build_module.cc @@ -23,6 +23,8 @@ #include #include #include +#include +#include #include "build_module.h" #include "pass/expr_alg_simplify.h" @@ -1059,13 +1061,32 @@ BuildRst BuildToFunc(const Schedule &inputs, const Array &in_args, cons } namespace { -void CreateCce(const std::string &code, const std::string &kernel_name) { - std::string file_name = kMsDavinciKernelPath; - file_name.append(kernel_name).append(".cce"); - std::ofstream of(file_name); - CHECK(of.is_open()) << "Failed to open " << file_name << " to dump cce."; - of << code << std::endl; - of.close(); +void CreateCode(const std::string &code, const std::string &kernel_name, const std::string &target_name) { + std::string file_path; + std::string file_suffix; + if (target_name.find("cce") != std::string::npos) { + file_path = std::string(kMsDavinciKernelPath); + file_suffix = ".cce"; + } else if (target_name.find("cuda") != std::string::npos) { + file_path = std::string(kMsGpuKernelPath) + "_" + std::to_string(getpid()) + "/"; + file_suffix = ".cu"; + } + + if (file_path.empty()) { + return; + } + + // Dump code to meta directory if it exists. + struct stat info; + if (stat(file_path.c_str(), &info) == 0) { + if (info.st_mode & S_IFDIR) { + std::string file_name = file_path + kernel_name + file_suffix; + std::ofstream of(file_name); + CHECK(of.is_open()) << "Failed to open " << file_name << " to dump code."; + of << code << std::endl; + of.close(); + } + } } } // namespace @@ -1115,7 +1136,7 @@ air::runtime::Module BuildToModule(const NodeRef &ref, const std::string &target auto mod0 = mhost->imports()[0]; CHECK(mod0.defined()); - CreateCce(mod0->GetSource(), build_rst->kernel_name); + CreateCode(mod0->GetSource(), build_rst->kernel_name, target_name); } return mhost; diff --git a/src/composite/util.h b/src/composite/util.h index 170c6ee7f7205d464a14bb172323ed1b1ded5c32..2b0ac58df1b0edbe978ee1007dd6ff3feddb0971 100644 --- a/src/composite/util.h +++ b/src/composite/util.h @@ -22,6 +22,7 @@ namespace akg { constexpr auto kMsDavinciKernelPath = "./kernel_meta/"; +constexpr auto kMsGpuKernelPath = "./cuda_meta"; static std::unordered_map type_mapping = { {"float32", air::Float(32)}, {"float16", air::Float(16)}, {"int32", air::Int(32)}, {"bool", air::Bool()}}; } // namespace akg