<commit_before># -*- coding: utf-8 -*-

"""
Use blaze.bkernel to assemble ckernels for evaluation.
"""

from __future__ import print_function, division, absolute_import
from pykit.ir import interp

import blaze
from blaze.bkernel import BlazeFuncDeprecated
from blaze.datashape.util import to_numba

#------------------------------------------------------------------------
# Interpreter
#------------------------------------------------------------------------

def jit_interp(func, env=None, exc_model=None, args=()):
    env = env or {}
    env.setdefault('interp.handlers', {}).update(handlers)
    return interp.run(func, env, exc_model, args=args)


def compile(func, env):
    # NOTE: A problem of using a DataDescriptor as part of KernelTree is that
    #       we can now only compile kernels when we have actual data. This is
    #       a problem for offline compilation strategies.
    return func

def run(func, args, **kwds):
    deferred_array = jit_interp(func, args=args)
    result = blaze.eval(deferred_array, **kwds)
    return result

#------------------------------------------------------------------------
# Handlers
#------------------------------------------------------------------------

def op_kernel(interp, funcname, *args):
    op = interp.op
    overload = op.metadata['overload']
    func = overload.func

    blaze_func = make_blazefunc(func)
    return blaze_func(*args)

def op_convert(interp, arg):
    op = interp.op
    dtype = op.type.measure
    blaze_func = make_blazefunc(converter(dtype))
    return blaze_func(arg)

def op_ret(interp, arg):
    interp.pc = -1
    return arg

handlers = {
    'kernel':   op_kernel,
    'convert':  op_convert,
    'ret':      op_ret,
}

#------------------------------------------------------------------------
# Utils
#------------------------------------------------------------------------

def make_blazefunc(f):
    return BlazeFuncDeprecated(f.__name__, template=f)

def converter(blaze_type):
    """
    Generate an element-wise conversion function that numba can jit-compile.
    """
    T = to_numba(blaze_type)
    def convert(value):
        return T(value)
    return convert<commit_msg>Work on generating ckernels from jit<commit_after># -*- coding: utf-8 -*-

"""
Use blaze.bkernel to assemble ckernels for evaluation.
"""

from __future__ import print_function, division, absolute_import
from pykit.ir import interp, visit

import blaze
from blaze.bkernel import BlazeFuncDeprecated
from blaze.bkernel.blaze_kernels import frompyfunc, fromctypes, BlazeElementKernel
from blaze.bkernel.kernel_tree import Argument, KernelTree
from blaze.datashape.util import to_numba

#------------------------------------------------------------------------
# Interpreter
#------------------------------------------------------------------------

def compile(func, env):
    # NOTE: A problem of using a DataDescriptor as part of KernelTree is that
    #       we can now only compile kernels when we have actual data. This is
    #       a problem for offline compilation strategies.
    jitted = jitter(func, env)
    condense_jitted_kernels(func, env, jitted)
    return func

def run(func, args, **kwds):
    deferred_array = jit_interp(func, args=args)
    result = blaze.eval(deferred_array, **kwds)
    return result

#------------------------------------------------------------------------
# Jit kernels
#------------------------------------------------------------------------

class Jitter(object):
    """Jit kernels"""

    def __init__(self, func):
        self.func = func
        self.jitted = {}

    def op_kernel(self, op):
        function = op.metadata['kernel']
        overload = op.metadata['overload']
        impl = construct_blaze_kernel(function, overload)
        if impl is not None:
            self.jitted[op] = impl

    def op_convert(self, op):
        # TODO: Use a heuristic to see whether we need to handle this, or
        #       someone else does
        [arg] = op.args

        dtype = op.type.measure
        blaze_func = make_blazefunc(converter(dtype))
        self.jitted[op] = blaze_func(arg)


def construct_blaze_kernel(function, overload):
    """
    Parameters
    ==========
    function: blaze.kernel.Kernel
    overload: blaze.overloading.Overload
    """
    func = overload.func
    polysig = overload.sig
    monosig = overload.resolved_sig

    numba_impls = function.find_impls(func, polysig, 'numba')
    llvm_impls = function.find_impls(func, polysig, 'llvm')

    if numba_impls:
        [impl] = numba_impls
        argtypes = [to_numba(a.measure) for a in monosig.argtypes]
        restype = to_numba(monosig.restype.measure)
        return frompyfunc(impl, (argtypes, restype), monosig.argtypes)
    elif llvm_impls:
        [impl] = llvm_impls
        return BlazeElementKernel(impl, monosig.argtypes)
    else:
        return None


def jitter(func, env=None):
    v = Jitter(func)
    visit(v, func)
    return v.jitted

#------------------------------------------------------------------------
# Fuse jitted kernels
#------------------------------------------------------------------------

class JitFuser(object):

    def __init__(self, func, jitted):
        self.func = func
        self.jitted = jitted
        self.trees = {}

    def op_kernel(self, op):
        function = op.metadata['kernel']
        overload = op.metadata['overload']
        jitted = self.jitted.get(op)
        if not jitted:
            return op

        consumers = self.func.uses[op]
        if len(consumers) > 1:
            # We have multiple consumers
            pass

        # TODO: Check external references in metadata in order to determine
        #       whether this is a fusion boundary

        args = []
        for arg in op.args[1:]:
            if arg in self.trees:
                tree = self.trees[arg]
            else:
                kernel = self.jitted[arg]
                tree_arg = Argument(arg, kernel.kinds[i],
                                    self.ranks[i], kernel.argtypes[i])
                children.append(tree_arg)


def condense_jitted_kernels(func, env, jitted):
    visit(JitFuser(func, jitted), func)

#------------------------------------------------------------------------
# Utils
#------------------------------------------------------------------------

def make_blazefunc(f):
    return BlazeFuncDeprecated(f.__name__, template=f)

def converter(blaze_type):
    """
    Generate an element-wise conversion function that numba can jit-compile.
    """
    T = to_numba(blaze_type)
    def convert(value):
        return T(value)
    return convert

def make_ckernel(blaze_func):
    raise NotImplementedError
{'fraction_non_alphanumeric': 0.23442402510085164, 'fraction_numerical': 0.0007470491558344539, 'mean_word_length': 1.6626889419252189, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5182530', 'n_tokens_mistral': 1886, 'n_tokens_neox': 1846, 'n_words': 486}