Tools : PyTeal optimization utility

pyteal/compiler/compiler.py

@@ -1,5 +1,8 @@
 from typing import List, Tuple, Set, Dict, Optional, cast
 
+from pyteal.compiler.optimizer import OptimizeOptions
+from .optimizer import optimizer
+
 from ..types import TealType
 from ..ast import (
     Expr,
@@ -13,7 +16,10 @@
 
 from .sort import sortBlocks
 from .flatten import flattenBlocks, flattenSubroutines
-from .scratchslots import assignScratchSlotsToSubroutines
+from .scratchslots import (
+    assignScratchSlotsToSubroutines,
+    collectUnoptimizedSlotIDs,
+)
 from .subroutines import (
     spillLocalSlotsDuringRecursion,
     resolveSubroutines,
@@ -31,9 +37,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 +117,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],
+    subroutineEndBlocks: Dict[Optional[SubroutineDefinition], TealBlock],
 ) -> None:
     currentSubroutine = (
         cast(SubroutineDeclaration, ast).subroutine
@@ -139,8 +147,8 @@ def compileSubroutine(
     verifyOpsForVersion(teal, options.version)
     verifyOpsForMode(teal, options.mode)
 
-    subroutineMapping[currentSubroutine] = teal
     subroutineBlocks[currentSubroutine] = start
+    subroutineEndBlocks[currentSubroutine] = end
 
     referencedSubroutines: Set[SubroutineDefinition] = set()
     for stmt in teal:
@@ -150,23 +158,38 @@ def compileSubroutine(
     if currentSubroutine is not None:
         subroutineGraph[currentSubroutine] = referencedSubroutines
 
-    newSubroutines = referencedSubroutines - subroutineMapping.keys()
+    newSubroutines = referencedSubroutines - subroutineBlocks.keys()
     for subroutine in sorted(newSubroutines, key=lambda subroutine: subroutine.id):
         compileSubroutine(
             subroutine.getDeclaration(),
             options,
-            subroutineMapping,
             subroutineGraph,
             subroutineBlocks,
+            subroutineEndBlocks,
         )
 
 
+def createTealMapping(
+    subroutineBlocks: Dict[Optional[SubroutineDefinition], TealBlock],
+    subroutineEndBlocks: Dict[Optional[SubroutineDefinition], TealBlock],
+) -> Dict[Optional[SubroutineDefinition], List[TealComponent]]:
+    subroutineMapping: Dict[
+        Optional[SubroutineDefinition], List[TealComponent]
+    ] = dict()
+    for subroutine, start in subroutineBlocks.items():
+        order = sortBlocks(start, subroutineEndBlocks[subroutine])
+        subroutineMapping[subroutine] = flattenBlocks(order)
+
+    return subroutineMapping
+
+
 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 +204,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 optimization will be applied.
 
     Returns:
         A TEAL assembly program compiled from the input expression.
@@ -199,20 +223,31 @@ 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()
+    subroutineEndBlocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
     compileSubroutine(
-        ast, options, subroutineMapping, subroutineGraph, subroutineBlocks
+        ast, options, subroutineGraph, subroutineBlocks, subroutineEndBlocks
     )
 
-    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_ids = collectUnoptimizedSlotIDs(subroutineBlocks)
+        print(str(options.optimize.skip_ids))
+        for start in subroutineBlocks.values():
+            optimizer.apply_global_optimizations(start, options.optimize)
+
+    localSlotAssignments = assignScratchSlotsToSubroutines(subroutineBlocks)
+
+    subroutineMapping: Dict[
+        Optional[SubroutineDefinition], List[TealComponent]
+    ] = createTealMapping(subroutineBlocks, subroutineEndBlocks)
 
     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,89 @@
+from typing import cast, Set
+from pyteal.ir.tealblock import TealBlock
+
+from pyteal.ir.tealop import TealOp
+from pyteal.ir.ops import Op
+
+
+class OptimizeOptions:
+    """An object which specifies the optimizations to be performed and relevant context.
+    Args:
+
+        scratch_slots (optional): cancel contiguous store/load operations
+            that have no load dependencies elsewhere.
+        skip_ids (optional): the slot ids 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."""
+
+    def __init__(self, *, scratch_slots: bool = False, skip_ids: Set[int] = None):
+        self.scratch_slots = scratch_slots
+        self.skip_ids: Set[int] = skip_ids if skip_ids is not None else set()
+
+
+def _remove_extraneous_slot_access(start: TealBlock, remove: Set[int]):
+    def keep_op(op: TealOp):
+        if type(op) != TealOp or (op.op != Op.store and op.op != Op.load):
+            return True
+
+        return cast(int, op.args[0]) not in remove
+
+    for block in TealBlock.Iterate(start):
+        block.ops = list(filter(keep_op, block.ops))
+
+
+def _has_load_dependencies(cur_block: TealBlock, start: TealBlock, slot: int, pos: int):
+    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 op.args[0] == slot:
+                return True
+
+    return False
+
+
+def _apply_slot_to_stack(cur_block: TealBlock, start: TealBlock, skip_ids: Set[int]):
+    slots_to_remove = set()
+    for i, op in enumerate(cur_block.ops[:-1]):
+        if type(op) != TealOp or op.op != Op.store:
+            continue
+
+        # do not optimize away reserved and global slots
+        if op.getSlots()[0].id in skip_ids:
+            continue
+
+        next_op = cur_block.ops[i + 1]
+        if type(next_op) != TealOp or next_op.op != Op.load:
+            continue
+
+        if op.args[0] != next_op.args[0]:
+            continue
+
+        if not _has_load_dependencies(cur_block, start, cast(int, op.args[0]), i + 1):
+            slots_to_remove.add(cast(int, op.args[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):
+        prev_ops = block.ops.copy()
+        for _ in range(len(block.ops)):
+            if options.scratch_slots:
+                _apply_slot_to_stack(block, start, options.skip_ids)
+
+            if prev_ops == block.ops:
+                break
+
+    return start