<commit_before>from libc.stdint cimport intptr_t
from libcpp.vector cimport vector

from ..cpp.func.callable cimport callable as _callable
from ..cpp.type cimport type as _type
from ..cpp.types.callable_type cimport make_callable
from ..cpp.func.elwise cimport elwise as _elwise
from ..cpp.func.reduction cimport reduction as _reduction
from ..cpp.array cimport array as _array

from ..config cimport translate_exception
from .callable cimport callable, dynd_nd_callable_from_cpp, dynd_nd_callable_to_cpp
from ..ndt.type cimport type, as_numba_type, from_numba_type, as_cpp_type

cdef extern from 'dynd/functional.hpp' namespace 'dynd::nd::functional':
    _callable _dispatch 'dynd::nd::functional::dispatch'[T](_type, T) \
        except +translate_exception
    _callable _multidispatch 'dynd::nd::functional::multidispatch'[T](_type, T, T) \
        except +translate_exception

cdef extern from 'functional.hpp':
    _callable _apply 'apply'(_type, object) except +translate_exception

cdef extern from "kernels/apply_jit_kernel.hpp" namespace "pydynd::nd::functional":
    _callable _apply_jit "pydynd::nd::functional::apply_jit"(const _type &tp, intptr_t)

    cdef cppclass jit_dispatcher:
        jit_dispatcher(object, object (*)(object, intptr_t, const _type *))

def _import_numba():
    try:
        import numba
    except ImportError:
        return False

    return True

cdef public object _jit(object func, intptr_t nsrc, const _type *src_tp):
    from llvmlite import ir

    CharType = ir.IntType(8)
    CharPointerType = CharType.as_pointer()
    Int32Type = ir.IntType(32)
    Int64Type = ir.IntType(64)

    def add_single_ir(ir_module):
        single = ir.Function(ir_module, ir.FunctionType(ir.VoidType(),
            [CharPointerType, CharPointerType.as_pointer()]),
            name = 'single')

        bb_entry = single.append_basic_block('entry')
        ir_builder = ir.IRBuilder(bb_entry)

        src = []
        for i, ir_type in enumerate(wrapped_func_ir_tp.args[-nsrc::]):
            src.append(ir_builder.load(ir_builder.bitcast(ir_builder.load(ir_builder.gep(single.args[1],
                [ir.Constant(Int32Type, i)])), ir_type.as_pointer())))

        status, dst = target_context.call_conv.call_function(ir_builder, wrapped_func,
            fndesc.restype, fndesc.argtypes, src)

        ir_builder.store(dst,
            ir_builder.bitcast(single.args[0], wrapped_func_ir_tp.args[0]))
        ir_builder.ret_void()

        return single

    # This is the Numba signature
    signature = tuple(as_numba_type(src_tp[i]) for i in range(nsrc))

    # Compile the function with Numba
    func.compile(signature)
    compile_res = func.overloads[signature]

    # Check if there is a corresponding return type in DyND
    cdef _type dst_tp = from_numba_type(compile_res.signature.return_type)

    # The following generates the wrapper function using LLVM IR
    fndesc = compile_res.fndesc
    target_context = compile_res.target_context
    library = target_context.jit_codegen().create_library(name = 'library')

    ir_module = library.create_ir_module(name = 'module')
#    memcpy = ir_module.declare_intrinsic('llvm.memcpy',
 #       [CharPointerType, CharPointerType, Int32Type])

    wrapped_func_ir_tp = target_context.call_conv.get_function_type(fndesc.restype,
        fndesc.argtypes)
    wrapped_func = ir_module.get_or_insert_function(wrapped_func_ir_tp,
        name = fndesc.llvm_func_name)

    single = add_single_ir(ir_module)

    library.add_ir_module(ir_module)
    library.finalize()

    cdef vector[_type] src_tp_copy
    for i in range(nsrc):
        src_tp_copy.push_back(src_tp[i])

    return dynd_nd_callable_from_cpp(_apply_jit(make_callable(dst_tp, src_tp_copy),
            library.get_pointer_to_function('single')))

def apply(func = None, jit = _import_numba(), *args, **kwds):
    from .. import ndt
    def make(type tp, func):
        if jit:<commit_msg>Finish fix for interop with new numba interface.<commit_after>from libc.stdint cimport intptr_t
from libcpp.vector cimport vector

