task: implement kernel fusion over distribution

.gitignore

@@ -12,3 +12,6 @@ CMakeCache.txt
 doc
 
 apps/tensor_times_vector/tensor_times_vector
+
+.CMakeCache
+.vscode

CMakeLists.txt

@@ -11,7 +11,7 @@ project(taco
 )
 option(CUDA "Build for NVIDIA GPU (CUDA must be preinstalled)" OFF)
 option(PYTHON "Build TACO for python environment" OFF)
-option(OPENMP "Build with OpenMP execution support" OFF)
+option(OPENMP "Build with OpenMP execution support" ON)
 option(COVERAGE "Build with code coverage analysis" OFF)
 set(TACO_FEATURE_CUDA 0)
 set(TACO_FEATURE_OPENMP 0)

include/taco/codegen/module.h

@@ -17,7 +17,7 @@ class Module {
 public:
   /// Create a module for some target
   Module(Target target=getTargetFromEnvironment())
-    : lib_handle(nullptr), moduleFromUserSource(false), target(target) {
+    : lib_handle(nullptr), so_lib_handle(nullptr), moduleFromUserSource(false), target(target) {
     setJITLibname();
     setJITTmpdir();
   }
@@ -44,19 +44,27 @@ class Module {
   /// before calling. If there's no function of this name then a nullptr is
   /// returned.
   void* getFuncPtr(std::string name);
+  void* getFuncPtr(std::string& sofile, std::string name);
 
   /// Call a raw function in this module and return the result
   int callFuncPackedRaw(std::string name, void** args);
+  int callFuncPackedRaw(std::string name, std::string&sofile, void** args);
 
   /// Call a raw function in this module and return the result
   int callFuncPackedRaw(std::string name, std::vector<void*> args) {
     return callFuncPackedRaw(name, args.data());
   }
+  int callFuncPackedRaw(std::string name, std::string& sofile, std::vector<void*> args) {
+    return callFuncPackedRaw(name, sofile, args.data());
+  }
 
   /// Call a function using the taco_tensor_t interface and return the result
   int callFuncPacked(std::string name, void** args) {
     return callFuncPackedRaw("_shim_"+name, args);
   }
+  int callFuncPacked(std::string name, std::string& sofile, void** args) {
+    return callFuncPackedRaw("__shim__"+name, sofile, args);
+  }
 
   /// Call a function using the taco_tensor_t interface and return the result
   int callFuncPacked(std::string name, std::vector<void*> args) {
@@ -72,6 +80,7 @@ class Module {
   std::string libname;
   std::string tmpdir;
   void* lib_handle;
+  void* so_lib_handle; // to pass the manually compiles so (shared object) code file
   std::vector<Stmt> funcs;
 
   // true iff the module was created from user-provided source

include/taco/index_notation/transformations.h

@@ -223,6 +223,19 @@ IndexStmt parallelizeOuterLoop(IndexStmt stmt);
  */
 IndexStmt reorderLoopsTopologically(IndexStmt stmt);
 
+/**
+ * @brief Transform topologically reordered iteration graph to 
+ * a branched version
+ * 
+ * @param stmt topologically sorted index statement
+ * @param assignment assignment statement of the base statement
+ * @param side side in which the operation is performed front/back
+ * @param iters number of iterations
+ * @return IndexStmt 
+ */
+IndexStmt loopFusionOverFission(IndexStmt stmt, Assignment assignment,
+  std::string side, int iters);
+
 /**
  * Performs scalar promotion so that reductions are done by accumulating into 
  * scalar temporaries whenever possible.

include/taco/tensor.h

@@ -413,6 +413,8 @@ class TensorBase {
 
   /// Compile the tensor expression.
   void compile();
+  void compute(std::ofstream& statfile);
+  void compute(std::ofstream& statfile, std::string& sofile);
 
   void compile(IndexStmt stmt, bool assembleWhileCompute=false);
 

src/codegen/module.cpp

@@ -138,6 +138,12 @@ string Module::compile() {
     prefix + file_ending + " " + shims_file + " " + 
     "-o " + fullpath + " -lm";
 
+  std::cout << "--- tmpdir: " << tmpdir << std::endl
+    << "--- libname: " << libname << std::endl
+    << "--- prefix: " << prefix << std::endl
+    << "--- fullpath: " << fullpath << std::endl
+    << "--- cmd: " << cmd << std::endl;
+
   // open the output file & write out the source
   compileToSource(tmpdir, libname);
 
@@ -172,6 +178,15 @@ void* Module::getFuncPtr(std::string name) {
   return dlsym(lib_handle, name.data());
 }
 
+void* Module::getFuncPtr(std::string& sofile, std::string name) {
+  std::cout << "opening shared object " << sofile << std::endl;
+  if (so_lib_handle) {
+    dlclose(so_lib_handle);
+  }
+  so_lib_handle = dlopen(sofile.data(), RTLD_NOW | RTLD_LOCAL);
+  return dlsym(so_lib_handle, name.data());
+}
+
 int Module::callFuncPackedRaw(std::string name, void** args) {
   typedef int (*fnptr_t)(void**);
   static_assert(sizeof(void*) == sizeof(fnptr_t),
@@ -210,5 +225,45 @@ int Module::callFuncPackedRaw(std::string name, void** args) {
   return ret;
 }
 
+int Module::callFuncPackedRaw(std::string name, std::string& sofile, void** args) {
+  typedef int (*fnptr_t)(void**);
+  static_assert(sizeof(void*) == sizeof(fnptr_t),
+    "Unable to cast dlsym() returned void pointer to function pointer");
+  void* v_func_ptr = getFuncPtr(sofile, name);
+  fnptr_t func_ptr;
+  *reinterpret_cast<void**>(&func_ptr) = v_func_ptr;
+
+#if USE_OPENMP
+  omp_sched_t existingSched;
+  ParallelSchedule tacoSched;
+  int existingChunkSize, tacoChunkSize;
+  int existingNumThreads = omp_get_max_threads();
+  omp_get_schedule(&existingSched, &existingChunkSize);
+  taco_get_parallel_schedule(&tacoSched, &tacoChunkSize);
+  switch (tacoSched) {
+    case ParallelSchedule::Static:
+      omp_set_schedule(omp_sched_static, tacoChunkSize);
+      break;
+    case ParallelSchedule::Dynamic:
+      omp_set_schedule(omp_sched_dynamic, tacoChunkSize);
+      break;
+    default:
+      break;
+  }
+  omp_set_num_threads(taco_get_num_threads());
+#endif
+
+  std::cout << "calling the function\n";
+  int ret = func_ptr(args);
+  std::cout << "function call completed\n";
+
+#if USE_OPENMP
+  omp_set_schedule(existingSched, existingChunkSize);
+  omp_set_num_threads(existingNumThreads);
+#endif
+
+  return ret;
+}
+
 } // namespace ir
 } // namespace taco