Commits

ebo committed c2e2659

Added better handling of locals

  • Participants
  • Parent commits dbf9fd2

Comments (0)

Files changed (1)

 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_binop.cc
-@@ -0,0 +1,416 @@
+@@ -0,0 +1,428 @@
 +#include "JIT/jit_binop.h"
 +#include "JIT/jit_types.h"
 +#include "JIT/jit_tracer.h"
 +
 +  class BinaryFloorDivideImpl : public BinaryOpImpl {
 +  public:
-+    BinaryFloorDivideImpl(JITValue *arg0, JITValue *arg1, JITStack &stack)
++    BinaryFloorDivideImpl(JITValue *arg0, JITValue *arg1, JITStack &stack, JITLocals &locals)
 +      : arg0_(arg0), arg1_(arg1), float_inline_(false),
-+        stack_(stack)
++        stack_(stack), locals_(locals)
 +    {
 +      if (arg0_->getStaticType()->isFloat() && arg1_->getStaticType()->isFloat()) {
 +        float_inline_ = true;
 +      builder.CreateCondBr(isnull, bail, trace_cont);
 +
 +      builder.SetInsertPoint(bail);
++      std::vector<JITValue *> &locals = locals_.locals;
++      for (size_t i = 0; i < locals.size(); ++i) {
++        if (locals[i] != NULL)
++          fbuilder.StoreLocal(i, locals[i]->getStaticType()->getValue(fbuilder, false, false));
++      }
 +      std::vector<JITValue *> stack = stack_.stack;;
 +      for (size_t i = 0; i < stack.size(); ++i) {
 +        fbuilder.Push(stack[i]->getStaticType()->getValue(fbuilder, false, false));
 +    JITValue *arg1_;
 +    bool float_inline_;
 +    JITStack stack_;
++    JITLocals locals_;
 +  };
 +
 +  class BinaryTrueDivideImpl : public BinaryOpImpl {
 +  public:
-+    BinaryTrueDivideImpl(JITValue *arg0, JITValue *arg1, JITStack &stack)
++    BinaryTrueDivideImpl(JITValue *arg0, JITValue *arg1, JITStack &stack, JITLocals &locals)
 +      : arg0_(arg0), arg1_(arg1), float_inline_(false),
-+        stack_(stack)
++        stack_(stack), locals_(locals)
 +    {
 +      if (arg0_->getStaticType()->isFloat() && arg1_->getStaticType()->isFloat()) {
 +        float_inline_ = true;
 +      builder.CreateCondBr(isnull, bail, trace_cont);
 +
 +      builder.SetInsertPoint(bail);
-+      std::vector<JITValue *> stack = stack_.stack;;
++      std::vector<JITValue *> &locals = locals_.locals;
++      for (size_t i = 0; i < locals.size(); ++i) {
++        if (locals[i] != NULL)
++          fbuilder.StoreLocal(i, locals[i]->getStaticType()->getValue(fbuilder, false, false));
++      }
++      std::vector<JITValue *> &stack = stack_.stack;;
 +      for (size_t i = 0; i < stack.size(); ++i) {
 +        fbuilder.Push(stack[i]->getStaticType()->getValue(fbuilder, false, false));
 +      }
 +    JITValue *arg1_;
 +    bool float_inline_;
 +    JITStack stack_;
++    JITLocals locals_;
 +  };
 +
 +  BinaryOpImpl *BinaryOpImpl::create(int opcode, JITValue *arg0, JITValue *arg1,
-+                                     JITStack &stack)
++                                     JITStack &stack, JITLocals &locals)
 +  {
 +    switch (opcode) {
 +    case INPLACE_ADD:
 +
 +    case INPLACE_FLOOR_DIVIDE:
 +    case BINARY_FLOOR_DIVIDE:
-+      return new BinaryFloorDivideImpl(arg0, arg1, stack);
++      return new BinaryFloorDivideImpl(arg0, arg1, stack, locals);
 +
 +    case INPLACE_TRUE_DIVIDE:
 +    case BINARY_TRUE_DIVIDE:
-+      return new BinaryTrueDivideImpl(arg0, arg1, stack);
++      return new BinaryTrueDivideImpl(arg0, arg1, stack, locals);
 +    default:
 +      break;
 +    }
 +  class BinaryOpImpl {
 +  public:
 +    static BinaryOpImpl *create(int opcode, JITValue *arg0, JITValue *arg1,
-+                                JITStack &stack);
++                                JITStack &stack, JITLocals &locals);
 +
 +    virtual bool inlined() const = 0;
 +
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_env.cc
-@@ -0,0 +1,49 @@
+@@ -0,0 +1,70 @@
 +#include "Python.h"
 +
 +#include "global_jit_data_fwd.h"
 +    return v;
 +  }
 +
++  JITValue *
++  JITLocals::Load(int num)
++  {
++    if (static_cast<size_t>(num) >= locals.size()) {
++      locals.resize(num + 1);
++    }
++    if (locals[num] == NULL) {
++      JITLocalValue *v = new JITLocalValue(num);
++      locals[num] = v;
++      loaded_locals.push_back(v);
++    }
++    return locals[num];
++  }
++
++  void
++  JITLocals::Store(int num, JITValue *val)
++  {
++    assert(locals.size() > static_cast<size_t>(num));
++    locals[num] = val;
++  }
++
 +}
 diff --git a/JIT/jit_env.h b/JIT/jit_env.h
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_env.h
-@@ -0,0 +1,36 @@
+@@ -0,0 +1,50 @@
 +// -*- C++ -*-
 +#ifndef PYTHON_JIT_ENV_H_
 +#define PYTHON_JIT_ENV_H_
 +
 +  class JITValue;
 +  class JITStackValue;
++  class JITLocalValue;
 +
 +  class JITStack {
 +  public:
 +    int stack_counter;
 +  };
 +
++  class JITLocals {
++  public:
++    JITLocals(int num)
++      : locals(num, NULL)
++    {}
++
++    JITValue *Load(int);
++    void Store(int, JITValue *);
++
++    std::vector<JITValue *> locals;
++    std::vector<JITLocalValue *> loaded_locals;
++  };
++
 +
 +}
 +
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_fbuilder.cc
-@@ -0,0 +1,171 @@
+@@ -0,0 +1,179 @@
 +#include "Python.h"
 +#include "opcode.h"
 +
 +    return top;
 +  }
 +
++  void
++  JITFunctionBuilder::StoreLocal(int num, llvm::Value *value)
++  {
++    Value *ptr = this->LocalPtr(num);
++    this->builder_.CreateStore(value, ptr);
++  }
++
++  llvm::Value *
++  JITFunctionBuilder::LoadLocal(int num)
++  {
++    Value *ptr = this->LocalPtr(num);
++    return this->builder_.CreateLoad(ptr);
++  }
++
 +  llvm::Value *
 +  JITFunctionBuilder::LocalPtr(int arg)
 +  {
 +  }
 +
 +  void
-+  JITFunctionBuilder::IsPythonTrue(llvm::Value *value)
-+  {
-+    // this->state_->CreateAllocaInEntryBlock();
-+  }
-+
-+  void
 +  JITFunctionBuilder::OutputInfo(Value *v)
 +  {
 +    Function *func = this->state_->GetGlobalFunction<void(PyObject*)>("output_info");
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_fbuilder.h
-@@ -0,0 +1,72 @@
+@@ -0,0 +1,73 @@
 +// -*- C++ -*-
 +#ifndef PYTHON_JIT_FBUILDER_H_
 +#define PYTHON_JIT_FBUILDER_H_
 +    llvm::Value *Top();
 +    void Push(llvm::Value *value);
 +
++    void StoreLocal(int num, llvm::Value *value);
++    llvm::Value *LoadLocal(int num);
++
 +    void SetLastI(int v);
 +    void SetBailReason(int v);
 +    void CreateBailBr();
 +    llvm::Value *LocalPtr(int arg);
 +    llvm::Value *ConstPtr(int arg);
 +
-+    void IsPythonTrue(llvm::Value *value);
-+
 +  private:
 +    void CreateBailBlock();
 +
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_tracer.cc
-@@ -0,0 +1,673 @@
+@@ -0,0 +1,691 @@
 +#include "Python.h"
 +#include "opcode.h"
 +
 +
 +  if (trace_.size() > 500) {
 +    suspend();
++    unsuspend();
 +    return;
 +  }
 +}
 +    for (size_t i = 0; i < loaded_stack_.size(); ++i) {
 +      loaded_stack_[i]->emit(builder);
 +    }
++    for (size_t i = 0; i < loaded_locals_.size(); ++i) {
++      loaded_locals_[i]->emit(builder);
++    }
 +    for (size_t i = 0; i < trace_.size(); ++i) {
 +      trace_[i]->emit(builder);
 +    }
 +  }
 +
 +  void