from ..cpp.func.callable cimport callable as _callable
from ..cpp.type cimport type as _type
from ..cpp.types.callable_type cimport make_callable
from ..cpp.func.elwise cimport elwise as _elwise
from ..cpp.func.reduction cimport reduction as _reduction
from ..cpp.array cimport array as _array

from ..config cimport translate_exception
from .callable cimport callable, dynd_nd_callable_from_cpp, dynd_nd_callable_to_cpp
from ..ndt.type cimport type, as_numba_type, from_numba_type, as_cpp_type

cdef extern from 'dynd/functional.hpp' namespace 'dynd::nd::functional':
    _callable _dispatch 'dynd::nd::functional::dispatch'[T](_type, T) \
        except +translate_exception
    _callable _multidispatch 'dynd::nd::functional::multidispatch'[T](_type, T, T) \
        except +translate_exception

cdef extern from 'functional.hpp':
    _callable _apply 'apply'(_type, object) except +translate_exception

cdef extern from "kernels/apply_jit_kernel.hpp" namespace "pydynd::nd::functional":
    _callable _apply_jit "pydynd::nd::functional::apply_jit"(const _type &tp, intptr_t)

    cdef cppclass jit_dispatcher:
        jit_dispatcher(object, object (*)(object, intptr_t, const _type *))

def _import_numba():
    try:
        import numba
    except ImportError:
        return False

    return True

cdef public object _jit(object func, intptr_t nsrc, const _type *src_tp):
    from llvmlite import ir

    CharType = ir.IntType(8)
    CharPointerType = CharType.as_pointer()
    Int32Type = ir.IntType(32)
    Int64Type = ir.IntType(64)

    def add_single_ir(ir_module):
        single = ir.Function(ir_module, ir.FunctionType(ir.VoidType(),
            [CharPointerType, CharPointerType.as_pointer()]),
            name = 'single')

        bb_entry = single.append_basic_block('entry')
        ir_builder = ir.IRBuilder(bb_entry)

        src = []
        for i, ir_type in enumerate(wrapped_func_ir_tp.args[-nsrc::]):
            src.append(ir_builder.load(ir_builder.bitcast(ir_builder.load(ir_builder.gep(single.args[1],
                [ir.Constant(Int32Type, i)])), ir_type.as_pointer())))

        status, dst = target_context.call_conv.call_function(ir_builder, wrapped_func,
            fndesc.restype, fndesc.argtypes, src)

        ir_builder.store(dst,
            ir_builder.bitcast(single.args[0], wrapped_func_ir_tp.args[0]))
        ir_builder.ret_void()

        return single

    # This is the Numba signature
    signature = tuple(as_numba_type(src_tp[i]) for i in range(nsrc))

    # Compile the function with Numba
    func.compile(signature)
    compile_res = func.overloads[signature]

    # Check if there is a corresponding return type in DyND
    cdef _type dst_tp = from_numba_type(compile_res.signature.return_type)

    # The following generates the wrapper function using LLVM IR
    fndesc = compile_res.fndesc
    target_context = compile_res.target_context
    library = target_context.codegen().create_library(name = 'library')

    ir_module = library.create_ir_module(name = 'module')
#    memcpy = ir_module.declare_intrinsic('llvm.memcpy',
 #       [CharPointerType, CharPointerType, Int32Type])

    wrapped_func_ir_tp = target_context.call_conv.get_function_type(fndesc.restype,
        fndesc.argtypes)
    wrapped_func = ir_module.get_or_insert_function(wrapped_func_ir_tp,
        name = fndesc.llvm_func_name)

    single = add_single_ir(ir_module)

    library.add_ir_module(ir_module)
    library.finalize()

    cdef vector[_type] src_tp_copy
    for i in range(nsrc):
        src_tp_copy.push_back(src_tp[i])

    return dynd_nd_callable_from_cpp(_apply_jit(make_callable(dst_tp, src_tp_copy),
            library.get_pointer_to_function('single')))

def apply(func = None, jit = _import_numba(), *args, **kwds):
    from .. import ndt
    def make(type tp, func):
        if jit:
{'fraction_non_alphanumeric': 0.09753320683111955, 'fraction_numerical': 0.004048070841239722, 'mean_word_length': 3.841396203306797, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4588369', 'n_tokens_mistral': 2684, 'n_tokens_neox': 2622, 'n_words': 554}