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();