-+  JITTraceAnalysis::createBinaryOp(bool inplace, JITStack &stack, TraceEntry &entry)
++  JITTraceAnalysis::createBinaryOp(bool inplace, JITStack &stack, JITLocals &locals, TraceEntry &entry)
 +  {
 +    int opcode = entry.opcode;
 +    JITValue *w = stack.Top();
 +    if (entry.arg1 != NULL)
 +      w->registerStaticType(entry.arg1);
 +
-+    BinaryOpImpl *impl = BinaryOpImpl::create(opcode, v, w, stack);
++    BinaryOpImpl *impl = BinaryOpImpl::create(opcode, v, w, stack, locals);
 +    assert(impl != NULL);
 +
 +    JITType *type_v = v->getStaticType();
 +    if (impl->inlined()) {
 +      if (!type_v->hasGuard()) {
 +        type_v->setGuard(true);
-+        this->createGuardType(v, BH_NO_BLACKHOLE, entry.pc, stack);
++        this->createGuardType(v, BH_NO_BLACKHOLE, entry.pc, stack, locals);
 +      }
 +
 +      if (!type_w->hasGuard()) {
 +        type_w->setGuard(true);
-+        this->createGuardType(w, BH_NO_BLACKHOLE, entry.pc, stack);
++        this->createGuardType(w, BH_NO_BLACKHOLE, entry.pc, stack, locals);
 +      }
 +    }
 +
 +    JITValue *x = this->createGenericBinOp(impl, inplace, v, w);
 +
 +    assert(x != NULL);
-+    this->createGuardNull(x, BH_BINARY_OP_FAIL, entry.pc, stack);
++    this->createGuardNull(x, BH_BINARY_OP_FAIL, entry.pc, stack, locals);
 +    stack.Push(x);
 +  }
 +
 +  JITTraceAnalysis::analyseTrace(std::vector<TraceEntry> &trace)
 +  {
 +    JITStack stack;
++    JITLocals locals(0);
 +    for (size_t i = 0; i < trace.size(); ++i) {
 +      TraceEntry &entry = trace[i];
 +      int opcode = entry.opcode;
 +      int flag = entry.flag;
 +
-+      this->createOpcodeInfo(entry.pc, stack);
++      this->createOpcodeInfo(entry.pc, stack, locals);
 +
 +      JITValue *w = NULL;
 +      JITValue *v = NULL;
 +      case POP_JUMP_IF_FALSE:
 +        v = stack.Pop();
 +        if (flag == TR_STANDARD) {
-+          this->createGuardTrue(v, BH_JUMP_ARG, entry.pc, stack);
++          this->createGuardTrue(v, BH_JUMP_ARG, entry.pc, stack, locals);
 +        }
 +        continue;
 +
 +      case FOR_ITER:
 +        v = stack.Top();
 +        x = this->createForIter(v);
-+        this->createGuardNull(x, BH_FOR_ITER_EXC, entry.pc, stack);
++        this->createGuardNull(x, BH_FOR_ITER_EXC, entry.pc, stack, locals);
 +        stack.Push(x);
 +        continue;
 +
 +      case STORE_FAST:
++        w = locals.Load(entry.oparg);
++        this->createXDecRef(w);
++
 +        v = stack.Pop();
 +        if (entry.arg0 != NULL)
 +          v->registerStaticType(entry.arg0);
-+        this->createStoreFast(v, entry.oparg);
++        locals.Store(entry.oparg, v);
 +        continue;
 +
 +      case LOAD_FAST:
-+        v = this->createLoadFast(entry.oparg);
-+        this->createGuardNull(v, BH_LOAD_FAST_FAIL, entry.pc, stack);
++        v = locals.Load(entry.oparg);
++        this->createGuardNull(v, BH_LOAD_FAST_FAIL, entry.pc, stack, locals);
 +        this->createIncRef(v);
 +        stack.Push(v);
 +        continue;
 +      case INPLACE_MULTIPLY:
 +      case INPLACE_FLOOR_DIVIDE:
 +      case INPLACE_TRUE_DIVIDE:
-+        this->createBinaryOp(true, stack, entry);
++        this->createBinaryOp(true, stack, locals, entry);
 +        break;
 +
 +      case BINARY_ADD:
 +      case BINARY_MULTIPLY:
 +      case BINARY_FLOOR_DIVIDE:
 +      case BINARY_TRUE_DIVIDE:
-+        this->createBinaryOp(false, stack, entry);
++        this->createBinaryOp(false, stack, locals, entry);
 +        break;
 +
 +      case COMPARE_OP:
 +        if (entry.arg1 != NULL)
 +          w->registerStaticType(entry.arg1);
 +        x = this->createCompareOp(v, w, entry.oparg);
-+        this->createGuardNull(x, BH_BINARY_OP_FAIL, entry.pc, stack);
++        this->createGuardNull(x, BH_BINARY_OP_FAIL, entry.pc, stack, locals);
 +        stack.Push(x);
 +        continue;
 +
 +        assert(false);
 +      }
 +    }
