Merge pull request #353 from tonybaloney/generator_assertions

Pyjion will infer that `range(n)` generates integers in iterator to improve unboxing
tonybaloney · Aug 5, 2021 · 4a43d19 · 4a43d19
2 parents bf35e82 + b535b6d
commit 4a43d19
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 22 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@
 
 * BINARY_MULTIPLY and BINARY_POWER will assume the resulting integer is a big integer (not unboxed)
 * Introduced two optimizations IntegerUnboxingMultiply and IntegerUnboxingPower which are applied at optimization level 2. Try level two if you work with integers, but at smaller values to see better performance.
+* Pyjion will infer that `range(n)` generates integers in iterator to improve unboxing
 * LOAD_BUILD_CLASS will infer a function type instead of Any (#42)
 * Instruction graphs will include the name of fast locals
 * Instruction graph const values are capped to 40 characters 

diff --git a/Tests/test_pgc.cpp b/Tests/test_pgc.cpp
@@ -330,12 +330,12 @@ TEST_CASE("Test PGC integer logic"){
                 "     x = 2\n"
                 "     z = y * y + x - y\n"
                 "     x *= z\n"
-                "  return\n"
+                "  return x\n"
                 );
         CHECK(t.pgcStatus() == PgcStatus::Uncompiled);
-        CHECK(t.returns() == "None");
+        CHECK(t.returns() == "19408");
         CHECK(t.pgcStatus() == PgcStatus::CompiledWithProbes);
-        CHECK(t.returns() == "None");
+        CHECK(t.returns() == "19408");
         CHECK(t.pgcStatus() == PgcStatus::Optimized);
     };
 }
