loopy/schedule/device_mapping.py
def restore_and_save_temporaries(kernel):
    """
    Add code that loads / spills the temporaries in the kernel which are
    live across sub-kernel calls.
    """
    # Compute live temporaries.
    live_in, live_out = compute_live_temporaries(kernel, kernel.schedule)

    # Create kernel variables based on live temporaries.
    inter_kernel_temporaries = set()

    call_count = 0
    for idx, sched_item in enumerate(kernel.schedule):
        if isinstance(sched_item, CallKernel):
            inter_kernel_temporaries |= filter_out_subscripts(live_in[idx])
            call_count += 1

    if call_count == 1:
        # Single call corresponds to a kernel which has not been split -
        # no need for restores / spills of temporaries.
        return kernel

    name_gen = kernel.get_var_name_generator()
    new_temporaries = determine_temporaries_to_promote(
        kernel, inter_kernel_temporaries, name_gen)

    # {{{ Insert loads and spills of new temporaries

    new_schedule = []
    new_instructions = []
    new_iname_to_tag = {}

    idx = 0
    schedule = kernel.schedule
    while idx < len(schedule):
        sched_item = schedule[idx]

        if not isinstance(sched_item, CallKernel):
            new_schedule.append(sched_item)
            idx += 1
            continue

        subkernel_prolog = []
        subkernel_epilog = []
        subkernel_schedule = []

        start_idx = idx
        idx += 1
        while not isinstance(schedule[idx], ReturnFromKernel):
            subkernel_schedule.append(schedule[idx])
            idx += 1

        subkernel_defs, subkernel_uses = \
            get_temporaries_defined_and_used_in_subrange(
                kernel, schedule, start_idx + 1, idx - 1)

        from loopy.kernel.data import temp_var_scope
        # Filter out temporaries that are global.
        subkernel_globals = set(
            tval for tval in subkernel_defs | subkernel_uses
            if kernel.temporary_variables[tval].scope == temp_var_scope.GLOBAL)

        tvals_to_spill = (subkernel_defs - subkernel_globals) & live_out[idx]
        # Need to load tvals_to_spill, to avoid overwriting entries that the
        # code doesn't touch when doing the spill.
        tvals_to_load = ((subkernel_uses - subkernel_globals)
            | tvals_to_spill) & live_in[start_idx]

        # Add new arguments.
        sched_item = sched_item.copy(
            extra_args=sched_item.extra_args
            + sorted(new_temporaries[tv].name
                     for tv in tvals_to_load | tvals_to_spill))

        # {{{ Add all the loads and spills.

        def insert_loads_or_spills(tvals, mode):
            assert mode in ["load", "spill"]
            local_temporaries = set()

            code_block = \
                subkernel_prolog if mode == "load" else subkernel_epilog

            new_kernel = kernel

            for tval in tvals:
                from loopy.kernel.tools import DomainChanger
                tval_hw_inames = new_temporaries[tval].hw_inames
                dchg = DomainChanger(kernel,
                    frozenset(sched_item.extra_inames + tval_hw_inames))
                domain = dchg.domain

                domain, hw_inames, dim_inames, itt = \
                    augment_domain_for_temporary_promotion(
                        new_kernel, domain, new_temporaries[tval], mode,
                        name_gen)
                new_iname_to_tag.update(itt)

                new_kernel = dchg.get_kernel_with(domain)

                # Add the load / spill instruction.
                insn_id = name_gen("{name}.{mode}".format(name=tval, mode=mode))

                def subscript_or_var(agg, subscript):
                    from pymbolic.primitives import Subscript, Variable
                    if len(subscript) == 0:
                        return Variable(agg)
                    else:
                        return Subscript(
                            Variable(agg),
                            tuple(map(Variable, subscript)))

                args = (
                    subscript_or_var(
                        tval, dim_inames),
                    subscript_or_var(
                        new_temporaries[tval].name, hw_inames + dim_inames))

                if mode == "spill":
                    args = reversed(args)

                from loopy.kernel.data import Assignment
                new_insn = Assignment(*args, id=insn_id)

                new_instructions.append(new_insn)

                loop_begin = [EnterLoop(iname=iname) for iname in dim_inames]
                loop_end = list(reversed([
                    LeaveLoop(iname=iname) for iname in dim_inames]))
                code_block.extend(
                    loop_begin +
                    [RunInstruction(insn_id=insn_id)] +
                    loop_end)
                if new_temporaries[tval].orig_temporary.is_local:
                    local_temporaries.add(new_temporaries[tval].name)

            # After loading / before spilling local temporaries, we need to
            # insert a barrier.
            if local_temporaries:
                if mode == "load":
                    subkernel_prolog.append(
                        Barrier(kind="local",
                                comment="for loads of {0}".format(
                                    ", ".join(sorted(local_temporaries)))))
                else:
                    subkernel_epilog.insert(0,
                        Barrier(kind="local",
                                comment="for spills of {0}".format(
                                    ", ".join(sorted(local_temporaries)))))
            return new_kernel

        kernel = insert_loads_or_spills(tvals_to_load, "load")
        kernel = insert_loads_or_spills(tvals_to_spill, "spill")

        # }}}

        new_schedule.extend(
            [sched_item] +
            subkernel_prolog +
            subkernel_schedule +
            subkernel_epilog +
            # ReturnFromKernel
            [schedule[idx]])

        # ReturnFromKernel
        idx += 1

    # }}}

    new_iname_to_tag.update(kernel.iname_to_tag)
    new_temporary_variables = dict(
        (t.name, t.as_variable()) for t in new_temporaries.values())
    new_temporary_variables.update(kernel.temporary_variables)

    kernel = kernel.copy(
        iname_to_tag=new_iname_to_tag,
        temporary_variables=new_temporary_variables,
        instructions=kernel.instructions + new_instructions,
        schedule=new_schedule
        )

    return kernel

    Add code that loads / spills the temporaries in the kernel which are
    live across sub-kernel calls.
    