-+    this->createTraceEnd(stack);
++    this->createTraceEnd(stack, locals);
++    loaded_locals_ = locals.loaded_locals;
 +    loaded_stack_ = stack.loaded_stack;
 +    stack_counter_ = stack.stack_counter;
 +  }
 +  }
 +
 +  void
++  JITLocalValue::emit(JITFunctionBuilder &fbuilder)
++  {
++    JITFunctionState::BuilderT &builder = fbuilder.builder();
++    llvm::Value *localptr = fbuilder.LocalPtr(num_);
++    value_ = builder.CreateLoad(localptr);
++    this->getStaticType()->init(value_);
++  }
++
++  void
 +  JITForIter::emit(JITFunctionBuilder &fbuilder)
 +  {
 +    JITFunctionState::BuilderT &builder = fbuilder.builder();
 +  }
 +
 +  void
-+  JITStoreFast::emit(JITFunctionBuilder &fbuilder)
-+  {
-+    JITFunctionState::BuilderT &builder = fbuilder.builder();
-+    llvm::Value *v = arg_->getValue(fbuilder);
-+    llvm::Value *localptr = fbuilder.LocalPtr(oparg_);
-+    llvm::Value *old = builder.CreateLoad(localptr);
-+    builder.CreateStore(v, localptr);
-+
-+    // This is correct. There is no static type for old.
-+    fbuilder.getState()->XDecRef(old);
-+  }
-+
-+  void
-+  JITLoadFast::emit(JITFunctionBuilder &fbuilder)
-+  {
-+    JITFunctionState::BuilderT &builder = fbuilder.builder();
-+    llvm::Value *localptr = fbuilder.LocalPtr(oparg_);
-+    value_ = builder.CreateLoad(localptr);
-+    this->getStaticType()->init(value_);
-+  }
-+
-+  void
 +  JITLoadConst::emit(JITFunctionBuilder &fbuilder)
 +  {
 +    JITFunctionState::BuilderT &builder = fbuilder.builder();
 +    builder.CreateCondBr(check, trace_cont, bail);
 +
 +    builder.SetInsertPoint(bail);
++    for (size_t i = 0; i < locals_.size(); ++i) {
++      if (locals_[i] != NULL)
++        fbuilder.StoreLocal(i, locals_[i]->getStaticType()->getValue(fbuilder, false, false));
++    }
 +    for (size_t i = 0; i < stack_.size(); ++i) {
 +      fbuilder.Push(stack_[i]->getStaticType()->getValue(fbuilder, false, false));
 +    }
 +                         bail, trace_cont);
 +
 +    builder.SetInsertPoint(bail);