diff --git a/src/pyjion/absint.cpp b/src/pyjion/absint.cpp
@@ -770,12 +770,10 @@ AbstractInterpreter::interpret(PyObject *builtins, PyObject *globals, PyjionCode
                     // about their deletion.
                     break;
                 case GET_ITER: {
-                    auto iteratorType = POP_VALUE();
-                    // TODO : Allow guarded/PGC sources to be optimized.
-                    auto source = AbstractValueWithSources(
-                            &Iterable,
-                            newSource(new IteratorSource(iteratorType.Value->needsGuard() ? AVK_Any: iteratorType.Value->kind(), curByte)));
-                    lastState.push(source);
+                    auto in = POP_VALUE();
+                    auto out = in.Value->iter(in.Sources);
+                    PUSH_INTERMEDIATE(out);
+                    break;
                 }
                     break;
                 case FOR_ITER: {
@@ -789,7 +787,8 @@ AbstractInterpreter::interpret(PyObject *builtins, PyObject *globals, PyjionCode
                     // When we compile this we don't actually leave the value on the stack,
                     // but the sequence of opcodes assumes that happens.  to keep our stack
                     // properly balanced we match what's really going on.
-                    PUSH_INTERMEDIATE(&Any);
+                    auto out = iterator.Value->next(iterator.Sources);
+                    PUSH_INTERMEDIATE(out);
                     break;
                 }
                 case POP_BLOCK: {

diff --git a/src/pyjion/absvalue.cpp b/src/pyjion/absvalue.cpp
@@ -53,6 +53,7 @@ MethodValue Method;
 CodeObjectValue CodeObject;
 EnumeratorValue Enumerator;
 RangeIteratorValue RangeIterator;
+RangeValue Range;
 MemoryViewValue MemoryView;
 ClassMethodValue ClassMethod;
 FilterValue Filter;
@@ -78,6 +79,14 @@ AbstractValue* AbstractValue::unary(AbstractSource* selfSources, int op) {
     return &Any;
 }
 
+AbstractValue* AbstractValue::iter(AbstractSource* selfSources) {
+    return &Any;
+}
+
+AbstractValue* AbstractValue::next(AbstractSource* selfSources) {
+    return &Any;
+}
+
 AbstractValue* AbstractValue::compare(AbstractSource* selfSources, int op, AbstractValueWithSources& other) {
     if (isKnownType(kind()) && isKnownType(other.Value->kind())) {
         // We know all of the known types don't have fancy rich comparison
@@ -1197,6 +1206,16 @@ const char *CodeObjectValue::describe() {
     return "codeobject";
 }
 
+// Iterators
+
+AbstractValue* RangeValue::iter(AbstractSource* selfSources) {
+    return &RangeIterator;
+}
+
+AbstractValue* RangeIteratorValue::next(AbstractSource* selfSources) {
+    return &Integer;
+}
+
 AbstractValueKind knownFunctionReturnType(AbstractValueWithSources source){
     // IS this a builtin?
     if (source.hasSource() && source.Sources->isBuiltin())
@@ -1260,6 +1279,8 @@ AbstractValue* avkToAbstractValue(AbstractValueKind kind){
             return &Method;
         case AVK_BigInteger:
             return &Integer;
+        case AVK_Range:
+            return &Range;
         case AVK_RangeIterator:
             return &RangeIterator;
         case AVK_MemoryView:
@@ -1363,6 +1384,9 @@ AbstractValueKind GetAbstractType(PyTypeObject* type, PyObject* value) {
     else if (type == &PyModule_Type) {
         return AVK_Module;
     }
+    else if (type == &PyRange_Type) {
+        return AVK_Range;
+    }
     else if (type == &PyRangeIter_Type) {
         return AVK_RangeIterator;
     }

diff --git a/src/pyjion/absvalue.h b/src/pyjion/absvalue.h
@@ -65,17 +65,18 @@ enum AbstractValueKind {
     AVK_Module = 21,
     AVK_Method = 22,
     AVK_BigInteger = 23,
-    AVK_RangeIterator = 24,
-    AVK_MemoryView = 25,
-    AVK_Classmethod = 26,
-    AVK_Filter = 27,
-    AVK_Property = 28,
-    AVK_Map = 29,
-    AVK_Baseobject = 30,
-    AVK_Reversed = 31,
-    AVK_Staticmethod = 32,
-    AVK_Super = 33,
-    AVK_Zip = 34,
+    AVK_Range = 24,
+    AVK_RangeIterator = 25,
+    AVK_MemoryView = 26,
+    AVK_Classmethod = 27,
+    AVK_Filter = 28,
+    AVK_Property = 29,
+    AVK_Map = 30,
+    AVK_Baseobject = 31,
+    AVK_Reversed = 32,
+    AVK_Staticmethod = 33,
+    AVK_Super = 34,
+    AVK_Zip = 35,
 };
 
 static bool isKnownType(AbstractValueKind kind) {
@@ -297,6 +298,8 @@ class AbstractValue {
     virtual AbstractValue* unary(AbstractSource* selfSources, int op);
     virtual AbstractValue* binary(AbstractSource* selfSources, int op, AbstractValueWithSources& other);
     virtual AbstractValue* compare(AbstractSource* selfSources, int op, AbstractValueWithSources& other);
+    virtual AbstractValue* iter(AbstractSource* selfSources);
+    virtual AbstractValue* next(AbstractSource* selfSources);
 
     virtual bool isAlwaysTrue() {
         return false;
@@ -566,13 +569,24 @@ class EnumeratorValue : public AbstractValue {
     const char* describe() override;
 };
 
+class RangeValue : public AbstractValue {
+    AbstractValueKind kind() override {
+        return AVK_Range;
+    }
+    const char* describe() override {
+        return "range";
+    }
+    AbstractValue* iter(AbstractSource* selfSources) override;
+};
+
 class RangeIteratorValue : public AbstractValue {
     AbstractValueKind kind() override {
         return AVK_RangeIterator;
     }
     const char* describe() override {
         return "range iterator";
     }
+    AbstractValue* next(AbstractSource* selfSources) override;
 };
 
 class MemoryViewValue : public AbstractValue {
@@ -740,6 +754,7 @@ extern MethodValue Method;
 extern CodeObjectValue CodeObject;
 extern EnumeratorValue Enumerator;
 extern RangeIteratorValue RangeIterator;
+extern RangeValue Range;
 extern MemoryViewValue MemoryView;
 extern ClassMethodValue ClassMethod;
 extern FilterValue Filter;

diff --git a/src/pyjion/knownmethods.h b/src/pyjion/knownmethods.h
@@ -54,7 +54,7 @@ unordered_map<const char*, AbstractValueKind> builtinReturnTypes = {
         {"pow",         AVK_Any},
         {"property",    AVK_Property},
         {"print",       AVK_None},
-        {"range",       AVK_RangeIterator},
+        {"range",       AVK_Range},
         {"repr",        AVK_String},
         {"reversed",    AVK_Reversed},
         {"round",       AVK_Any},