From 7f8c843f5bd8456c7b4359641dfe98a8c6b76558 Mon Sep 17 00:00:00 2001 From: Sam Kaplan Date: Tue, 13 Sep 2022 09:45:20 -0400 Subject: [PATCH] Add persistent kernel cache on disk This adds another layer of kernel caching, that lives on disk and persists between program runs. This includes the generated source, header, and compiled object files for each kernel. The file name is a hash of the names, types, and formats of each tensor, plus a serialization of the IndexStmt. This should be sufficient to uniquely identify a kernel. When the environment variable TACO_CACHE_DIR is set, the value is used as the cache directory. If unset, the disk-level cache is disabled. --- include/taco/codegen/module.h | 15 ++++- include/taco/index_notation/index_notation.h | 6 ++ src/codegen/module.cpp | 63 ++++++++++++++++---- src/index_notation/index_notation.cpp | 16 +++++ src/tensor.cpp | 2 +- 5 files changed, 86 insertions(+), 16 deletions(-) diff --git a/include/taco/codegen/module.h b/include/taco/codegen/module.h index 36eb34f1a..ff77b8202 100644 --- a/include/taco/codegen/module.h +++ b/include/taco/codegen/module.h @@ -6,6 +6,7 @@ #include #include #include +#include #include "taco/target.h" #include "taco/ir/ir.h" @@ -16,10 +17,18 @@ namespace ir { class Module { public: /// Create a module for some target - Module(Target target=getTargetFromEnvironment()) - : lib_handle(nullptr), moduleFromUserSource(false), target(target) { + Module(const std::string& cacheStr = "", Target target=getTargetFromEnvironment()) + : cacheStr(cacheStr), lib_handle(nullptr), moduleFromUserSource(false), target(target) { setJITLibname(); setJITTmpdir(); + + if (cacheStr != "") { + std::hash hasher; + size_t hash = hasher(cacheStr); + std::ostringstream s; + s << std::hex << hash; + cacheStrHashed = s.str(); + } } /// Compile the source into a library, returning its full path @@ -67,6 +76,8 @@ class Module { void setSource(std::string source); private: + std::string cacheStr; + std::string cacheStrHashed; std::stringstream source; std::stringstream header; std::string libname; diff --git a/include/taco/index_notation/index_notation.h b/include/taco/index_notation/index_notation.h index dd451b337..6928c8035 100644 --- a/include/taco/index_notation/index_notation.h +++ b/include/taco/index_notation/index_notation.h @@ -814,6 +814,9 @@ class IndexStmt : public util::IntrusivePtr { SubType as() { return to(*this); } + + /// Get string to use for this index statement in cache. + std::string getCacheString() const; }; /// Check if two index statements are isomorphic. @@ -1232,6 +1235,9 @@ class TensorVar : public util::Comparable { friend bool operator==(const TensorVar&, const TensorVar&); friend bool operator<(const TensorVar&, const TensorVar&); + /// Get string to use for this tensor in cache. + std::string getCacheString() const; + private: struct Content; std::shared_ptr content; diff --git a/src/codegen/module.cpp b/src/codegen/module.cpp index 89738f22d..df06bdcda 100644 --- a/src/codegen/module.cpp +++ b/src/codegen/module.cpp @@ -68,11 +68,15 @@ void Module::compileToSource(string path, string prefix) { ofstream source_file; string file_ending = should_use_CUDA_codegen() ? ".cu" : ".c"; source_file.open(path+prefix+file_ending); + taco_uassert((bool)source_file) << "Could not open file '" + << path+prefix+file_ending << "' for writing"; source_file << source.str(); source_file.close(); ofstream header_file; header_file.open(path+prefix+".h"); + taco_uassert((bool)header_file) << "Could not open file '" + << path+prefix+".h" << "' for writing"; header_file << header.str(); header_file.close(); } @@ -97,9 +101,13 @@ void writeShims(vector funcs, string path, string prefix) { ofstream shims_file; if (should_use_CUDA_codegen()) { shims_file.open(path+prefix+"_shims.cpp"); + taco_uassert((bool)shims_file) << "Could not open file '" + << path+prefix+"_shims.cpp" << "' for writing"; } else { shims_file.open(path+prefix+".c", ios::app); + taco_uassert((bool)shims_file) << "Could not open file '" + << path+prefix+".c" << "' for writing"; } shims_file << "#include \"" << path << prefix << ".h\"\n"; shims_file << shims.str(); @@ -109,7 +117,14 @@ void writeShims(vector funcs, string path, string prefix) { } // anonymous namespace string Module::compile() { - string prefix = tmpdir+libname; + string cache_dir = util::getFromEnv("TACO_CACHE_DIR", ""); + bool use_cache = cache_dir != "" && cacheStr != ""; + if (use_cache && cache_dir.back() != '/') { + cache_dir += '/'; + } + string dir = use_cache ? cache_dir : tmpdir; + string name = use_cache ? cacheStrHashed : libname; + string prefix = dir+name; string fullpath = prefix + ".so"; string cc; @@ -140,21 +155,43 @@ string Module::compile() { file_ending = ".c"; shims_file = ""; } + + string source_file; + string object_file; + if (use_cache) { + source_file = cache_dir + name + file_ending; + object_file = cache_dir + name + ".so"; + } else { + source_file = prefix + file_ending; + object_file = fullpath; + } string cmd = cc + " " + cflags + " " + - prefix + file_ending + " " + shims_file + " " + - "-o " + fullpath + " -lm"; + source_file + " " + shims_file + " " + + "-o " + object_file + " -lm"; - // open the output file & write out the source - compileToSource(tmpdir, libname); - - // write out the shims - writeShims(funcs, tmpdir, libname); - - // now compile it - int err = system(cmd.data()); - taco_uassert(err == 0) << "Compilation command failed:\n" << cmd - << "\nreturned " << err; + bool cached = false; + if (use_cache) { + // first check if this file already exists in cache + ifstream cached_source_file(source_file); + ifstream cached_header_file(cache_dir + name + ".h"); + ifstream cached_object_file(object_file); + // only run codegen if the files don't already exist + cached = cached_source_file.good() && cached_header_file.good() && cached_object_file.good(); + } + + if (!cached) { + // open the output file & write out the source + compileToSource(dir, name); + + // write out the shims + writeShims(funcs, dir, name); + + // now compile it + int err = system(cmd.data()); + taco_uassert(err == 0) << "Compilation command failed:\n" << cmd + << "\nreturned " << err; + } // use dlsym() to open the compiled library if (lib_handle) { diff --git a/src/index_notation/index_notation.cpp b/src/index_notation/index_notation.cpp index e38e3d2d3..2c214f9a1 100644 --- a/src/index_notation/index_notation.cpp +++ b/src/index_notation/index_notation.cpp @@ -2074,6 +2074,15 @@ IndexStmt IndexStmt::wsaccel(TensorVar& ws, bool shouldAccel, const std::vector< return *this; } +std::string IndexStmt::getCacheString() const { + std::ostringstream s; + for (auto& var : getTensorVars(*this)) { + s << var.getCacheString() << "|"; + } + s << *this; + return s.str(); +} + std::ostream& operator<<(std::ostream& os, const IndexStmt& expr) { if (!expr.defined()) return os << "IndexStmt()"; IndexNotationPrinter printer(os); @@ -2694,6 +2703,13 @@ std::ostream& operator<<(std::ostream& os, const TensorVar& var) { return os << var.getName() << " : " << var.getType(); } +string TensorVar::getCacheString() const { + ostringstream s; + s << getName() << ":"; + s << getType().getDataType() << ";"; + s << getFormat(); + return s.str(); +} static bool isValid(Assignment assignment, string* reason) { if (reason == nullptr) { diff --git a/src/tensor.cpp b/src/tensor.cpp index 257c396c3..310ba8e3b 100644 --- a/src/tensor.cpp +++ b/src/tensor.cpp @@ -666,7 +666,7 @@ void TensorBase::compile(taco::IndexStmt stmt, bool assembleWhileCompute) { // If we have to recompile the kernel, we need to create a new Module. Since // the module we are holding on to could have been retrieved from the cache, // we can't modify it. - content->module = make_shared(); + content->module = make_shared(stmtToCompile.getCacheString()); content->module->addFunction(content->assembleFunc); content->module->addFunction(content->computeFunc); content->module->compile();