Fuse temp

.gitignore

@@ -12,3 +12,7 @@ CMakeCache.txt
 doc
 
 apps/tensor_times_vector/tensor_times_vector
+
+.cache
+.vscode
+compile_commands.json

CMakeLists.txt

@@ -11,10 +11,10 @@ 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)
+set(TACO_FEATURE_OPENMP 1)
 set(TACO_FEATURE_PYTHON 0)
 if(CUDA)
   message("-- Searching for CUDA Installation")

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,11 +44,16 @@ 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, std::string& sofile, void** args);
   int callFuncPackedRaw(std::string name, void** args);
 
   /// Call a raw function in this module and return the result
+  int callFuncPackedRaw(std::string name, std::string& sofile, std::vector<void*> args) {
+    return callFuncPackedRaw(name, sofile, args.data());
+  }
   int callFuncPackedRaw(std::string name, std::vector<void*> args) {
     return callFuncPackedRaw(name, args.data());
   }
@@ -57,6 +62,10 @@ class Module {
   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 +81,7 @@ class Module {
   std::string libname;
   std::string tmpdir;
   void* lib_handle;
+  void* so_lib_handle;
   std::vector<Stmt> funcs;
 
   // true iff the module was created from user-provided source

include/taco/index_notation/transformations.h

@@ -223,6 +223,9 @@ IndexStmt parallelizeOuterLoop(IndexStmt stmt);
  */
 IndexStmt reorderLoopsTopologically(IndexStmt stmt);
 
+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/codegen.cpp

@@ -229,6 +229,49 @@ string CodeGen::printTensorProperty(string varname, const GetProperty* op, bool
   return ret.str();
 }
 
+string CodeGen::getUnpackedTensorArgument(string varname, const GetProperty* op,
+                            bool is_output_prop) {
+  stringstream ret;
+  ret << "";
+
+  auto tensor = op->tensor.as<Var>();
+  if (op->property == TensorProperty::Values) {
+    // for the values, it's in the last slot
+    ret << "uniform " << printType(tensor->type, false) << " " << varname << "[]";
+    return ret.str();
+  } else if (op->property == TensorProperty::ValuesSize) {
+    ret << "int32 " << varname;
+    return ret.str();
+  }
+
+  // for a Dense level, nnz is an int
+  // for a Fixed level, ptr is an int
+  // all others are int*
+  if (op->property == TensorProperty::Dimension) {
+    if (op->type == Int32) {
+      ret << "uniform int32 ";
+    } else if (op->type == Int64) {
+      ret << "uniform int64 ";
+    } else {
+      ret << "int ";
+    }
+    ret << varname;
+
+  } else {
+    taco_iassert(op->property == TensorProperty::Indices);
+    if (op->type == Int32) {
+      ret << "uniform int32 ";
+    } else if (op->type == Int64) {
+      ret << "uniform int64 ";
+    } else {
+      ret << "uniform int ";
+    }
+    ret << varname << "[]";
+  }
+
+  return ret.str();
+}
+
 string CodeGen::unpackTensorProperty(string varname, const GetProperty* op,
                             bool is_output_prop) {
   stringstream ret;
@@ -310,13 +353,9 @@ string CodeGen::pointTensorProperty(std::string varname) {
   return ret.str();
 }
 
-// helper to print declarations
-string CodeGen::printDecls(map<Expr, string, ExprCompare> varMap,
-                           vector<Expr> inputs, vector<Expr> outputs) {
-  stringstream ret;
-  unordered_set<string> propsAlreadyGenerated;
-
-  vector<const GetProperty*> sortedProps;
+void CodeGen::getSortedProps(map<Expr, string, ExprCompare> &varMap,
+              vector<const GetProperty*> &sortedProps, vector<Expr> &inputs,
+              vector<Expr> &outputs) {
 
   for (auto const& p: varMap) {
     if (p.first.as<GetProperty>())
@@ -355,6 +394,17 @@ string CodeGen::printDecls(map<Expr, string, ExprCompare> varMap,
          return a->index < b->index;
        });
 
+}
+
+// helper to print declarations
+string CodeGen::printDecls(map<Expr, string, ExprCompare> varMap,
+                           vector<Expr> inputs, vector<Expr> outputs) {
+  stringstream ret;
+  unordered_set<string> propsAlreadyGenerated;
+
+  vector<const GetProperty*> sortedProps;
+  getSortedProps(varMap, sortedProps, inputs, outputs);
+
   for (auto prop: sortedProps) {
     bool isOutputProp = (find(outputs.begin(), outputs.end(),
                               prop->tensor) != outputs.end());
@@ -375,7 +425,6 @@ string CodeGen::printDecls(map<Expr, string, ExprCompare> varMap,
   return ret.str();
 }
 
-
 string CodeGen::printPack(map<tuple<Expr, TensorProperty, int, int>,
         string> outputProperties, vector<Expr> outputs) {
   stringstream ret;

src/codegen/codegen.h

@@ -16,7 +16,8 @@ class CodeGen : public IRPrinter {
   enum CodeGenType { C, CUDA };
 
   CodeGen(std::ostream& stream, CodeGenType type) : IRPrinter(stream), codeGenType(type) {};
-  CodeGen(std::ostream& stream, bool color, bool simplify, CodeGenType type) : IRPrinter(stream, color, simplify), codeGenType(type) {};
+  CodeGen(std::ostream& stream, bool color, bool simplify, CodeGenType type) 
+    : IRPrinter(stream, color, simplify), codeGenType(type) {};
   /// Initialize the default code generator
   static std::shared_ptr<CodeGen> init_default(std::ostream &dest, OutputKind outputKind);
 
@@ -26,6 +27,9 @@ class CodeGen : public IRPrinter {
 protected:
   static bool checkForAlloc(const Function *func);
   static int countYields(const Function *func);
+  void getSortedProps(std::map<Expr, std::string, ExprCompare> &varMap,
+              std::vector<const GetProperty*> &sortedProps, std::vector<Expr> &inputs,
+              std::vector<Expr> &outputs);
 
   static std::string printCType(Datatype type, bool is_ptr);
   static std::string printCUDAType(Datatype type, bool is_ptr);
@@ -52,6 +56,10 @@ class CodeGen : public IRPrinter {
   std::string printFuncName(const Function *func, 
           std::map<Expr, std::string, ExprCompare> inputMap={}, 
           std::map<Expr, std::string, ExprCompare> outputMap={});
+
+  std::string printTensorProperty(std::string varname, const GetProperty* op, bool is_ptr);
+  std::string getUnpackedTensorArgument(std::string varname, const GetProperty* op,
+                              bool is_output_prop); 
 
   void resetUniqueNameCounters();
   std::string genUniqueName(std::string name);
@@ -61,9 +69,8 @@ class CodeGen : public IRPrinter {
 private:
   virtual std::string restrictKeyword() const { return ""; }
 
-  std::string printTensorProperty(std::string varname, const GetProperty* op, bool is_ptr);
   std::string unpackTensorProperty(std::string varname, const GetProperty* op,
-                              bool is_output_prop);
+                              bool is_output_prop); 
   std::string packTensorProperty(std::string varname, Expr tnsr, TensorProperty property,
                             int mode, int index);
   std::string pointTensorProperty(std::string varname);

src/codegen/codegen_c.cpp

@@ -34,6 +34,7 @@ const string cHeaders =
   "#include <math.h>\n"
   "#include <complex.h>\n"
   "#include <string.h>\n"
+  "#include <omp.h>\n"
   "#if _OPENMP\n"
   "#include <omp.h>\n"
   "#endif\n"
@@ -308,6 +309,7 @@ void CodeGen_C::visit(const Function* func) {
   // output body
   print(func->body);
 
+
   // output repack only if we allocated memory
   if (checkForAlloc(func))
     out << endl << printPack(varFinder.outputProperties, func->outputs);
@@ -403,6 +405,7 @@ static string getAtomicPragma() {
 // Docs for vectorization pragmas:
 // http://clang.llvm.org/docs/LanguageExtensions.html#extensions-for-loop-hint-optimizations
 void CodeGen_C::visit(const For* op) {
+
   switch (op->kind) {
     case LoopKind::Vectorized:
       doIndent();

src/codegen/codegen_c.h

@@ -28,23 +28,24 @@ class CodeGen_C : public CodeGen {
 protected:
   using IRPrinter::visit;
 
-  void visit(const Function*);
-  void visit(const VarDecl*);
-  void visit(const Yield*);
-  void visit(const Var*);
-  void visit(const For*);
-  void visit(const While*);
-  void visit(const GetProperty*);
-  void visit(const Min*);
-  void visit(const Max*);
-  void visit(const Allocate*);
-  void visit(const Sqrt*);
-  void visit(const Store*);
-  void visit(const Assign*);
+  virtual void visit(const Function*);
+  virtual void visit(const VarDecl*);
+  virtual void visit(const Yield*);
+  virtual void visit(const Var*);
+  virtual void visit(const For*);
+  virtual void visit(const While*);
+  virtual void visit(const GetProperty*);
+  virtual void visit(const Min*);
+  virtual void visit(const Max*);
+  virtual void visit(const Allocate*);
+  virtual void visit(const Sqrt*);
+  virtual void visit(const Store*);
+  virtual void visit(const Assign*);
 
   std::map<Expr, std::string, ExprCompare> varMap;
   std::vector<Expr> localVars;
   std::ostream &out;
+  int count = 0;
 
   OutputKind outputKind;
 

src/codegen/module.cpp

@@ -42,6 +42,7 @@ void Module::addFunction(Stmt func) {
 
 void Module::compileToSource(string path, string prefix) {
   if (!moduleFromUserSource) {
+    std::cout << "module not from user source\n";
 
     // create a codegen instance and add all the funcs
     bool didGenRuntime = false;
@@ -109,6 +110,7 @@ void writeShims(vector<Stmt> funcs, string path, string prefix) {
 } // anonymous namespace
 
 string Module::compile() {
+  std::cout << "Module::compile\n";
   string prefix = tmpdir+libname;
   string fullpath = prefix + ".so";
 
@@ -137,12 +139,24 @@ string Module::compile() {
   string cmd = cc + " " + cflags + " " +
     prefix + file_ending + " " + shims_file + " " + 
     "-o " + fullpath + " -lm";
+  std::cout << "--------------------------------------------------------------------------------tmpdir: " << tmpdir << std::endl;
+  std::cout << "--------------------------------------------------------------------------------libname: " << libname << std::endl;
+  std::cout << "--------------------------------------------------------------------------------prefix: " << prefix << std::endl;
+  std::cout << "--------------------------------------------------------------------------------fullpath: " << fullpath << std::endl;
+  std::cout << "--------------------------------------------------------------------------------cmd: " << cmd << std::endl;
 
   // open the output file & write out the source
   compileToSource(tmpdir, libname);
+
 
   // write out the shims
   writeShims(funcs, tmpdir, libname);
+  for (auto &statement : funcs) {
+    std::cout << "----- statement --------" << std::endl;
+    // std::cout << statement;
+    std::cout << std::endl;
+  }
+  std::cout << tmpdir << std::endl << libname << std::endl;
 
   // now compile it
   int err = system(cmd.data());
@@ -168,10 +182,61 @@ string Module::getSource() {
   return source.str();
 }
 
+void* Module::getFuncPtr(std::string& sofile, std::string name) {
+  std::cout << "opening shared object 1\n";
+  if (so_lib_handle) {
+    dlclose(so_lib_handle);
+  }
+  std::cout << "opening shared object 2\n";
+  so_lib_handle = dlopen(sofile.data(), RTLD_NOW | RTLD_LOCAL);
+  std::cout << "opening shared object : " << sofile << std::endl;
+  return dlsym(so_lib_handle, name.data());
+}
+
 void* Module::getFuncPtr(std::string name) {
   return dlsym(lib_handle, name.data());
 }
 
+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;
+}
+
 int Module::callFuncPackedRaw(std::string name, void** args) {
   typedef int (*fnptr_t)(void**);
   static_assert(sizeof(void*) == sizeof(fnptr_t),
@@ -200,7 +265,9 @@ int Module::callFuncPackedRaw(std::string name, void** args) {
   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);