-+
++    for (size_t i = 0; i < locals_.size(); ++i) {
++      if (locals_[i] != NULL)
++        fbuilder.StoreLocal(i, locals_[i]->getStaticType()->getValue(fbuilder, false, false));
++    }
 +    for (size_t i = 0; i < stack_.size(); ++i) {
 +      fbuilder.Push(stack_[i]->getStaticType()->getValue(fbuilder, false, false));
 +    }
 +    builder.SetInsertPoint(bail_cont);
 +
 +    arg_->DecRef(fbuilder);
++    for (size_t i = 0; i < locals_.size(); ++i) {
++      if (locals_[i] != NULL)
++        fbuilder.StoreLocal(i, locals_[i]->getStaticType()->getValue(fbuilder, false, false));
++    }
 +    for (size_t i = 0; i < stack_.size(); ++i) {
 +      fbuilder.Push(stack_[i]->getStaticType()->getValue(fbuilder, false, false));
 +    }
 +  JITTraceEnd::emit(JITFunctionBuilder &fbuilder)
 +  {
 +    // This is in the trace
++    for (size_t i = 0; i < locals_.size(); ++i) {
++      if (locals_[i] != NULL)
++        fbuilder.StoreLocal(i, locals_[i]->getValue(fbuilder));
++    }
++
 +    for (size_t i = 0; i < stack_.size(); ++i) {
 +      fbuilder.Push(stack_[i]->getValue(fbuilder));
 +    }
 +    arg_->IncRef(fbuilder);
 +  }
 +
