Tools : PyTeal optimization utility (#247)

* Optimization added for repeated int constants under 2**7 w/ tests

* fixed type problem and formatted

* Expanded test and added comment for clarification

* implement optimization utility with simple slot store/load canceling

* minor refactor

* reformat code

* Update pyteal/compiler/optimizer/optimizer.py

Co-authored-by: Michael Diamant <[email protected]>

* Update pyteal/compiler/optimizer/optimizer.py

Co-authored-by: Michael Diamant <[email protected]>

* Adding exponentiation to arithmatic ops docs (#134)

Add missing exponentiation operation in document

* updating to use new syntax for seq (#135)

* updating to use new syntax for seq

* rewording

* Make pylance recognize wildcard imports (#133)

* adding exports directly to top level __all__

* apply black formatter

* adding initial generate script

* fmt

* rm all from all

* adding check to travis

* reading in original __init__ and using its imports, dont write to filesystem if --check is passed

* make messages more profesh

* fix flags after black formatted them

* y

* flippin black formatter

* help text fix

* asdfasdf

* Include pyi files in build (#137)

* Revert "Optimization for constant assembly (#128)"

This reverts commit 5636ccd.

* Revert "String optimization and addition of Suffix() (#126)"

This reverts commit 7cb7b9a.

* Update to v0.9.1 (#138)

* Revert "Revert "String optimization and addition of Suffix() (#126)""

This reverts commit 564e602.

* Revert "Revert "Optimization for constant assembly (#128)""

This reverts commit cc405a5.

* Update examples.rst (#140)

* Fix type for App.globalGetEx in docs (#142)

* up max teal version (#146)

* up max teal version

* make test fail if its greater than version defined as MAX_TEAL_VERSION

* Fmt

* hardcode to 7

* Add version 6 test

* Formatting subroutines with name and newline (#148)

* using the subroutine name for the label

* adding newline before label declaration, fix tests to account for newline

* remove commented name, fix test

* only add newline for subroutines with comment

* naming with suffix

* adding test for invalid name

* Call type_of() in require_type() for better exception messages (#151)

* call type_of in require_type to catch exceptions

* fix formatting for types.py and types_test.py

* `method` pseudo-op support for ABI methods (#153)

- Add support for `method` pseudo-opcode in PyTeal.
- Add `name` field in `subroutine` to override __name__ from function implementation, for readability in generated code.

* Print diff of `__init__.pyi` (#166)

* Print diff of __init__.pyi

* Format

* Undo travis change

* C2C Feature Support (#149)

- `itxn_next` implementation / test
- `itxn_field` support for array field setting
- `gitxn / gitxna` implementation / test
- `gloadss` implementation / test

* Add BytesSqrt (#163)

* Add BytesSqrt

* Update pyteal/ast/unaryexpr_test.py

Co-authored-by: Jason Paulos <[email protected]>

Co-authored-by: Jason Paulos <[email protected]>

* adding new globals from teal6 (#168)

* adding new globals from teal6

* fmt

* Acct params get (#165)

* Adding account param getter

* Add to init

* fix op names and type

* adding tests

* allow bytes to be passed

* tweak docs, add require check for any

* Change Subroutine Wrapped Callable to a class with call method (#171)

Allows for more information (name, return type, has return) about the subroutine extractable from wrapped fnImpl by subroutine

* Subroutine Type Annotations (#182)

This PR requires that any type annotation of a Subroutine parameter or return value be of type `Expr`. Missing annotations are assumed to be `Expr`'s as well. In a follow up PR #183 this restriction will be loosened.

* fix docs referencing what apps should eval to (#191)

* Move from Travis to Github Actions (#190)

* MultiValue expression implemented to support opcodes that return multiple values (#196)

* Optimization added for repeated int constants under 2**7 w/ tests

* fixed type problem and formatted

* Expanded test and added comment for clarification

* add multivalue expr and change maybevalue to derive from multivalue

* updated tests and formatting

* reorder output slots to reflect stack ordering

* add additional assertion in MaybeValue test to enforce slot ordering

* Support TEAL 6 txn fields LastLog, StateProofPK and opcodes divw, itxnas, gitxnas (#174)

* adding new teal6 ops, no pyteal expressions defined for them yet

* Add opcode support for divw

* Add opcode support for divw (#192)

* Add opcode support for itxnas and gitxnas (#193)

* Add opcode support for itxnas and gitxnas

* Update stale reference to inner transaction limit

* Fix allowed types for GitxnaExpr txnIndex

* Remove obsolete logic for handling GitxnaExpr.teal construction

* Remove unnecessary cast and fix gitxna runtime type checking

* Move type validation to constructors for gtxn and gitxn variants

* Add missed tests from prior commit

* Fix duplicate test case

* Move index validation from subclasses to TxnaExpr

* Inline validation functions per PR feedback

* Remove unused imports

* Refactor to isinstance tupled check

* Remove TEAL v1 min version test per PR feedback

* Fix constructor type checking for GtxnExpr

* Refactor to remove duplicate type check function

* Update last_log docstring

Co-authored-by: Jason Paulos <[email protected]>

* Expose state_proof_pk txn field

* Update transaction field docs to reflect TEAL v6

* Update transaction field docs to reflect TEAL v6

Co-authored-by: michaeldiamant <[email protected]>
Co-authored-by: Jason Paulos <[email protected]>

* Fixed typo (#202)

* Add Github action to generate docset (#201)

* Add build docset step

* non-slim container

* Update docs to group transaction field tables like go-algorand (#204)

* Update accessing_transaction_field.rst to fix typo (#207)

* Add docs README to explain docs/ testing procedure (#205)

* v0.10.0 (#206)

* Update to v0.10.0

* Add latest commits to changelog

* fixing github actions to run on tags (#208)

* Update build.yml

* Update build.yml

* Fix typos in docstrings and error messages (#211)

* Test on Python 3.10 (#212)

* Update versions.rst (#210)

* Update versions.rst

content of [https://github.com/algorand/pyteal/releases] is not shown in [https://pyteal.readthedocs.io/en/latest/versions.html]

* Update docs/versions.rst

Co-authored-by: Jason Paulos <[email protected]>

Co-authored-by: Jason Paulos <[email protected]>

* Pass-by-Ref / Dynamic Scratch Variables via the `loads` and `stores` opcodes (#198)

* Pass-by-Reference Semantics
* Use a Dynamic ScratchVar to "iterate" over other ScratchVar's
* Another approach for E2E testing

* Fix build script invocation (#223)

* Bring #225 to master (#227)

* Ignore tests generating TEAL file outputs used for expected comparisons (#228)

* Fix typo in CONTRIBUTING.md (#229)

* Fix subroutine mutual recursion with different argument counts bug (#234)

* Fix mutual recursion bug

* Remove usage of set.pop

* Revert "Pass-by-Ref / Dynamic Scratch Variables via the `loads` and `stores` opcodes (#198)"

This reverts commit cf95165.

* v0.10.1 (#237)

* Revert "Revert "Pass-by-Ref / Dynamic Scratch Variables via the `loads` and `stores` opcodes (#198)""

This reverts commit 51ec8c9.

* Update user guide docs to reflect addition of DynamicScratchVar (#226)

* Update CONTRIBUTING.md on PEP 8 naming conventions policy (#241)

* implement optimization utility with simple slot store/load canceling

* minor refactor

* reformat code

* correct import format to match convention

* slot optimization awareness of reserved ids added

* fix typo

* remove dataclass usage

* slight reorg of compiler process in order to perform optimization on cfg

* clean up imports

* updated documentation and reformatted with new version of black

* remove unused imports and comments

* reformatting

* add additional optimizer unit tests

* improve testing and slight refactoring

* more renaming

* documentation and import changes

* fixed typos in docs

Co-authored-by: Michael Diamant <[email protected]>
Co-authored-by: Ben Guidarelli <[email protected]>
Co-authored-by: Jason Paulos <[email protected]>
Co-authored-by: Edward D Gaudio <[email protected]>
Co-authored-by: Joe Polny <[email protected]>
Co-authored-by: Hang Su <[email protected]>
Co-authored-by: Łukasz Ptak <[email protected]>
Co-authored-by: Zeph Grunschlag <[email protected]>
Co-authored-by: Jack <[email protected]>
Co-authored-by: Glory Agatevure <[email protected]>
Co-authored-by: Adriano Di Luzio <[email protected]>
Co-authored-by: PabloLION <[email protected]>

Author: 13 people

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