Add code that loads / spills the temporaries in the kernel which are
live across sub-kernel calls.
Add code that loads / spills the temporaries in the kernel which are live across sub - kernel calls .
def restore_and_save_temporaries(kernel):
    live_in, live_out = compute_live_temporaries(kernel, kernel.schedule)
    inter_kernel_temporaries = set()
    call_count = 0
    for idx, sched_item in enumerate(kernel.schedule):
        if isinstance(sched_item, CallKernel):
            inter_kernel_temporaries |= filter_out_subscripts(live_in[idx])
            call_count += 1
    if call_count == 1:
        return kernel
    name_gen = kernel.get_var_name_generator()
    new_temporaries = determine_temporaries_to_promote(
        kernel, inter_kernel_temporaries, name_gen)
    new_schedule = []
    new_instructions = []
    new_iname_to_tag = {}
    idx = 0
    schedule = kernel.schedule
    while idx < len(schedule):
        sched_item = schedule[idx]
        if not isinstance(sched_item, CallKernel):
            new_schedule.append(sched_item)
            idx += 1
            continue
        subkernel_prolog = []
        subkernel_epilog = []
        subkernel_schedule = []
        start_idx = idx
        idx += 1
        while not isinstance(schedule[idx], ReturnFromKernel):
            subkernel_schedule.append(schedule[idx])
            idx += 1
        subkernel_defs, subkernel_uses = \
            get_temporaries_defined_and_used_in_subrange(
                kernel, schedule, start_idx + 1, idx - 1)
        from loopy.kernel.data import temp_var_scope
        subkernel_globals = set(
            tval for tval in subkernel_defs | subkernel_uses
            if kernel.temporary_variables[tval].scope == temp_var_scope.GLOBAL)
        tvals_to_spill = (subkernel_defs - subkernel_globals) & live_out[idx]
        tvals_to_load = ((subkernel_uses - subkernel_globals)
            | tvals_to_spill) & live_in[start_idx]
        sched_item = sched_item.copy(
            extra_args=sched_item.extra_args
            + sorted(new_temporaries[tv].name
                     for tv in tvals_to_load | tvals_to_spill))
        def insert_loads_or_spills(tvals, mode):
            assert mode in ["load", "spill"]
            local_temporaries = set()
            code_block = \
                subkernel_prolog if mode == "load" else subkernel_epilog
            new_kernel = kernel
            for tval in tvals:
                from loopy.kernel.tools import DomainChanger
                tval_hw_inames = new_temporaries[tval].hw_inames
                dchg = DomainChanger(kernel,
                    frozenset(sched_item.extra_inames + tval_hw_inames))
                domain = dchg.domain
                domain, hw_inames, dim_inames, itt = \
                    augment_domain_for_temporary_promotion(
                        new_kernel, domain, new_temporaries[tval], mode,
                        name_gen)
                new_iname_to_tag.update(itt)
                new_kernel = dchg.get_kernel_with(domain)
                insn_id = name_gen("{name}.{mode}".format(name=tval, mode=mode))
                def subscript_or_var(agg, subscript):
                    from pymbolic.primitives import Subscript, Variable
                    if len(subscript) == 0:
                        return Variable(agg)
                    else:
                        return Subscript(
                            Variable(agg),
                            tuple(map(Variable, subscript)))
                args = (
                    subscript_or_var(
                        tval, dim_inames),
                    subscript_or_var(
                        new_temporaries[tval].name, hw_inames + dim_inames))
                if mode == "spill":
                    args = reversed(args)
                from loopy.kernel.data import Assignment
                new_insn = Assignment(*args, id=insn_id)
                new_instructions.append(new_insn)
                loop_begin = [EnterLoop(iname=iname) for iname in dim_inames]
                loop_end = list(reversed([
                    LeaveLoop(iname=iname) for iname in dim_inames]))
                code_block.extend(
                    loop_begin +
                    [RunInstruction(insn_id=insn_id)] +
                    loop_end)
                if new_temporaries[tval].orig_temporary.is_local:
                    local_temporaries.add(new_temporaries[tval].name)
            if local_temporaries:
                if mode == "load":
                    subkernel_prolog.append(
                        Barrier(kind="local",
                                comment="for loads of {0}".format(
                                    ", ".join(sorted(local_temporaries)))))
                else:
                    subkernel_epilog.insert(0,
                        Barrier(kind="local",
                                comment="for spills of {0}".format(
                                    ", ".join(sorted(local_temporaries)))))
            return new_kernel
        kernel = insert_loads_or_spills(tvals_to_load, "load")
        kernel = insert_loads_or_spills(tvals_to_spill, "spill")
        new_schedule.extend(
            [sched_item] +
            subkernel_prolog +
            subkernel_schedule +
            subkernel_epilog +
            [schedule[idx]])
        idx += 1
    new_iname_to_tag.update(kernel.iname_to_tag)
    new_temporary_variables = dict(
        (t.name, t.as_variable()) for t in new_temporaries.values())
    new_temporary_variables.update(kernel.temporary_variables)
    kernel = kernel.copy(
        iname_to_tag=new_iname_to_tag,
        temporary_variables=new_temporary_variables,
        instructions=kernel.instructions + new_instructions,
        schedule=new_schedule
        )
    return kernel