++  void
++  JITXDecRef::emit(JITFunctionBuilder &fbuilder)
++  {
++    arg_->XDecRef(fbuilder);
++  }
++
 +  const int JITOpcodeInfo::ID = 0;
 +  const int JITStackValue::ID = 0;
++  const int JITLocalValue::ID = 0;
 +  const int JITForIter::ID = 0;
-+  const int JITStoreFast::ID = 0;
-+  const int JITLoadFast::ID = 0;
 +  const int JITLoadConst::ID = 0;
 +  const int JITCompareOp::ID = 0;
 +  const int JITIncRef::ID = 0;
++  const int JITXDecRef::ID = 0;
 +  const int JITGuardNull::ID = 0;
 +  const int JITGuardType::ID = 0;
 +  const int JITGuardTrue::ID = 0;
 new file mode 100644
 --- /dev/null
 +++ b/JIT/jit_tracer.h
-@@ -0,0 +1,735 @@
+@@ -0,0 +1,730 @@
 +// -*- C++ -*-
 +#ifndef PYTHON_JIT_TRACER_H_
 +#define PYTHON_JIT_TRACER_H_
 +
 +  class JITOpcodeInfo : public JITOpcode {
 +  public:
-+    JITOpcodeInfo(int lasti, std::vector<JITValue *> &stack)
-+      : lasti_(lasti), stack_(stack)
++    JITOpcodeInfo(int lasti, std::vector<JITValue *> &stack, std::vector<JITValue *> &locals)
++      : lasti_(lasti), stack_(stack), locals_(locals)
 +    {
 +    }
 +
 +  private:
 +    int lasti_;
 +    std::vector<JITValue *> stack_;
++    std::vector<JITValue *> locals_;
 +  };
 +
