<commit_before>from __future__ import print_function, division, absolute_import
from numba import ir
from numba.ir_utils import *
#from numba.annotations import type_annotations
from numba import types, config
from numba.typing import npydecl

import numpy
MAP_TYPES = [numpy.ufunc]

class ArrayAnalysis(object):
    """Analyzes Numpy array computations for properties such as shapes
    and equivalence classes.
    """

    def __init__(self, func_ir, type_annotation):
        """Constructor for the Rewrite class."""
        self.func_ir = func_ir
        self.type_annotation = type_annotation
        self.next_eq_class = 1
        # equivalence classes for each dimension of each array are saved
        # string to tuple of class numbers
        # example: {'A':[1,2]}
        #          {1:[n,a],2:[k,m,3]}
        self.array_shape_classes = {}
        # class zero especial and is size 1 for constants
        # and added broadcast dimensions
        # -1 class means unknown
        self.class_sizes = {0:[1]}
        # size variable to use for each array dimension
        self.array_size_vars = {}
        # keep a list of numpy Global variables to find numpy calls
        self.numpy_globals = []
        # calls that are essentially maps like DUFunc
        self.map_calls = []
        # keep numpy call variables with their call names
        self.numpy_calls = {}
        # keep attr calls to arrays like t=A.sum() as {t:('sum',A)}
        self.array_attr_calls = {}
        # keep tuple builds like {'t':[a,b],}
        self.tuple_table = {}

    def run(self):
        # TODO: ignoring CFG for now
        for (key, block) in self.func_ir.blocks.items():
            self._analyze_block(block)
        if config.DEBUG_ARRAY_OPT==1:
            print("classes: ", self.array_shape_classes)
            print("class sizes: ", self.class_sizes)
            print("numpy globals ", self.numpy_globals)
            print("numpy calls ", self.numpy_calls)
            print("array attr calls ", self.array_attr_calls)
            print("tuple table ", self.tuple_table)

    def _analyze_block(self, block):
        out_body = []
        for inst in block.body:
            # instructions can generate extra size calls to be appended
            # if an array doesn't have a size variable, it should be generated
            # when it is created
            generated_size_calls = self._analyze_inst(inst)
            out_body.append(inst)
            for node in generated_size_calls:
                out_body.append(node)
        block.body = out_body

    def _analyze_inst(self, inst):
        if isinstance(inst, ir.Assign):
            return self._analyze_assign(inst)
        return []

    def _analyze_assign(self, assign):

        #if isinstance(assign.value, ir.Arg) and self._isarray(assign.value.name):
        #    self._add_array_corr(assign.value.name)
        # lhs is always var?
        assert isinstance(assign.target, ir.Var)
        lhs = assign.target.name
        rhs = assign.value
        if isinstance(rhs, ir.Global):
            for T in MAP_TYPES:
                if isinstance(rhs.value, T):
                    self.map_calls.append(lhs)
            if rhs.value.__name__=='numpy':
                self.numpy_globals.append(lhs)
        if isinstance(rhs, ir.Expr) and rhs.op=='getattr':
            if rhs.value.name in self.numpy_globals:
                self.numpy_calls[lhs] = rhs.attr
            elif self._isarray(rhs.value.name):
                self.array_attr_calls[lhs] = (rhs.attr, rhs.value.name)
        if isinstance(rhs, ir.Expr) and rhs.op=='build_tuple':
            self.tuple_table[lhs] = rhs.items

        #rhs_class_out = self._analyze_rhs_classes(rhs)
        size_calls = []
        if self._isarray(lhs):
            rhs_corr = self._analyze_rhs_classes(rhs).copy()
            self.array_shape_classes[lhs] = rhs_corr
            self.array_size_vars[lhs] = [-1]*self._get_ndims(lhs)
            # make sure output lhs array has size variables for each dimension
            for (i,corr) in enumerate(rhs_corr):
                # if corr unknown or new
                if corr==-1 or corr not in self.class_sizes.keys():
                    # generate size call nodes for this dimension
                    nodes = self._gen_size_call(assign.target, i)
                    size_calls += nodes
                    assert isinstance(nodes[-1], ir.Assign)
                    size_var = nodes[-1].target
                    if corr!=-1:
                        self.class_sizes[corr] = [size_var]
                    self.array_size_vars[lhs][i] = size_var
                else:
                    # reuse a size variable from this correlation
                    # TODO: consider CFG?
                    self.array_size_vars[lhs][i] = self.class_sizes[corr][0]

        #print(self.array_shape_classes)
        return size_calls

    def _gen_size_call(self, var, i):
        out = []
        ndims = self._get_ndims(var.name)
        # attr call: A_sh_attr = getattr(A, shape)
        shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
        attr_var = ir.Var(var.scope, var.name+"_sh_attr"+str(i), var.loc)
        self.type_annotation.typemap[attr_var.name] = types.containers.UniTuple(INT_TYPE, ndims)
        attr_assign = ir.Assign(shape_attr_call, attr_var, var.loc)
        out.append(attr_assign)
        # const var for dim: $constA0 = Const(0)
        const_node = ir.Const(i, var.loc)
        const_var = ir.Var(var.scope, "$const"+var.name+str(i), var.loc)
        self.type_annotation.typemap[const_var.name] = INT_TYPE
        const_assign = ir.Assign(const_node, const_var, var.loc)
        out.append(const_assign)
        # get size: Asize0 = A_sh_attr[0]
        size_var = ir.Var(var.scope, var.name+"size"+str(i), var.loc)
        self.type_annotation.typemap[size_var.name] = INT_TYPE
        getitem_node = ir.Expr.static_getitem(attr_var, i, const_var, var.loc)
        self.type_annotation.calltypes[getitem_node] = None
        getitem_assign = ir.Assign(getitem_node, size_var, var.loc)
        out.append(getitem_assign)
        return out

    # lhs is array so rhs has to return array
    def _analyze_rhs_classes(self, node):
        if isinstance(node, ir.Arg):
            assert self._isarray(node.name)
            return self._add_array_corr(node.name)
        elif isinstance(node, ir.Var):
            return self.array_shape_classes[node.name].copy()
        elif isinstance(node, ir.Expr):
            if node.op=='unary' and node.fn in UNARY_MAP_OP:
                assert isinstance(node.value, ir.Var)
                in_var = node.value.name<commit_msg>unique names for array size vars<commit_after>from __future__ import print_function, division, absolute_import
