numba/numba

Support jit decorator flags in cc.export (e.g. fastmath)

Open

#7,179 opened on Jul 3, 2021

View on GitHub
 (2 comments) (0 reactions) (0 assignees)Python (1,132 forks)batch import
AOTfeature_requestgood first issue

Repository metrics

Stars
 (9,177 stars)
PR merge metrics
 (Avg merge 14d 16h) (30 merged PRs in 30d)

Description

From #5664:

Adding the possibility of using @njit flags in @cc.export would be great! Is it possible?

and:

I have managed to add the 'fastmath' flag to the ExportEntry object:

class ExportEntry(object):
    """
    A simple record for exporting symbols.
    """

    def __init__(self, symbol, signature, function, fastmath=False):
        self.symbol = symbol
        self.signature = signature
        self.function = function
        self.fastmath = fastmath

Then, the _cull_exports methodof the _ModuleCompiler object shall me modified with the following code:

for entry in self.export_entries:
            flags_aux = flags.copy()
            if entry.fastmath:
                flags_aux.set('fastmath', value=cpu.FastMathOptions(True))

            cres = compile_extra(self.typing_context, self.context,
                                 entry.function,
                                 entry.signature.args,
                                 entry.signature.return_type, flags_aux,
                                 locals={}, library=library)

With this modifications, I am able to use the @cc.export decorator like this:

@cc.export('foo', 'float64(float64[::1], float64[::1])', fastmath=True)

And it is working for me right now. I guess this code may be improved, but this simple patch works, so I would expect that 'fastmath' and 'parallel' flags could be added to @cc.export in an "official" way.

The above could be worked up into a patch with tests added to complete the feature imp[lementation.

Contributor guide