Make Python group operations

Viewed 77
def f1(x):
    for i in range(1, 100):
        x *= 2
        x /= 3.14159
        x *= i**.25
    return x

def f2(x):
    for i in range(1, 100):
        x *= 2 / 3.14159 * i**.25
    return x

Both functions compute exactly the same, but f1 takes 3x longer to do so, even with @numba.njit. Can Python be made to recognize the equivalence in compilation, just like it optimizes in other ways seen with dis by e.g. throwing out unused assignments?

Note, I'm aware floating point arithmetic cares about order, so the two functions may output slightly differently, but if anything more separate edits to array values are less accurate, so it'd be a 2-in-1 optimization.


x = np.random.randn(10000, 1000)
%timeit f1(x.copy())        # 2.68 s ± 50.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit f2(x.copy())        # 894 ms ± 36.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit njit(f1)(x.copy())  # 2.59 s ± 65.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit njit(f2)(x.copy())  # 901 ms ± 41.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
2 Answers

Using numba.jit is probably the best optimization you will get for the moment for this kind of function. You may also want to try pypy and do some benchmark comparisons.

Although, I want to point out why the two functions are not equivalent and so you should not expect f1 to be reduced to f2.

The order of operation goes as follows for f1:

x1 = (x * 2)            # First binary operation
x2 = (x1 / 3.14159      # Second binary operation
x3 = x2 * (i ** 0.25)   # Third and fourth binary operation

# Order: Multiplication, division, exponent, multiplication

Which is not the same as for f2:

x *= ((2 / 3.14159) * (i ** 0.25))
#  ^     ^          ^     ^
#  |     |          |     |
#  4     1          3     2

# Order: Division, exponent, multiplication, multiplication

Since floating-point arithmetic is not associative, those may not yield the same result. For that reason, it would be wrong for a compiler or interpreter to do the optimization you expected unless it is meant to optimize floating-point precision.

I am not aware of a Python tool which is meant to do this specific kind of optimization.

Probably can't do it with the jit. I have tried fastmath and nogil kwarg specified in the api: https://numba.pydata.org/numba-doc/latest/reference/jit-compilation.html

f0 is still slightly slower than f1 after getting rid of overflow or denormal number. plot

from timeit import default_timer as timer
import numpy as np
import matplotlib.pyplot as plt
import numba as nb


def f0(x):
    for i in range(1, 1000):
        x *= 3.000001
        x /= 3
    return x


def f1(x):
    for i in range(1, 1000):
        x *= 3.000001 / 3
    return x


def timing(f, **kwarg):
    x = np.ones(1000, dtype=np.float32)
    times = []
    n_iter = list(range(100, 1000, 100))
    f2 = nb.njit(f, **kwarg)
    for i in n_iter:
        print(i)
        s = timer()
        for j in range(i):
            f2(x)
        e = timer()
        times.append(e - s)
    print(x)
    m, b = np.polyfit(n_iter, times, 1)
    return times, m, b, n_iter


def main():
    results = []
    for fastmath in [True, False]:
        for i, f in enumerate([f0, f1]):
            kwarg = {
                "fastmath": fastmath,
                "nogil": True
            }
            r1, m, b, n_iter = timing(f, **kwarg)
            label = "f%d with %s" % (i, kwarg)
            plt.plot(n_iter, r1, label=label)
            results.append((m, b, label))
    for m, b, kwarg in results:
        print(m * 1e5, b, kwarg)
    plt.legend(loc="upper left")
    plt.xlabel("n iterations")
    plt.ylabel("timing")
    plt.show()
    plt.close()


if __name__ == '__main__':
    main()

Related