def restore_and_save_temporaries ( kernel ) : live_in , live_out = compute_live_temporaries ( kernel , kernel . schedule ) inter_kernel_temporaries = set ( ) call_count = 0 for idx , sched_item in enumerate ( kernel . schedule ) : if isinstance ( sched_item , CallKernel ) : inter_kernel_temporaries |= filter_out_subscripts ( live_in [ idx ] ) call_count += 1 if call_count == 1 : return kernel name_gen = kernel . get_var_name_generator ( ) new_temporaries = determine_temporaries_to_promote ( kernel , inter_kernel_temporaries , name_gen ) new_schedule = [ ] new_instructions = [ ] new_iname_to_tag = { } idx = 0 schedule = kernel . schedule while idx < len ( schedule ) : sched_item = schedule [ idx ] if not isinstance ( sched_item , CallKernel ) : new_schedule . append ( sched_item ) idx += 1 continue subkernel_prolog = [ ] subkernel_epilog = [ ] subkernel_schedule = [ ] start_idx = idx idx += 1 while not isinstance ( schedule [ idx ] , ReturnFromKernel ) : subkernel_schedule . append ( schedule [ idx ] ) idx += 1 subkernel_defs , subkernel_uses = get_temporaries_defined_and_used_in_subrange ( kernel , schedule , start_idx + 1 , idx - 1 ) from loopy . kernel . data import temp_var_scope subkernel_globals = set ( tval for tval in subkernel_defs | subkernel_uses if kernel . temporary_variables [ tval ] . scope == temp_var_scope . GLOBAL ) tvals_to_spill = ( subkernel_defs - subkernel_globals ) & live_out [ idx ] tvals_to_load = ( ( subkernel_uses - subkernel_globals ) | tvals_to_spill ) & live_in [ start_idx ] sched_item = sched_item . copy ( extra_args = sched_item . extra_args + sorted ( new_temporaries [ tv ] . name for tv in tvals_to_load | tvals_to_spill ) ) def insert_loads_or_spills ( tvals , mode ) : assert mode in [ "load" , "spill" ] local_temporaries = set ( ) code_block = subkernel_prolog if mode == "load" else subkernel_epilog new_kernel = kernel for tval in tvals : from loopy . kernel . tools import DomainChanger tval_hw_inames = new_temporaries [ tval ] . hw_inames dchg = DomainChanger ( kernel , frozenset ( sched_item . extra_inames + tval_hw_inames ) ) domain = dchg . domain domain , hw_inames , dim_inames , itt = augment_domain_for_temporary_promotion ( new_kernel , domain , new_temporaries [ tval ] , mode , name_gen ) new_iname_to_tag . update ( itt ) new_kernel = dchg . get_kernel_with ( domain ) insn_id = name_gen ( "{name}.{mode}" . format ( name = tval , mode = mode ) ) def subscript_or_var ( agg , subscript ) : from pymbolic . primitives import Subscript , Variable if len ( subscript ) == 0 : return Variable ( agg ) else : return Subscript ( Variable ( agg ) , tuple ( map ( Variable , subscript ) ) ) args = ( subscript_or_var ( tval , dim_inames ) , subscript_or_var ( new_temporaries [ tval ] . name , hw_inames + dim_inames ) ) if mode == "spill" : args = reversed ( args ) from loopy . kernel . data import Assignment new_insn = Assignment ( * args , id = insn_id ) new_instructions . append ( new_insn ) loop_begin = [ EnterLoop ( iname = iname ) for iname in dim_inames ] loop_end = list ( reversed ( [ LeaveLoop ( iname = iname ) for iname in dim_inames ] ) ) code_block . extend ( loop_begin + [ RunInstruction ( insn_id = insn_id ) ] + loop_end ) if new_temporaries [ tval ] . orig_temporary . is_local : local_temporaries . add ( new_temporaries [ tval ] . name ) if local_temporaries : if mode == "load" : subkernel_prolog . append ( Barrier ( kind = "local" , comment = "for loads of {0}" . format ( ", " . join ( sorted ( local_temporaries ) ) ) ) ) else : subkernel_epilog . insert ( 0 , Barrier ( kind = "local" , comment = "for spills of {0}" . format ( ", " . join ( sorted ( local_temporaries ) ) ) ) ) return new_kernel kernel = insert_loads_or_spills ( tvals_to_load , "load" ) kernel = insert_loads_or_spills ( tvals_to_spill , "spill" ) new_schedule . extend ( [ sched_item ] + subkernel_prolog + subkernel_schedule + subkernel_epilog + [ schedule [ idx ] ] ) idx += 1 new_iname_to_tag . update ( kernel . iname_to_tag ) new_temporary_variables = dict ( ( t . name , t . as_variable ( ) ) for t in new_temporaries . values ( ) ) new_temporary_variables . update ( kernel . temporary_variables ) kernel = kernel . copy ( iname_to_tag = new_iname_to_tag , temporary_variables = new_temporary_variables , instructions = kernel . instructions + new_instructions , schedule = new_schedule ) return kernel
Add code that loads / spills the temporaries in the kernel which are
live across sub-kernel calls.
Add code that loads / spills the temporaries in the kernel which are live across sub - kernel calls .
"""
    Add code that loads / spills the temporaries in the kernel which are
    live across sub-kernel calls.
    """ # Compute live temporaries. # Create kernel variables based on live temporaries. # Single call corresponds to a kernel which has not been split - # no need for restores / spills of temporaries. # {{{ Insert loads and spills of new temporaries # Filter out temporaries that are global. # Need to load tvals_to_spill, to avoid overwriting entries that the # code doesn't touch when doing the spill. # Add new arguments. # {{{ Add all the loads and spills. # Add the load / spill instruction. # After loading / before spilling local temporaries, we need to # insert a barrier. # }}} # ReturnFromKernel # ReturnFromKernel # }}}
kernel