Repository: numba/numba

Issue:
Not implemented opcode: EXTENDED_ARG
Trying to compile a code with an opcode longer than 64K instructions causes Numba to raise a `NotImplementedError`, the `EXTENDED_ARG` opcode is not supported. In our case, we need this to compile an auto generated scientific code.

A good explanation of what the EXTENDED_ARG opcode does, can be found here: https://mail.python.org/pipermail/python-dev/2008-March/078104.html

``` python
code = """from numba import jit, void
@jit(void(), nopython=True)
def run():
    val = 0
    while val < 100000:
"""
for i in range(10000):
    code += "        val += 1\n"
with open("huge.py", "w") as fp:
    fp.write(code)
import huge
```

The raised Exception:

```
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<...>/generate_huge.py in <module>()
      9 with open("huge.py", "w") as fp:
     10     fp.write(code)
---> 11 import huge

<...>/huge.py in <module>()
      1 from numba import jit, void
----> 2 @jit(void(), nopython=True)
      3 def run():
      4     val = 0
      5     while val < 100000:

/opt/local/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/site-packages/numba/decorators.pyc in wrapper(func)
    166                           targetoptions=targetoptions)
    167         for sig in sigs:
--> 168             disp.compile(sig)
    169         disp.disable_compile()
    170         return disp

/opt/local/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/site-packages/numba/dispatcher.pyc in compile(self, sig)
    320                                           self.py_func,
    321                                           args=args, return_type=return_type,
--> 322                                           flags=flags, locals=self.locals)
    323 
    324             # Check typing error if object mode is used

/opt/local/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/site-packages/numba/compiler.pyc in compile_extra(typingctx, targetctx, func, args, return_type, flags, locals, library)
    592     pipeline = Pipeline(typingctx, targetctx, library,
    593                         args, return_type, flags, locals)
--> 594     return pipeline.compile_extra(func)
    595 
    596 

/opt/local/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/site-packages/numba/compiler.pyc in compile_extra(self, func)
    313                 return self.stage_compile_interp_mode()
    314             else:
--> 315                 raise e
    316 
    317         return self.compile_bytecode(bc, func_attr=self.func_attr)

NotImplementedError: offset=6 opcode=0x91 opname=EXTENDED_ARG
```

I'm running Python 2.7, Numba 0.20.0 on OSX



Discussion:
The EXTENDED_ARG opcode itself is easy to support but you must be prepared for very long compilation times if you are JIT-compiling such a function. Right now Numba isn't prepared to deal with such long functions.

I'm aware of this, but we have a very long running simulation, so we gain anyway.


Resolving patch:
diff --git a/numba/bytecode.py b/numba/bytecode.py
--- a/numba/bytecode.py
+++ b/numba/bytecode.py
@@ -94,6 +94,7 @@ def _make_bytecode_table():
                     ('COMPARE_OP', 2),
                     ('DELETE_ATTR', 2),
                     ('DUP_TOP', 0),
+                    ('EXTENDED_ARG', 2),
                     ('FOR_ITER', 2),
                     ('GET_ITER', 0),
                     ('INPLACE_ADD', 0),
@@ -153,6 +154,7 @@ def _as_opcodes(seq):
 JABS_OPS = frozenset(dis.hasjabs)
 JUMP_OPS = JREL_OPS | JABS_OPS
 TERM_OPS = frozenset(_as_opcodes(['RETURN_VALUE', 'RAISE_VARARGS']))
+EXTENDED_ARG = dis.EXTENDED_ARG
 
 
 class ByteCodeInst(object):
@@ -225,7 +227,7 @@ def __init__(self, code):
     def __iter__(self):
         return self
 
-    def next(self):
+    def _fetch_opcode(self):
         offset, opcode = next(self.iter)
         try:
             info = BYTECODE_TABLE[opcode]
@@ -237,6 +239,14 @@ def next(self):
             arg = self.read_arg(info.argsize)
         else:
             arg = None
+        return offset, opcode, arg
+
+    def next(self):
+        offset, opcode, arg = self._fetch_opcode()
+        if opcode == EXTENDED_ARG:
+            hi_arg = arg
+            offset, opcode, lo_arg = self._fetch_opcode()
+            arg = (hi_arg << 16) + lo_arg
         return offset, ByteCodeInst(offset=offset, opcode=opcode, arg=arg)
 
     __next__ = next


Test patch:
diff --git a/numba/tests/support.py b/numba/tests/support.py
--- a/numba/tests/support.py
+++ b/numba/tests/support.py
@@ -312,6 +312,31 @@ def compile_function(name, code, globs):
     eval(co, globs, ns)
     return ns[name]
 
+def tweak_code(func, codestring=None, consts=None):
+    """
+    Tweak the code object of the given function by replacing its
+    *codestring* (a bytes object) and *consts* tuple, optionally.
+    """
+    co = func.__code__
+    tp = type(co)
+    if codestring is None:
+        codestring = co.co_code
+    if consts is None:
+        consts = co.co_consts
+    if sys.version_info >= (3,):
+        new_code = tp(co.co_argcount, co.co_kwonlyargcount, co.co_nlocals,
+                      co.co_stacksize, co.co_flags, codestring,
+                      consts, co.co_names, co.co_varnames,
+                      co.co_filename, co.co_name, co.co_firstlineno,
+                      co.co_lnotab)
+    else:
+        new_code = tp(co.co_argcount, co.co_nlocals,
+                      co.co_stacksize, co.co_flags, codestring,
+                      consts, co.co_names, co.co_varnames,
+                      co.co_filename, co.co_name, co.co_firstlineno,
+                      co.co_lnotab)
+    func.__code__ = new_code
+
 
 # From CPython
 
diff --git a/numba/tests/test_extended_arg.py b/numba/tests/test_extended_arg.py
new file mode 100644
--- /dev/null
+++ b/numba/tests/test_extended_arg.py
@@ -0,0 +1,42 @@
+from __future__ import print_function
+
+import numba.unittest_support as unittest
+
+import dis
+import struct
+import sys
+
+from numba import jit
+from .support import TestCase, tweak_code
+
+
+class TestExtendedArg(TestCase):
+    """
+    Test support for the EXTENDED_ARG opcode.
+    """
+
+    def get_extended_arg_load_const(self):
+        """
+        Get a function with a EXTENDED_ARG opcode before a LOAD_CONST opcode.
+        """
+        def f():
+            x = 5
+            return x
+
+        b = bytearray(f.__code__.co_code)
+        consts = f.__code__.co_consts
+        consts = consts + (None,) * 65535 + (42,)
+        b[:0] = struct.pack("<BH", dis.EXTENDED_ARG, 1)
+
+        tweak_code(f, codestring=bytes(b), consts=consts)
+        return f
+
+    def test_extended_arg_load_const(self):
+        pyfunc = self.get_extended_arg_load_const()
+        self.assertPreciseEqual(pyfunc(), 42)
+        cfunc = jit(nopython=True)(pyfunc)
+        self.assertPreciseEqual(cfunc(), 42)
+
+
+if __name__ == '__main__':
+    unittest.main()