numba icon indicating copy to clipboard operation
numba copied to clipboard

Add ability to @jit to disable auto-vectorization

Open itamarst opened this issue 2 years ago • 3 comments

Two use cases:

  1. 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.
  2. 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()

itamarst avatar Jan 17 '24 01:01 itamarst

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)?

stuartarchibald avatar Jan 17 '24 15:01 stuartarchibald

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.

itamarst avatar Jan 17 '24 16:01 itamarst

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!

stuartarchibald avatar Jan 17 '24 16:01 stuartarchibald