-+  class JITStackValue : public JITValue {
++
++  class JITLocalValue : public JITValue {
 +  public:
-+    JITStackValue()
-+      : value_(NULL), type_(new JITOpaqueObject)
++    JITLocalValue(int num)
++      : value_(NULL), type_(new JITOpaqueObject),
++        num_(num)
 +    {}
-+    ~JITStackValue()
++    ~JITLocalValue()
 +    {
 +      delete type_;
 +    }
 +
 +    virtual void dump()
 +    {
-+      std::cout << "Stack Value\n";
-+    }
-+    virtual void emit(JITFunctionBuilder &builder);
++      std::cout << "Local Value " << num_ << "\n";
++    }
++    virtual void emit(JITFunctionBuilder &fbuilder);
 +    virtual llvm::Value *getValue(JITFunctionBuilder &fbuilder) const
 +    {
 +      return getStaticType()->getValue(fbuilder);
 +    }
-+
 +    virtual llvm::Value *getNakedValue(JITFunctionBuilder &fbuilder) const
 +    {
 +      return getStaticType()->getNakedValue(fbuilder);
 +  private:
 +    llvm::Value *value_;
 +    JITType *type_;
++    int num_;
 +  };
 +
-+  class JITForIter : public JITValue {
++  class JITStackValue : public JITValue {
 +  public:
-+    JITForIter(JITValue *arg) : arg_(arg), value_(NULL), type_(new JITOpaqueObject)
-+    {
-+    }
-+
-+    ~JITForIter()
++    JITStackValue()
++      : value_(NULL), type_(new JITOpaqueObject)
++    {}
++    ~JITStackValue()
 +    {
 +      delete type_;
 +    }
 +
 +    virtual void dump()
 +    {
-+      std::cout << "FOR_ITER\n";
-+    }
-+
++      std::cout << "Stack Value\n";
++    }
 +    virtual void emit(JITFunctionBuilder &builder);
-+
 +    virtual llvm::Value *getValue(JITFunctionBuilder &fbuilder) const
 +    {
 +      return getStaticType()->getValue(fbuilder);
 +    static const int ID;
 +
 +  private:
-+    JITValue *arg_;
 +    llvm::Value *value_;
 +    JITType *type_;
 +  };
 +
-+  class JITStoreFast : public JITOpcode {
++  class JITForIter : public JITValue {
 +  public:
-+    JITStoreFast(JITValue *arg, int oparg)
-+      : arg_(arg), oparg_(oparg)
++    JITForIter(JITValue *arg) : arg_(arg), value_(NULL), type_(new JITOpaqueObject)
 +    {
 +    }
 +
++    ~JITForIter()
++    {
++      delete type_;
++    }
++
 +    virtual void dump()
 +    {
-+      std::cout << "STORE_FAST\n";
-+    }
-+
-+    virtual void emit(JITFunctionBuilder &builder);
-+
-+    virtual uintptr_t type()
-+    {
-+      return reinterpret_cast<uintptr_t>(&ID);
-+    }
-+    static const int ID;
-+
-+  private:
-+    JITValue *arg_;
-+    int oparg_;
-+  };
-+
-+  class JITLoadFast : public JITValue {
-+  public:
-+    JITLoadFast(int oparg)
-+      : oparg_(oparg), value_(NULL), type_(new JITOpaqueObject)
-+    {}
-+
-+    ~JITLoadFast()
-+    {
-+      delete type_;
-+    }
-+
-+    virtual void dump()
-+    {
-+      std::cout << "LOAD_FAST\n";
++      std::cout << "FOR_ITER\n";
 +    }
 +
 +    virtual void emit(JITFunctionBuilder &builder);
 +    static const int ID;
 +
 +  private:
-+    int oparg_;
++    JITValue *arg_;
 +    llvm::Value *value_;
 +    JITType *type_;
 +  };
 +    JITValue *arg_;
 +  };
 +
++  class JITXDecRef : public JITOpcode {
++  public:
++    JITXDecRef(JITValue *arg) : arg_(arg)
++    {}
++
++    virtual void dump()
++    {
++      std::cout << "XDECREF\n";
++    }
++
++    virtual void emit(JITFunctionBuilder &builder);
++
++    virtual uintptr_t type()
++    {
++      return reinterpret_cast<uintptr_t>(&ID);
++    }
++    static const int ID;
++
++  private:
++    JITValue *arg_;
++  };
++
 +  class JITGuardType : public JITOpcode {
 +  public:
 +    JITGuardType(JITValue *arg, BLACKHOLE reason, int lasti,
-+                 std::vector<JITValue *> &stack)
++                 std::vector<JITValue *> &stack,
++                 std::vector<JITValue *> &locals)
 +      : arg_(arg), reason_(reason), lasti_(lasti),
-+        stack_(stack)
++        stack_(stack), locals_(locals)
 +    {}
 +
 +    virtual void dump()
 +    BLACKHOLE reason_;
 +    int lasti_;
 +    std::vector<JITValue *> stack_;
++    std::vector<JITValue *> locals_;
 +  };
 +
 +  class JITGuardNull : public JITOpcode {
 +  public:
 +    JITGuardNull(JITValue *arg, BLACKHOLE reason, int lasti,
-+                 std::vector<JITValue *> &stack)
++                 std::vector<JITValue *> &stack, std::vector<JITValue *> &locals)
 +      : arg_(arg), reason_(reason), lasti_(lasti),
-+        stack_(stack)
++        stack_(stack), locals_(locals)
 +    {}
 +
 +    virtual void dump()
 +    BLACKHOLE reason_;
 +    int lasti_;
 +    std::vector<JITValue *> stack_;
++    std::vector<JITValue *> locals_;
 +  };
 +
 +  class JITGuardTrue : public JITOpcode {
 +  public:
 +    JITGuardTrue(JITValue *arg, BLACKHOLE reason, int lasti,
-+                 std::vector<JITValue *> &stack)
++                 std::vector<JITValue *> &stack, std::vector<JITValue *> &locals)
 +      : arg_(arg), reason_(reason), lasti_(lasti),
-+        stack_(stack)
++        stack_(stack), locals_(locals)
 +    {}
 +
 +    virtual void dump()
 +    BLACKHOLE reason_;
 +    int lasti_;
 +    std::vector<JITValue *> stack_;
