Tools : PyTeal optimization utility

docs/compiler_optimization.rst

@@ -0,0 +1,30 @@
+.. _compiler_optimization:
+
+Compiler Optimization
+========================
+**The optimizer is at an early stage and is disabled by default. Backwards compatability cannot be
+guaranteed at this point.**
+
+The optimizer is a tool for improving performance and reducing resource consumption. In this context,
+the terms *performance* and *resource* can apply across multiple dimensions, including but not limited
+to: compiled code size, scratch slot usage, opcode cost, etc. 
+
+Optimizer Usage
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The compiler determines which optimizations to apply based on the provided :any:`OptimizeOptions` object as
+shown in the code block below. The :any:`OptimizeOptions` constructor receives a set of keyword arguments 
+representing flags corresponding to particular optimizations. If arguments are not provided to the
+constructor or no :any:`OptimizeOptions` object is passed to :any:`compileTeal` then the default behavior is
+that no optimizations are applied.
+
+============================== ================================================================================ ===========================
+Optimization Flag              Description                                                                      Default
+============================== ================================================================================ ===========================
+:code:`scratch_slots`          A boolean describing whether or not scratch slot optimization should be applied. :code:`False`
+============================== ================================================================================ ===========================
+
+.. code-block:: python
+
+    optimize_options = OptimizeOptions(scratch_slots=True)
+    compileTeal(approval_program(), mode=Mode.Application, version=4, optimize=optimize_options)

docs/index.rst

@@ -43,6 +43,7 @@ PyTeal **hasn't been security audited**. Use it at your own risk.
    state
    assets
    versions
+   compiler_optimization
 
 .. toctree::
    :maxdepth: 3

pyteal/__init__.py

@@ -8,6 +8,7 @@
     DEFAULT_TEAL_VERSION,
     CompileOptions,
     compileTeal,
+    OptimizeOptions,
 )
 from .types import TealType
 from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
@@ -23,6 +24,7 @@
         "DEFAULT_TEAL_VERSION",
         "CompileOptions",
         "compileTeal",
+        "OptimizeOptions",
         "TealType",
         "TealInternalError",
         "TealTypeError",

pyteal/__init__.pyi

@@ -11,6 +11,7 @@ from .compiler import (
     DEFAULT_TEAL_VERSION,
     CompileOptions,
     compileTeal,
+    OptimizeOptions,
 )
 from .types import TealType
 from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
@@ -122,6 +123,7 @@ __all__ = [
     "Not",
     "OnComplete",
     "Op",
+    "OptimizeOptions",
     "Or",
     "Pop",
     "Reject",

pyteal/compiler/__init__.py

@@ -6,10 +6,13 @@
     compileTeal,
 )
 
+from .optimizer import OptimizeOptions
+
 __all__ = [
     "MAX_TEAL_VERSION",
     "MIN_TEAL_VERSION",
     "DEFAULT_TEAL_VERSION",
     "CompileOptions",
     "compileTeal",
+    "OptimizeOptions",
 ]

pyteal/compiler/compiler.py

@@ -1,5 +1,7 @@
 from typing import List, Tuple, Set, Dict, Optional, cast
 