from numba import ir
from numba.ir_utils import *
#from numba.annotations import type_annotations
from numba import types, config
from numba.typing import npydecl

import numpy
MAP_TYPES = [numpy.ufunc]

class ArrayAnalysis(object):
    """Analyzes Numpy array computations for properties such as shapes
    and equivalence classes.
    """

    def __init__(self, func_ir, type_annotation):
        """Constructor for the Rewrite class."""
        self.func_ir = func_ir
        self.type_annotation = type_annotation
        self.next_eq_class = 1
        # equivalence classes for each dimension of each array are saved
        # string to tuple of class numbers
        # example: {'A':[1,2]}
        #          {1:[n,a],2:[k,m,3]}
        self.array_shape_classes = {}
        # class zero especial and is size 1 for constants
        # and added broadcast dimensions
        # -1 class means unknown
        self.class_sizes = {0:[1]}
        # size variable to use for each array dimension
        self.array_size_vars = {}
        # keep a list of numpy Global variables to find numpy calls
        self.numpy_globals = []
        # calls that are essentially maps like DUFunc
        self.map_calls = []
        # keep numpy call variables with their call names
        self.numpy_calls = {}
        # keep attr calls to arrays like t=A.sum() as {t:('sum',A)}
        self.array_attr_calls = {}
        # keep tuple builds like {'t':[a,b],}
        self.tuple_table = {}

    def run(self):
        # TODO: ignoring CFG for now
        for (key, block) in self.func_ir.blocks.items():
            self._analyze_block(block)
        if config.DEBUG_ARRAY_OPT==1:
            print("classes: ", self.array_shape_classes)
            print("class sizes: ", self.class_sizes)
            print("numpy globals ", self.numpy_globals)
            print("numpy calls ", self.numpy_calls)
            print("array attr calls ", self.array_attr_calls)
            print("tuple table ", self.tuple_table)

    def _analyze_block(self, block):
        out_body = []
        for inst in block.body:
            # instructions can generate extra size calls to be appended
            # if an array doesn't have a size variable, it should be generated
            # when it is created
            generated_size_calls = self._analyze_inst(inst)
            out_body.append(inst)
            for node in generated_size_calls:
                out_body.append(node)
        block.body = out_body

    def _analyze_inst(self, inst):
        if isinstance(inst, ir.Assign):
            return self._analyze_assign(inst)
        return []

    def _analyze_assign(self, assign):

        #if isinstance(assign.value, ir.Arg) and self._isarray(assign.value.name):
        #    self._add_array_corr(assign.value.name)
        # lhs is always var?
        assert isinstance(assign.target, ir.Var)
        lhs = assign.target.name
        rhs = assign.value
        if isinstance(rhs, ir.Global):
            for T in MAP_TYPES:
                if isinstance(rhs.value, T):
                    self.map_calls.append(lhs)
            if rhs.value.__name__=='numpy':
                self.numpy_globals.append(lhs)
        if isinstance(rhs, ir.Expr) and rhs.op=='getattr':
            if rhs.value.name in self.numpy_globals:
                self.numpy_calls[lhs] = rhs.attr
            elif self._isarray(rhs.value.name):
                self.array_attr_calls[lhs] = (rhs.attr, rhs.value.name)
        if isinstance(rhs, ir.Expr) and rhs.op=='build_tuple':
            self.tuple_table[lhs] = rhs.items

        #rhs_class_out = self._analyze_rhs_classes(rhs)
        size_calls = []
        if self._isarray(lhs):
            rhs_corr = self._analyze_rhs_classes(rhs).copy()
            self.array_shape_classes[lhs] = rhs_corr
            self.array_size_vars[lhs] = [-1]*self._get_ndims(lhs)
            # make sure output lhs array has size variables for each dimension
            for (i,corr) in enumerate(rhs_corr):
                # if corr unknown or new
                if corr==-1 or corr not in self.class_sizes.keys():
                    # generate size call nodes for this dimension
                    nodes = self._gen_size_call(assign.target, i)
                    size_calls += nodes
                    assert isinstance(nodes[-1], ir.Assign)
                    size_var = nodes[-1].target
                    if corr!=-1:
                        self.class_sizes[corr] = [size_var]
                    self.array_size_vars[lhs][i] = size_var
                else:
                    # reuse a size variable from this correlation
                    # TODO: consider CFG?
                    self.array_size_vars[lhs][i] = self.class_sizes[corr][0]

        #print(self.array_shape_classes)
        return size_calls

    def _gen_size_call(self, var, i):
        out = []
        ndims = self._get_ndims(var.name)
        # attr call: A_sh_attr = getattr(A, shape)
        shape_attr_call = ir.Expr.getattr(var, "shape", var.loc)
        attr_var = ir.Var(var.scope, mk_unique_var(var.name+"_sh_attr"+str(i)), var.loc)
        self.type_annotation.typemap[attr_var.name] = types.containers.UniTuple(INT_TYPE, ndims)
        attr_assign = ir.Assign(shape_attr_call, attr_var, var.loc)
        out.append(attr_assign)
        # const var for dim: $constA0 = Const(0)
        const_node = ir.Const(i, var.loc)
        const_var = ir.Var(var.scope, mk_unique_var("$const"+var.name+str(i)), var.loc)
        self.type_annotation.typemap[const_var.name] = INT_TYPE
        const_assign = ir.Assign(const_node, const_var, var.loc)
        out.append(const_assign)
        # get size: Asize0 = A_sh_attr[0]
        size_var = ir.Var(var.scope, mk_unique_var(var.name+"size"+str(i)), var.loc)
        self.type_annotation.typemap[size_var.name] = INT_TYPE
        getitem_node = ir.Expr.static_getitem(attr_var, i, const_var, var.loc)
        self.type_annotation.calltypes[getitem_node] = None
        getitem_assign = ir.Assign(getitem_node, size_var, var.loc)
        out.append(getitem_assign)
        return out

    # lhs is array so rhs has to return array
    def _analyze_rhs_classes(self, node):
        if isinstance(node, ir.Arg):
            assert self._isarray(node.name)
            return self._add_array_corr(node.name)
        elif isinstance(node, ir.Var):
            return self.array_shape_classes[node.name].copy()
        elif isinstance(node, ir.Expr):
            if node.op=='unary' and node.fn in UNARY_MAP_OP:
                assert isinstance(node.value, ir.Var)
                in_var = node.value.name
{'fraction_non_alphanumeric': 0.09879302637460885, 'fraction_numerical': 0.0031291908806437194, 'mean_word_length': 2.1613283089967026, 'pattern_counts': {'":': 0, '<': 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': '7425495', 'n_tokens_mistral': 4090, 'n_tokens_neox': 4033, 'n_words': 966}