++    std::vector<JITValue *> locals_;
 +  };
 +
 +  class JITTraceEnd : public JITOpcode {
 +  public:
-+    JITTraceEnd(std::vector<JITValue *> &stack)
-+      : stack_(stack)
++    JITTraceEnd(std::vector<JITValue *> &stack, std::vector<JITValue *> &locals)
++      : stack_(stack), locals_(locals)
 +    {
 +    }
 +
 +
 +  private:
 +    std::vector<JITValue *> stack_;
++    std::vector<JITValue *> locals_;
 +  };
 +
 +  struct JITTraceAnalysis {
 +
 +    void emit(JITFunctionState &state);
 +
-+    void createBinaryOp(bool inplace, JITStack &stack, TraceEntry &entry);
-+
-+    JITOpcodeInfo *createOpcodeInfo(int lasti, JITStack &stack)
++    void createBinaryOp(bool inplace, JITStack &stack, JITLocals &locals, TraceEntry &entry);
++
++    JITOpcodeInfo *createOpcodeInfo(int lasti, JITStack &stack, JITLocals &locals)
 +    {
-+      JITOpcodeInfo *r = new JITOpcodeInfo(lasti, stack.stack);
++      JITOpcodeInfo *r = new JITOpcodeInfo(lasti, stack.stack, locals.locals);
 +      trace_.push_back(r);
 +      return r;
 +    }
 +      return r;
 +    }
 +
-+    JITGuardNull *createGuardNull(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack)
++    JITGuardNull *createGuardNull(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack, JITLocals &locals)
 +    {
-+      JITGuardNull *r = new JITGuardNull(v, reason, lasti, stack.stack);
++      JITGuardNull *r = new JITGuardNull(v, reason, lasti, stack.stack, locals.locals);
 +      trace_.push_back(r);
 +      return r;
 +    }
 +
-+    JITGuardTrue *createGuardTrue(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack)
++    JITGuardTrue *createGuardTrue(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack, JITLocals &locals)
 +    {
-+      JITGuardTrue *r = new JITGuardTrue(v, reason, lasti, stack.stack);
++      JITGuardTrue *r = new JITGuardTrue(v, reason, lasti, stack.stack, locals.locals);
 +      trace_.push_back(r);
 +      return r;
 +    }
 +
-+    JITGuardType *createGuardType(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack)
++    JITGuardType *createGuardType(JITValue *v, BLACKHOLE reason, int lasti, JITStack &stack, JITLocals &locals)
 +    {
-+      JITGuardType *r = new JITGuardType(v, reason, lasti, stack.stack);
-+      trace_.push_back(r);
-+      return r;
-+    }
-+
-+    JITStoreFast *createStoreFast(JITValue *v, int oparg)
-+    {
-+      JITStoreFast *r = new JITStoreFast(v, oparg);
-+      trace_.push_back(r);
-+      return r;
-+    }
-+
-+    JITLoadFast *createLoadFast(int oparg)
-+    {
-+      JITLoadFast *r = new JITLoadFast(oparg);
++      JITGuardType *r = new JITGuardType(v, reason, lasti, stack.stack, locals.locals);
 +      trace_.push_back(r);
 +      return r;
 +    }
 +      return r;
 +    }
 +
-+    JITTraceEnd *createTraceEnd(JITStack &stack)
++    JITXDecRef *createXDecRef(JITValue *v)
 +    {
-+      JITTraceEnd *r = new JITTraceEnd(stack.stack);
++      JITXDecRef *r = new JITXDecRef(v);
 +      trace_.push_back(r);
 +      return r;
 +    }
 +
++    JITTraceEnd *createTraceEnd(JITStack &stack, JITLocals &locals)
++    {
++      JITTraceEnd *r = new JITTraceEnd(stack.stack, locals.locals);
++      trace_.push_back(r);
++      return r;
++    }
++
 +  private:
++    std::vector<JITLocalValue *> loaded_locals_;
 +    std::vector<JITStackValue *> loaded_stack_;
 +    int stack_counter_;
 +
  AC_MSG_CHECKING(for build directories)
  for dir in $SRCDIRS; do
      if test ! -d $dir; then
