Ejemplo n.º 1
0
def create_prange_closure(env, prange_node, body, target):
    # Find referenced and assigned variables
    v = VariableFindingVisitor()
    v.visitlist(body)

    # Determine privates and reductions. Shared variables will be handled by
    # the closure support.
    privates = set(v.assigned) - set(v.reductions)
    reductions = v.reductions

    if isinstance(target, ast.Name) and target.id in reductions:
        # Remove target variable from reductions if present
        reductions.pop(target.id)
    privates.add(target.id)

    privates_struct_type = numba.struct([])
    privates_struct = ast.Name('__numba_privates', ast.Param())

    args = [privates_struct]
    func_def = ast.FunctionDef(name=templating.temp_name("prange_body"),
                               args=ast.arguments(args=args,
                                                  vararg=None,
                                                  kwarg=None,
                                                  defaults=[]),
                               body=copy.deepcopy(body),
                               decorator_list=[])

    # Update outlined prange body closure

    func_signature = void(privates_struct_type.ref())
    # func_signature.struct_by_reference = True
    need_closure_wrapper = False
    locals_dict = {'__numba_privates': privates_struct_type.ref()}

    func_env = env.translation.make_partial_env(
        func_def,
        func_signature=func_signature,
        need_closure_wrapper=need_closure_wrapper,
        locals=locals_dict,
    )

    # Update prange node
    prange_node.func_env = func_env
    prange_node.privates_struct_type = privates_struct_type
    prange_node.privates = privates
    prange_node.reductions = reductions
    prange_node.func_def = func_def
Ejemplo n.º 2
0
def create_prange_closure(env, prange_node, body, target):
    # Find referenced and assigned variables
    v = VariableFindingVisitor()
    v.visitlist(body)

    # Determine privates and reductions. Shared variables will be handled by
    # the closure support.
    privates = set(v.assigned) - set(v.reductions)
    reductions = v.reductions

    if isinstance(target, ast.Name) and target.id in reductions:
        # Remove target variable from reductions if present
        reductions.pop(target.id)
    privates.add(target.id)

    privates_struct_type = numba.struct([])
    privates_struct = ast.Name('__numba_privates', ast.Param())

    args = [privates_struct]
    func_def = ast.FunctionDef(name=templating.temp_name("prange_body"),
                               args=ast.arguments(args=args, vararg=None,
                                                  kwarg=None, defaults=[]),
                               body=copy.deepcopy(body),
                               decorator_list=[])

    # Update outlined prange body closure

    func_signature = void(privates_struct_type.ref())
    # func_signature.struct_by_reference = True
    need_closure_wrapper = False
    locals_dict = { '__numba_privates': privates_struct_type.ref() }

    func_env = env.translation.make_partial_env(
        func_def,
        func_signature=func_signature,
        need_closure_wrapper=need_closure_wrapper,
        locals=locals_dict,
    )

    # Update prange node
    prange_node.func_env = func_env
    prange_node.privates_struct_type = privates_struct_type
    prange_node.privates = privates
    prange_node.reductions = reductions
    prange_node.func_def = func_def
Ejemplo n.º 3
0
    def register_array_expression(self, node, lhs=None):
        super(ArrayExpressionRewriteNative, self).register_array_expression(
            node, lhs)

        # llvm_module = llvm.core.Module.new(temp_name("array_expression_module"))
        # llvm_module = self.env.llvm_context.module

        lhs_type = lhs.type if lhs else node.type
        is_expr = lhs is None

        if node.type.is_array and lhs_type.ndim < node.type.ndim:
            # TODO: this is valid in NumPy if the leading dimensions of the
            # TODO: RHS have extent 1
            raise error.NumbaError(
                node, "Right hand side must have a "
                      "dimensionality <= %d" % lhs_type.ndim)

        # Create ufunc scalar kernel
        ufunc_ast, signature, ufunc_builder = get_py_ufunc_ast(self.env, lhs, node)

        # Compile ufunc scalar kernel with numba
        ast.fix_missing_locations(ufunc_ast)
        # func_env = self.env.crnt.inherit(
        #     func=None, ast=ufunc_ast, func_signature=signature,
        #     wrap=False, #link=False, #llvm_module=llvm_module,
        # )
        # pipeline.run_env(self.env, func_env) #, pipeline_name='codegen')

        func_env, (_, _, _) = pipeline.run_pipeline2(
            self.env, None, ufunc_ast, signature,
            function_globals=self.env.crnt.function_globals,
            wrap=False, link=False, nopython=True,
            #llvm_module=llvm_module, # pipeline_name='codegen',
        )
        llvm_module = func_env.llvm_module

        operands = ufunc_builder.operands
        operands = [nodes.CloneableNode(operand) for operand in operands]

        if lhs is not None:
            lhs = nodes.CloneableNode(lhs)
            broadcast_operands = [lhs] + operands
            lhs = lhs.clone
        else:
            broadcast_operands = operands[:]

        shape = slicenodes.BroadcastNode(lhs_type, broadcast_operands)
        operands = [op.clone for op in operands]

        if lhs is None and self.nopython:
            raise error.NumbaError(
                node, "Cannot allocate new memory in nopython context")
        elif lhs is None:
            # TODO: determine best output order at runtime
            shape = shape.cloneable
            lhs = nodes.ArrayNewEmptyNode(lhs_type, shape.clone,
                                          lhs_type.is_f_contig).cloneable

        # Build minivect wrapper kernel
        context = NumbaStaticArgsContext()
        context.llvm_module = llvm_module
        # context.llvm_ee = self.env.llvm_context.execution_engine

        b = context.astbuilder
        variables = [b.variable(name_node.type, "op%d" % i)
                     for i, name_node in enumerate([lhs] + operands)]
        miniargs = [b.funcarg(variable) for variable in variables]
        body = miniutils.build_kernel_call(func_env.lfunc.name, signature,
                                           miniargs, b)

        minikernel = b.function_from_numpy(
            temp_name("array_expression"), body, miniargs)
        lminikernel, = context.run_simple(minikernel,
                                          specializers.StridedSpecializer)
        # lminikernel.linkage = llvm.core.LINKAGE_LINKONCE_ODR

        # pipeline.run_env(self.env, func_env, pipeline_name='post_codegen')
        # llvm_module.verify()
        del func_env

        assert lminikernel.module is llvm_module
        # print("---------")
        # print(llvm_module)
        # print("~~~~~~~~~~~~")
        lminikernel = self.env.llvm_context.link(lminikernel)

        # Build call to minivect kernel
        operands.insert(0, lhs)
        args = [shape]
        scalar_args = []
        for operand in operands:
            if operand.type.is_array:
                data_p = self.array_attr(operand, 'data')
                data_p = nodes.CoercionNode(data_p,
                                            operand.type.dtype.pointer())
                if not isinstance(operand, nodes.CloneNode):
                    operand = nodes.CloneNode(operand)
                strides_p = self.array_attr(operand, 'strides')
                args.extend((data_p, strides_p))
            else:
                scalar_args.append(operand)

        args.extend(scalar_args)
        result = nodes.NativeCallNode(minikernel.type, args, lminikernel)

        # Use native slicing in array expressions
        slicenodes.mark_nopython(ast.Suite(body=result.args))

        if not is_expr:
            # a[:] = b[:] * c[:]
            return result

        # b[:] * c[:], return new array as expression
        return nodes.ExpressionNode(stmts=[result], expr=lhs.clone)
