diff --git a/doc/ref_transform.rst b/doc/ref_transform.rst index 740c5cb5848dbb7c6f657011bfc23fa88ca173ec..57d33b53999e06cbb07cc8363bbc46c091033cb3 100644 --- a/doc/ref_transform.rst +++ b/doc/ref_transform.rst @@ -118,7 +118,7 @@ Finishing up .. autofunction:: generate_loop_schedules -.. autofunction:: get_one_scheduled_kernel +.. autofunction:: get_one_linearized_kernel .. autofunction:: save_and_reload_temporaries diff --git a/doc/tutorial.rst b/doc/tutorial.rst index 753b09b5da42835b88a000bc0400fa18a254d80f..1b017f701f8161e93c4fdc1c14644dfe4b4fa74c 100644 --- a/doc/tutorial.rst +++ b/doc/tutorial.rst @@ -1204,16 +1204,16 @@ Here is what happens when we try to generate code for the kernel: This happens due to the kernel splitting done by :mod:`loopy`. The splitting happens when the instruction schedule is generated. To see the schedule, we -should call :func:`loopy.get_one_scheduled_kernel`: +should call :func:`loopy.get_one_linearized_kernel`: - >>> knl = lp.get_one_scheduled_kernel(lp.preprocess_kernel(knl)) + >>> knl = lp.get_one_linearized_kernel(lp.preprocess_kernel(knl)) >>> print(knl) --------------------------------------------------------------------------- KERNEL: rotate_v2 --------------------------------------------------------------------------- ... --------------------------------------------------------------------------- - SCHEDULE: + LINEARIZATION: 0: CALL KERNEL rotate_v2(extra_args=[], extra_inames=[]) 1: tmp = arr[i_inner + i_outer*16] {id=maketmp} 2: RETURN FROM KERNEL rotate_v2 @@ -1233,12 +1233,12 @@ goes for local temporaries). :func:`loopy.save_and_reload_temporaries` for the purpose of handling the task of saving and restoring temporary values across global barriers. This function adds instructions to the kernel without scheduling them. That means -that :func:`loopy.get_one_scheduled_kernel` needs to be called one more time to +that :func:`loopy.get_one_linearized_kernel` needs to be called one more time to put those instructions into the schedule. - >>> knl = lp.get_one_scheduled_kernel(lp.preprocess_kernel(knl)) + >>> knl = lp.get_one_linearized_kernel(lp.preprocess_kernel(knl)) >>> knl = lp.save_and_reload_temporaries(knl) - >>> knl = lp.get_one_scheduled_kernel(knl) # Schedule added instructions + >>> knl = lp.get_one_linearized_kernel(knl) # Schedule added instructions >>> print(knl) --------------------------------------------------------------------------- KERNEL: rotate_v2 @@ -1251,7 +1251,7 @@ put those instructions into the schedule. --------------------------------------------------------------------------- ... --------------------------------------------------------------------------- - SCHEDULE: + LINEARIZATION: 0: CALL KERNEL rotate_v2(extra_args=['tmp_save_slot'], extra_inames=[]) 1: tmp = arr[i_inner + i_outer*16] {id=maketmp} 2: tmp_save_slot[tmp_save_hw_dim_0_rotate_v2, tmp_save_hw_dim_1_rotate_v2] = tmp {id=tmp.save} diff --git a/loopy/__init__.py b/loopy/__init__.py index b60de6e2dcd35c1c167bf5e303401f2c6242ebec..67ad7de87decdb8694d7098e3446edf5d043df30 100644 --- a/loopy/__init__.py +++ b/loopy/__init__.py @@ -123,7 +123,8 @@ from loopy.transform.add_barrier import add_barrier from loopy.type_inference import infer_unknown_types from loopy.preprocess import preprocess_kernel, realize_reduction -from loopy.schedule import generate_loop_schedules, get_one_scheduled_kernel +from loopy.schedule import ( + generate_loop_schedules, get_one_scheduled_kernel, get_one_linearized_kernel) from loopy.statistics import (ToCountMap, CountGranularity, stringify_stats_mapping, Op, MemAccess, get_op_poly, get_op_map, get_lmem_access_poly, get_DRAM_access_poly, get_gmem_access_poly, get_mem_access_map, @@ -248,7 +249,8 @@ __all__ = [ "infer_unknown_types", "preprocess_kernel", "realize_reduction", - "generate_loop_schedules", "get_one_scheduled_kernel", + "generate_loop_schedules", + "get_one_scheduled_kernel", "get_one_linearized_kernel", "GeneratedProgram", "CodeGenerationResult", "PreambleInfo", "generate_code", "generate_code_v2", "generate_body", diff --git a/loopy/kernel/__init__.py b/loopy/kernel/__init__.py index 9096edcc0d7eb0111c393eb4cc5ed78405dba408..c3cd1738d7160c7feb6ef5d1042e3d41e19cdfdb 100644 --- a/loopy/kernel/__init__.py +++ b/loopy/kernel/__init__.py @@ -98,10 +98,25 @@ class _UniqueVarNameGenerator(UniqueNameGenerator): # {{{ loop kernel object +class _deprecated_KernelState_SCHEDULED(object): # noqa + def __init__(self, f): + self.f = f + + def __get__(self, obj, klass): + warn( + "'KernelState.SCHEDULED' is deprecated. " + "Use 'KernelState.LINEARIZED'.", + DeprecationWarning, stacklevel=2) + return self.f() + class KernelState: # noqa INITIAL = 0 PREPROCESSED = 1 - SCHEDULED = 2 + LINEARIZED = 2 + + @_deprecated_KernelState_SCHEDULED + def SCHEDULED(): # pylint:disable=no-method-argument + return KernelState.LINEARIZED # {{{ kernel_state, KernelState compataibility @@ -227,7 +242,9 @@ class LoopKernel(ImmutableRecordWithoutPickling): # {{{ constructor - def __init__(self, domains, instructions, args=None, schedule=None, + def __init__(self, domains, instructions, args=None, + schedule=None, + linearization=None, name="loopy_kernel", preambles=None, preamble_generators=None, @@ -336,6 +353,23 @@ class LoopKernel(ImmutableRecordWithoutPickling): ]: raise ValueError("invalid value for 'state'") + # `linearization` is replacing `schedule`, but we're not changing + # this under the hood yet, so for now, store it inside `schedule` + # and raise deprecation warning anyway + if schedule is not None: + if linearization is not None: + # these should not both be present + raise ValueError( + "received both `schedule` and `linearization` args, " + "'LoopKernel.schedule' is deprecated. " + "Use 'LoopKernel.linearization'.") + warn( + "'LoopKernel.schedule' is deprecated. " + "Use 'LoopKernel.linearization'.", + DeprecationWarning, stacklevel=2) + elif linearization is not None: + schedule = linearization + from collections import defaultdict assert not isinstance(iname_to_tags, defaultdict) @@ -1344,7 +1378,7 @@ class LoopKernel(ImmutableRecordWithoutPickling): if "schedule" in what and kernel.schedule is not None: lines.extend(sep) if show_labels: - lines.append("SCHEDULE:") + lines.append("LINEARIZATION:") from loopy.schedule import dump_schedule lines.append(dump_schedule(kernel, kernel.schedule)) @@ -1394,6 +1428,14 @@ class LoopKernel(ImmutableRecordWithoutPickling): # }}} + # {{{ handle linearization variable that doesn't yet exist + + @property + def linearization(self): + return self.schedule + + # }}} + # {{{ direct execution def __call__(self, *args, **kwargs): diff --git a/loopy/schedule/__init__.py b/loopy/schedule/__init__.py index f145c7122b9fd6e9e516d0becf3d4461fc0cce8c..0983c5e0d513d51a04a3f6cf3033904435ef1412 100644 --- a/loopy/schedule/__init__.py +++ b/loopy/schedule/__init__.py @@ -2029,6 +2029,15 @@ def _get_one_scheduled_kernel_inner(kernel): def get_one_scheduled_kernel(kernel): + warn_with_kernel( + kernel, "get_one_scheduled_kernel_deprecated", + "get_one_scheduled_kernel is deprecated. " + "Use get_one_linearized_kernel instead.", + DeprecationWarning) + return get_one_linearized_kernel(kernel) + + +def get_one_linearized_kernel(kernel): from loopy import CACHING_ENABLED sched_cache_key = kernel