+diff --git a/dotests.py b/dotests.py
+new file mode 100644
+--- /dev/null
++++ b/dotests.py
+@@ -0,0 +1,188 @@
++import unittest
++import _jit
++
++
++class Utilities:
++    def setUp(self):
++        self.seq_float = [float(i) for i in range(20)]
++        self.seq_int = list(range(20))
++
++    def analyse(self, func, seq):
++        a = func(seq)
++        _jit.activate()
++        b = func(seq)
++        _jit.deactivate()
++        self.assertTrue(func.__code__._has_jit())
++        c = func(seq)
++
++        self.assertAlmostEqual(a, b)
++        self.assertAlmostEqual(a, c)
++
++
++class TestJIT(Utilities, unittest.TestCase):
++    def testMultiTracesPerCode(self):
++        def doit():
++            a = 0
++            for i in range(10):
++                a += i
++            for i in range(10):
++                a += i
++            return a
++
++        _jit.activate()
++        doit()
++        _jit.deactivate()
++        self.assertTrue(doit.__code__._has_jit())
++        ftxt = doit.__code__._dump_jit()
++        self.assertEqual(ftxt.count("define i32 @"), 2)
++
++    def testLocalsSimple(self):
++        def doit(seq):
++            acc = 0.0
++            for i in seq:
++                i += i
++                a = i + 2.0
++                acc += a
++
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++
++
++class TestCompares(Utilities, unittest.TestCase):
++    def testComparesSimple(self):
++        def doit(seq):
++            acc = 0.0
++            for i in seq:
++                acc += i < 1.0
++                acc += i > 1.0
++                acc += i is 1.0
++                acc += i is not 1.0
++
++        self.analyse(doit, self.seq_float)
++
++
++class TestBinops(Utilities, unittest.TestCase):
++    def testBinOpsFloat(self):
++        def doit(seq):
++            acc = 0.0
++            for i in seq:
++                acc -= i * 2.0
++            return acc
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++        def doit(seq):
++            acc = 1.0
++            for i in seq:
++                acc *= i + 1.0
++            return acc
++
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++        def doit(seq):
++            acc = 1.0
++            for i in seq:
++                acc += i * i + 1.0 / 2.0
++            return acc
++
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++    def testFloatFloorDiv(self):
++        def doit(seq):
++            acc = 1.0
++            for i in seq:
++                acc += 1.0 // (i + 1.0)
++            return acc
++        self.analyse(doit, self.seq_float)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++        # Test bail on zero divisor
++        def doit(seq):
++            acc = 0.0
++            for i in seq:
++                acc += 1.0 // i
++            return acc
++        self.analyse(doit, [1.0, 1.0, 1.0])
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++        try:
++            doit([1.0, 0.0])
++        except ZeroDivisionError:
++            pass
++        else:
++            self.fail("Compiled function did not raise ZeroDivisionError")
++
++    def testFloatTrueDiv(self):
++        def doit(seq):
++            acc = 1.0
++            for i in seq:
++                acc += 1.0 / (i + 1.0)
++            return acc
++        self.analyse(doit, self.seq_float)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++        self.analyse(doit, self.seq_int)
++
++        # Test bail on zero divisor
++        def doit(seq):
++            acc = 0.0
++            for i in seq:
++                acc += 1.0 / i
++            return acc
++        self.analyse(doit, [1.0, 1.0, 1.0])
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" not in ftxt)
++
++        try:
++            doit([1.0, 0.0])
++        except ZeroDivisionError:
++            pass
++        else:
++            self.fail("Compiled function did not raise ZeroDivisionError")
++
++    def testBinOpsInt(self):
++        def doit(seq):
++            acc = 0
++            for i in seq:
++                acc -= i + 2
++            return acc
++
++        self.analyse(doit, self.seq_int)
++        self.analyse(doit, self.seq_float)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" in ftxt)
++
++        def doit(seq):
++            acc = 1
++            for i in seq:
++                acc //= i + 1
++            return acc
++
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" in ftxt)
++
++        def doit(seq):
++            acc = 1
++            for i in seq:
++                acc /= i + 1
++            return acc
++
++        self.analyse(doit, self.seq_float)
++        self.analyse(doit, self.seq_int)
++        ftxt = doit.__code__._dump_jit()
++        self.assertTrue("PyNumber" in ftxt)
++
++if __name__ == "__main__":
++    unittest.main()
 diff --git a/pyconfig.h.in b/pyconfig.h.in
 --- a/pyconfig.h.in
 +++ b/pyconfig.h.in