numba
numba copied to clipboard
Add ability to @jit to disable auto-vectorization
Two use cases:
- More significantly, sometimes auto-vectorization with SIMD makes a function slower. There are the environment variables to disable auto-vectorization, but that affects all code in the process. Even if there's only a single function, restarting the process for performance comparisons is annoying.
- For educational purposes, I am disabling it and comparing functions with and without SIMD (which is how I discovered the above).
A workaround
Here's what I'm doing now:
from contextlib import contextmanager
import sys
import os
@contextmanager
def disabled_simd():
"""
Temporarily disable SIMD when creating an `@njit`-decorated function.
This is a horrible hack, but Numba doesn't have any other way of doing this
at the time of writing.
Usage::
with disabled_simd() as njit_no_simd:
@njit_no_simd
def this_will_not_use_simd():
# ... your Numba code ...
# Make sure to compile it inside the with block by calling it with appropriate inputs!
_ = this_will_not_use_simd()
"""
def clear_numba():
"""Remove already imported modules for Numba and LLVMLite."""
for mod in list(sys.modules):
if mod.startswith("numba") or mod.startswith("llvmlite"):
del sys.modules[mod]
clear_numba()
os.environ["NUMBA_LOOP_VECTORIZE"] = "0"
try:
from numba import njit
yield njit
finally:
os.environ["NUMBA_LOOP_VECTORIZE"] = "1"
clear_numba()
Thanks for the request @itamarst. Just to be clear, are you looking to disable loop vectorization or to disable the issuing of SIMD instructions entirely (or something else)?
Mmm. Disable loop auto-vectorization I guess. And once SLP is re-enabled, have a flag disabling that. Basically a programmatic equivalent to NUMBA_LOOP_VECTORIZE.
Mmm. Disable loop auto-vectorization I guess. And once SLP is re-enabled, have a flag disabling that. Basically a programmatic equivalent to
NUMBA_LOOP_VECTORIZE.
Thanks for clarifying!