Ejemplo n.º 4
0
    def register_array_expression(self, node, lhs=None):
        super(ArrayExpressionRewriteNative, self).register_array_expression(
            node, lhs)

        lhs_type = lhs.type if lhs else node.type
        is_expr = lhs is None

        if node.type.is_array and lhs_type.ndim < node.type.ndim:
            # TODO: this is valid in NumPy if the leading dimensions of the
            # TODO: RHS have extent 1
            raise error.NumbaError(
                node, "Right hand side must have a "
                      "dimensionality <= %d" % lhs_type.ndim)

        # Create ufunc scalar kernel
        ufunc_ast, signature, ufunc_builder = self.get_py_ufunc_ast(lhs, node)
        signature.struct_by_reference = True

        # Compile ufunc scalar kernel with numba
        ast.fix_missing_locations(ufunc_ast)
        func_env, (_, _, _) = pipeline.run_pipeline2(
            self.env, None, ufunc_ast, signature,
            function_globals={},
        )

        # Manual linking
        lfunc = func_env.lfunc

        # print lfunc
        operands = ufunc_builder.operands
        functions.keep_alive(self.func, lfunc)

        operands = [nodes.CloneableNode(operand) for operand in operands]

        if lhs is not None:
            lhs = nodes.CloneableNode(lhs)
            broadcast_operands = [lhs] + operands
            lhs = lhs.clone
        else:
            broadcast_operands = operands[:]

        shape = slicenodes.BroadcastNode(lhs_type, broadcast_operands)
        operands = [op.clone for op in operands]

        if lhs is None and self.nopython:
            raise error.NumbaError(
                node, "Cannot allocate new memory in nopython context")
        elif lhs is None:
            # TODO: determine best output order at runtime
            shape = shape.cloneable
            lhs = nodes.ArrayNewEmptyNode(lhs_type, shape.clone,
                                          lhs_type.is_f_contig).cloneable

        # Build minivect wrapper kernel
        context = NumbaproStaticArgsContext()
        context.llvm_module = self.env.llvm_context.module
        # context.debug = True
        context.optimize_broadcasting = False
        b = context.astbuilder

        variables = [b.variable(name_node.type, "op%d" % i)
                     for i, name_node in enumerate([lhs] + operands)]
        miniargs = [b.funcarg(variable) for variable in variables]
        body = miniutils.build_kernel_call(lfunc.name, signature, miniargs, b)

        minikernel = b.function_from_numpy(
            templating.temp_name("array_expression"), body, miniargs)
        lminikernel, ctypes_kernel = context.run_simple(
            minikernel, specializers.StridedSpecializer)

        # Build call to minivect kernel
        operands.insert(0, lhs)
        args = [shape]
        scalar_args = []
        for operand in operands:
            if operand.type.is_array:
                data_p = self.array_attr(operand, 'data')
                data_p = nodes.CoercionNode(data_p,
                                            operand.type.dtype.pointer())
                if not isinstance(operand, nodes.CloneNode):
                    operand = nodes.CloneNode(operand)
                strides_p = self.array_attr(operand, 'strides')
                args.extend((data_p, strides_p))
            else:
                scalar_args.append(operand)

        args.extend(scalar_args)
        result = nodes.NativeCallNode(minikernel.type, args, lminikernel)

        # Use native slicing in array expressions
        slicenodes.mark_nopython(ast.Suite(body=result.args))

        if not is_expr:
            # a[:] = b[:] * c[:]
            return result

        # b[:] * c[:], return new array as expression
        return nodes.ExpressionNode(stmts=[result], expr=lhs.clone)