+from .optimizer import OptimizeOptions, apply_global_optimizations
+
 from ..types import TealType
 from ..ast import (
     Expr,
@@ -13,7 +15,10 @@
 
 from .sort import sortBlocks
 from .flatten import flattenBlocks, flattenSubroutines
-from .scratchslots import assignScratchSlotsToSubroutines
+from .scratchslots import (
+    assignScratchSlotsToSubroutines,
+    collect_unoptimized_slots,
+)
 from .subroutines import (
     spillLocalSlotsDuringRecursion,
     resolveSubroutines,
@@ -31,9 +36,11 @@ def __init__(
         *,
         mode: Mode = Mode.Signature,
         version: int = DEFAULT_TEAL_VERSION,
+        optimize: OptimizeOptions = None,
     ) -> None:
         self.mode = mode
         self.version = version
+        self.optimize = optimize if optimize is not None else OptimizeOptions()
 
         self.currentSubroutine: Optional[SubroutineDefinition] = None
 
@@ -109,9 +116,9 @@ def verifyOpsForMode(teal: List[TealComponent], mode: Mode):
 def compileSubroutine(
     ast: Expr,
     options: CompileOptions,
-    subroutineMapping: Dict[Optional[SubroutineDefinition], List[TealComponent]],
     subroutineGraph: Dict[SubroutineDefinition, Set[SubroutineDefinition]],
-    subroutineBlocks: Dict[Optional[SubroutineDefinition], TealBlock],
+    subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
+    subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
 ) -> None:
     currentSubroutine = (
         cast(SubroutineDeclaration, ast).subroutine
@@ -139,8 +146,8 @@ def compileSubroutine(
     verifyOpsForVersion(teal, options.version)
     verifyOpsForMode(teal, options.mode)
 
-    subroutineMapping[currentSubroutine] = teal
-    subroutineBlocks[currentSubroutine] = start
+    subroutine_start_blocks[currentSubroutine] = start
+    subroutine_end_blocks[currentSubroutine] = end
 
     referencedSubroutines: Set[SubroutineDefinition] = set()
     for stmt in teal:
@@ -150,23 +157,38 @@ def compileSubroutine(
     if currentSubroutine is not None:
         subroutineGraph[currentSubroutine] = referencedSubroutines
 
-    newSubroutines = referencedSubroutines - subroutineMapping.keys()
+    newSubroutines = referencedSubroutines - subroutine_start_blocks.keys()
     for subroutine in sorted(newSubroutines, key=lambda subroutine: subroutine.id):
         compileSubroutine(
             subroutine.getDeclaration(),
             options,
-            subroutineMapping,
             subroutineGraph,
-            subroutineBlocks,
+            subroutine_start_blocks,
+            subroutine_end_blocks,
         )
 
 
+def sort_subroutine_blocks(
+    subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
+    subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
+) -> Dict[Optional[SubroutineDefinition], List[TealComponent]]:
+    subroutine_mapping: Dict[
+        Optional[SubroutineDefinition], List[TealComponent]
+    ] = dict()
+    for subroutine, start in subroutine_start_blocks.items():
+        order = sortBlocks(start, subroutine_end_blocks[subroutine])
+        subroutine_mapping[subroutine] = flattenBlocks(order)
+
+    return subroutine_mapping
+
+
 def compileTeal(
     ast: Expr,
     mode: Mode,
     *,
     version: int = DEFAULT_TEAL_VERSION,
     assembleConstants: bool = False,
+    optimize: OptimizeOptions = None,
 ) -> str:
     """Compile a PyTeal expression into TEAL assembly.
 
@@ -181,6 +203,7 @@ def compileTeal(
             constants will be assembled in the most space-efficient way, so enabling this may reduce
             the compiled program's size. Enabling this option requires a minimum TEAL version of 3.
             Defaults to false.
+        optimize (optional): OptimizeOptions that determine which optimizations will be applied.
 
     Returns:
         A TEAL assembly program compiled from the input expression.
@@ -199,20 +222,32 @@ def compileTeal(
             )
         )
 
-    options = CompileOptions(mode=mode, version=version)
+    options = CompileOptions(mode=mode, version=version, optimize=optimize)
 
-    subroutineMapping: Dict[
-        Optional[SubroutineDefinition], List[TealComponent]
-    ] = dict()
     subroutineGraph: Dict[SubroutineDefinition, Set[SubroutineDefinition]] = dict()
-    subroutineBlocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
+    subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
+    subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
     compileSubroutine(
-        ast, options, subroutineMapping, subroutineGraph, subroutineBlocks
+        ast, options, subroutineGraph, subroutine_start_blocks, subroutine_end_blocks
     )
 
-    localSlotAssignments = assignScratchSlotsToSubroutines(
-        subroutineMapping, subroutineBlocks
-    )
+    # note: optimizations are off by default, in which case, apply_global_optimizations
+    # won't make any changes. Because the optimizer is invoked on a subroutine's
+    # control flow graph, the optimizer requires context across block boundaries. This
+    # is necessary for the dependency checking of local slots. Global slots, slots
+    # used by DynamicScratchVar, and reserved slots are not optimized.
+    if options.optimize.scratch_slots:
+        options.optimize._skip_slots = collect_unoptimized_slots(
+            subroutine_start_blocks
+        )
+        for start in subroutine_start_blocks.values():
+            apply_global_optimizations(start, options.optimize)
+
+    localSlotAssignments = assignScratchSlotsToSubroutines(subroutine_start_blocks)
+
+    subroutineMapping: Dict[
+        Optional[SubroutineDefinition], List[TealComponent]
+    ] = sort_subroutine_blocks(subroutine_start_blocks, subroutine_end_blocks)
 
     spillLocalSlotsDuringRecursion(
         version, subroutineMapping, subroutineGraph, localSlotAssignments

pyteal/compiler/optimizer/__init__.py

@@ -0,0 +1 @@
+from .optimizer import OptimizeOptions, apply_global_optimizations

pyteal/compiler/optimizer/optimizer.py

@@ -0,0 +1,102 @@
+from typing import Set
+from ...ast import ScratchSlot
+from ...ir import TealBlock, TealOp, Op
+from ...errors import TealInternalError
+
+
+class OptimizeOptions:
+    """An object which specifies the optimizations to be performed and relevant context.
+
+    _skip_slots: the slots that should be skipped during optimization. At the moment this includes:
+            1. reserved slots because they may have dependencies outside
+            the current application. For example, the 'gloads' opcode can
+            access the slots of other applications in the tx group.
+            2. global slots because they're outside the scope of global
+            optimizations, which only apply to the control flow graph of
+            a single subroutine.
+            3. slots used with dynamic scratch vars. These slots use
+            indirection by means of the 'stores' opcode and dependencies
+            can only be determined at runtime.
+
+    Args:
+
+        scratch_slots (optional): cancel contiguous store/load operations
+            that have no load dependencies elsewhere.
+    """
+
+    def __init__(self, *, scratch_slots: bool = False):
+        self.scratch_slots = scratch_slots
+        self._skip_slots: Set[ScratchSlot] = set()
+
+
+def _remove_extraneous_slot_access(start: TealBlock, remove: Set[ScratchSlot]):
+    def keep_op(op: TealOp) -> bool:
+        if type(op) != TealOp or (op.op != Op.store and op.op != Op.load):
+            return True
+
+        return not set(op.getSlots()).issubset(remove)
+
+    for block in TealBlock.Iterate(start):
+        block.ops = list(filter(keep_op, block.ops))
+
+
+# Very dumb, overly eager dependency checking. A "dependency" is considered
+# any time the slot is loaded from in the entire control flow graph. This
+# can definitely be improved in the future.
+def _has_load_dependencies(
+    cur_block: TealBlock, start: TealBlock, slot: ScratchSlot, pos: int
+) -> bool:
+    for block in TealBlock.Iterate(start):
+        for i, op in enumerate(block.ops):
+            if block == cur_block and i == pos:
+                continue
+
+            if type(op) == TealOp and op.op == Op.load and slot in set(op.getSlots()):
+                return True
+
+    return False
+
+
+def _apply_slot_to_stack(
+    cur_block: TealBlock, start: TealBlock, skip_slots: Set[ScratchSlot]
+):
+    slots_to_remove = set()
+    # surprisingly, this slicing is totally safe - even if the list is empty.
+    for i, op in enumerate(cur_block.ops[:-1]):
+        if type(op) != TealOp or op.op != Op.store:
+            continue
+
+        if set(op.getSlots()).issubset(skip_slots):
+            continue
+
+        next_op = cur_block.ops[i + 1]
+        if type(next_op) != TealOp or next_op.op != Op.load:
+            continue
+
+        cur_slots, next_slots = op.getSlots(), next_op.getSlots()
+        if len(cur_slots) != 1 or len(next_slots) != 1:
+            raise TealInternalError(
+                "load/store op does not have exactly one slot argument"
+            )
+        if cur_slots[0] != next_slots[0]:
+            continue
+
+        if not _has_load_dependencies(cur_block, start, cur_slots[0], i + 1):
+            slots_to_remove.add(cur_slots[0])
+
+    _remove_extraneous_slot_access(start, slots_to_remove)
+
+
+def apply_global_optimizations(start: TealBlock, options: OptimizeOptions) -> TealBlock:
+    # limit number of iterations to length of teal program to avoid potential
+    # infinite loops.
+    for block in TealBlock.Iterate(start):
+        for _ in range(len(block.ops)):
+            prev_ops = block.ops.copy()
+            if options.scratch_slots:
+                _apply_slot_to_stack(block, start, options._skip_slots)
+
+            if prev_ops == block.ops:
+                break
+
+    return start