Jupyter quanto fixes (#22)

* Add methods to check for passes/analyses

* Fix xeus build

* clang-format

* Optimize JIT

* Allow jit mode to take data from file

* Fix OpenMP in JIT mode

* Update JIT main

* Update cloning

Co-authored-by: Ibrahim Numanagić <ibrahimpasa@gmail.com>

cmake/deps.cmake

@@ -155,7 +155,8 @@ if(CODON_JUPYTER)
         OPTIONS "BUILD_EXAMPLES OFF"
                 "XEUS_BUILD_SHARED_LIBS OFF"
                 "XEUS_STATIC_DEPENDENCIES ON"
-                "CMAKE_POSITION_INDEPENDENT_CODE ON")
+                "CMAKE_POSITION_INDEPENDENT_CODE ON"
+                "XEUS_DISABLE_ARCH_NATIVE ON")
     if (xeus_ADDED)
         install(TARGETS nlohmann_json EXPORT xeus-targets)
     endif()

codon/app/main.cpp

@@ -207,9 +207,26 @@ std::string jitExec(codon::jit::JIT *jit, const std::string &code) {
   }
   return *result;
 }
+
+void jitLoop(codon::jit::JIT *jit, std::istream &fp) {
+  std::string code;
+  for (std::string line; std::getline(fp, line);) {
+    if (line != "#%%") {
+      code += line + "\n";
+    } else {
+      fmt::print("{}[done]\n", jitExec(jit, code));
+      code = "";
+      fflush(stdout);
+    }
+  }
+  if (!code.empty())
+    fmt::print("{}[done]\n", jitExec(jit, code));
+}
 } // namespace
 
 int jitMode(const std::vector<const char *> &args) {
+  llvm::cl::opt<std::string> input(llvm::cl::Positional, llvm::cl::desc("<input file>"),
+                                   llvm::cl::init("-"));
   llvm::cl::list<std::string> plugins("plugin",
                                       llvm::cl::desc("Load specified plugin"));
   llvm::cl::opt<std::string> log("log", llvm::cl::desc("Enable given log streams"));
@@ -233,18 +250,12 @@ int jitMode(const std::vector<const char *> &args) {
 
   llvm::cantFail(jit.init());
   fmt::print(">>> Codon JIT v{} <<<\n", CODON_VERSION);
-  std::string code;
+  if (input == "-") {
-  for (std::string line; std::getline(std::cin, line);) {
+    jitLoop(&jit, std::cin);
-    if (line != "#%%") {
-      code += line + "\n";
   } else {
-      fmt::print("{}[done]\n", jitExec(&jit, code));
+    std::ifstream fileInput(input);
-      code = "";
+    jitLoop(&jit, fileInput);
-      fflush(stdout);
   }
-  }
-  if (!code.empty())
-    fmt::print("{}[done]\n", jitExec(&jit, code));
   return EXIT_SUCCESS;
 }
 

codon/compiler/engine.cpp

@@ -15,7 +15,7 @@ llvm::Expected<llvm::orc::ThreadSafeModule>
 Engine::optimizeModule(llvm::orc::ThreadSafeModule module,
                        const llvm::orc::MaterializationResponsibility &R) {
   module.withModuleDo([](llvm::Module &module) {
-    ir::optimize(&module, /*debug=*/true, /*jit=*/true);
+    ir::optimize(&module, /*debug=*/false, /*jit=*/true);
   });
   return std::move(module);
 }

codon/parser/visitors/simplify/simplify_stmt.cpp

@@ -1067,7 +1067,6 @@ void SimplifyVisitor::visit(ClassStmt *stmt) {
     seqassert(c, "not a class AST for {}", canonicalName);
     preamble->globals.push_back(c->clone());
     c->suite = clone(suite);
-
   }
   stmts[0] = N<ClassStmt>(canonicalName, std::vector<Param>{}, N<SuiteStmt>(),
                           Attr({Attr::Extend}), std::vector<ExprPtr>{},

codon/sir/transform/manager.h

@@ -109,6 +109,28 @@ public:
   explicit PassManager(bool debug = false, std::vector<std::string> disabled = {})
       : PassManager(debug ? Init::DEBUG : Init::RELEASE, std::move(disabled)) {}
 
+  /// Checks if the given pass is included in this manager.
+  /// @param key the pass key
+  /// @return true if manager has the given pass
+  bool hasPass(const std::string &key) {
+    for (auto &pair : passes) {
+      if (pair.first == key)
+        return true;
+    }
+    return false;
+  }
+
+  /// Checks if the given analysis is included in this manager.
+  /// @param key the analysis key
+  /// @return true if manager has the given analysis
+  bool hasAnalysis(const std::string &key) {
+    for (auto &pair : analyses) {
+      if (pair.first == key)
+        return true;
+    }
+    return false;
+  }
+
   /// Registers a pass and appends it to the execution order.
   /// @param pass the pass
   /// @param insertBefore insert pass before the pass with this given key

codon/sir/util/cloning.cpp

@@ -1,8 +1,61 @@
 #include "cloning.h"
 
+#include "codon/sir/util/operator.h"
+
 namespace codon {
 namespace ir {
 namespace util {
+namespace {
+struct GatherLocals : public util::Operator {
+  std::vector<Var *> locals;
+  void preHook(Node *node) override {
+    for (auto *v : node->getUsedVariables()) {
+      if (!v->isGlobal())
+        locals.push_back(v);
+    }
+  }
+};
+} // namespace
+
+Value *CloneVisitor::clone(const Value *other, BodiedFunc *cloneTo) {
+  if (!other)
+    return nullptr;
+
+  if (cloneTo) {
+    auto *M = cloneTo->getModule();
+    GatherLocals gl;
+    const_cast<Value *>(other)->accept(gl);
+    for (auto *v : gl.locals) {
+      auto *clonedVar = M->N<Var>(v, v->getType(), v->isGlobal(), v->getName());
+      cloneTo->push_back(clonedVar);
+      forceRemap(v, clonedVar);
+    }
+  }
+
+  auto id = other->getId();
+  if (ctx.find(id) == ctx.end()) {
+    other->accept(*this);
+    ctx[id] = result;
+
+    for (auto it = other->attributes_begin(); it != other->attributes_end(); ++it) {
+      const auto *attr = other->getAttribute(*it);
+      if (attr->needsClone()) {
+        ctx[id]->setAttribute(attr->clone(*this), *it);
+      }
+    }
+  }
+  return cast<Value>(ctx[id]);
+}
+
+Var *CloneVisitor::clone(const Var *other) {
+  if (!other)
+    return nullptr;
+  auto id = other->getId();
+  if (ctx.find(id) != ctx.end())
+    return cast<Var>(ctx[id]);
+
+  return const_cast<Var *>(other);
+}
 
 void CloneVisitor::visit(const Var *v) {
   result = module->N<Var>(v, v->getType(), v->isGlobal(), v->getName());

codon/sir/util/cloning.h

@@ -12,7 +12,7 @@ private:
   /// the clone context
   std::unordered_map<int, Node *> ctx;
   /// the result
-  Node *result = nullptr;
+  Node *result;
   /// the module
   Module *module;
   /// true if break/continue loops should be cloned
@@ -22,7 +22,8 @@ public:
   /// Constructs a clone visitor.
   /// @param module the module
   /// @param cloneLoop true if break/continue loops should be cloned
-  explicit CloneVisitor(Module *module, bool cloneLoop = true) : module(module) {}
+  explicit CloneVisitor(Module *module, bool cloneLoop = true)
+      : ctx(), result(nullptr), module(module), cloneLoop(cloneLoop) {}
 
   virtual ~CloneVisitor() noexcept = default;
 
@@ -69,38 +70,14 @@ public:
 
   /// Clones a value, returning the previous value if other has already been cloned.
   /// @param other the original
+  /// @param cloneTo the function to clone locals to, or null if none
   /// @return the clone
-  Value *clone(const Value *other) {
+  Value *clone(const Value *other, BodiedFunc *cloneTo = nullptr);
-    if (!other)
-      return nullptr;
-
-    auto id = other->getId();
-    if (ctx.find(id) == ctx.end()) {
-      other->accept(*this);
-      ctx[id] = result;
-
-      for (auto it = other->attributes_begin(); it != other->attributes_end(); ++it) {
-        const auto *attr = other->getAttribute(*it);
-        if (attr->needsClone()) {
-          ctx[id]->setAttribute(attr->clone(*this), *it);
-        }
-      }
-    }
-    return cast<Value>(ctx[id]);
-  }
 
   /// Returns the original unless the variable has been force cloned.
   /// @param other the original
   /// @return the original or the previous clone
-  Var *clone(const Var *other) {
+  Var *clone(const Var *other);
-    if (!other)
-      return nullptr;
-    auto id = other->getId();
-    if (ctx.find(id) != ctx.end())
-      return cast<Var>(ctx[id]);
-
-    return const_cast<Var *>(other);
-  }
 
   /// Clones a flow, returning the previous value if other has already been cloned.
   /// @param other the original

codon/sir/util/irtools.cpp

@@ -113,8 +113,10 @@ VarValue *makeVar(Value *x, SeriesFlow *flow, BodiedFunc *parent, bool prepend)
   const bool global = (parent == nullptr);
   auto *M = x->getModule();
   auto *v = M->Nr<Var>(x->getType(), global);
-  if (global)
+  if (global) {
-    v->setName("_anon_global");
+    static int counter = 1;
+    v->setName("_anon_global_" + std::to_string(counter++));
+  }
   auto *a = M->Nr<AssignInstr>(v, x);
   if (prepend) {
     flow->insert(flow->begin(), a);

stdlib/openmp.codon

@@ -323,23 +323,29 @@ def _fork_call(microtask: cobj, args):
 
 
 def _static_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
+    @nonpure
     def _loop_step():
         return 1
 
+    @nonpure
     def _loop_loc_and_gtid(
         loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int
     ):
         pass
 
+    @nonpure
     def _loop_body_stub(i, args):
         pass
 
+    @nonpure
     def _loop_schedule():
         return (1 << 30) | 35  # nonmonotonic, dynamic chunked
 
+    @nonpure
     def _loop_shared_updates(args):
         pass
 
+    @nonpure
     def _loop_reductions(args):
         pass
 
@@ -371,23 +377,29 @@ def _static_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
 
 
 def _static_chunked_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
+    @nonpure
     def _loop_step():
         return 1
 
+    @nonpure
     def _loop_loc_and_gtid(
         loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int
     ):
         pass
 
+    @nonpure
     def _loop_body_stub(i, args):
         pass
 
+    @nonpure
     def _loop_schedule():
         return (1 << 30) | 35  # nonmonotonic, dynamic chunked
 
+    @nonpure
     def _loop_shared_updates(args):
         pass
 
+    @nonpure
     def _loop_reductions(args):
         pass
 
@@ -425,26 +437,33 @@ def _static_chunked_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32]
 
 
 def _dynamic_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
+    @nonpure
     def _loop_step():
         return 1
 
+    @nonpure
     def _loop_loc_and_gtid(
         loc_ref: Ptr[Ident], reduction_loc_ref: Ptr[Ident], gtid: int
     ):
         pass
 
+    @nonpure
     def _loop_body_stub(i, args):
         pass
 
+    @nonpure
     def _loop_schedule():
         return (1 << 30) | 35  # nonmonotonic, dynamic chunked
 
+    @nonpure
     def _loop_shared_updates(args):
         pass
 
+    @nonpure
     def _loop_reductions(args):
         pass
 
+    @nonpure
     def _loop_ordered():
         return False
 
@@ -510,6 +529,7 @@ def _spawn_and_run_task(
 # spawns a new task for each loop iteration.
 def _task_loop_outline_template(gtid_ptr: Ptr[i32], btid_ptr: Ptr[i32], args):
     def _routine_stub(gtid: i32, data: cobj, P: type, S: type):
+        @nonpure
         def _task_loop_body_stub(priv, shared):
             pass