Skip to content

Commit 2e59bbf

Browse files
algoidurovicmichaeldiamantbarnjaminjasonpaulosedwardgaudio
authored
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]>
1 parent 18814d3 commit 2e59bbf

11 files changed

+852
-113
lines changed

docs/compiler_optimization.rst

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
.. _compiler_optimization:
2+
3+
Compiler Optimization
4+
========================
5+
**The optimizer is at an early stage and is disabled by default. Backwards compatability cannot be
6+
guaranteed at this point.**
7+
8+
The optimizer is a tool for improving performance and reducing resource consumption. In this context,
9+
the terms *performance* and *resource* can apply across multiple dimensions, including but not limited
10+
to: compiled code size, scratch slot usage, opcode cost, etc.
11+
12+
Optimizer Usage
13+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
14+
15+
The compiler determines which optimizations to apply based on the provided :any:`OptimizeOptions` object as
16+
shown in the code block below. The :any:`OptimizeOptions` constructor receives a set of keyword arguments
17+
representing flags corresponding to particular optimizations. If arguments are not provided to the
18+
constructor or no :any:`OptimizeOptions` object is passed to :any:`compileTeal` then the default behavior is
19+
that no optimizations are applied.
20+
21+
============================== ================================================================================ ===========================
22+
Optimization Flag Description Default
23+
============================== ================================================================================ ===========================
24+
:code:`scratch_slots` A boolean describing whether or not scratch slot optimization should be applied. :code:`False`
25+
============================== ================================================================================ ===========================
26+
27+
.. code-block:: python
28+
29+
optimize_options = OptimizeOptions(scratch_slots=True)
30+
compileTeal(approval_program(), mode=Mode.Application, version=4, optimize=optimize_options)

docs/index.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ PyTeal **hasn't been security audited**. Use it at your own risk.
4343
state
4444
assets
4545
versions
46+
compiler_optimization
4647

4748
.. toctree::
4849
:maxdepth: 3

pyteal/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
DEFAULT_TEAL_VERSION,
99
CompileOptions,
1010
compileTeal,
11+
OptimizeOptions,
1112
)
1213
from .types import TealType
1314
from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
@@ -23,6 +24,7 @@
2324
"DEFAULT_TEAL_VERSION",
2425
"CompileOptions",
2526
"compileTeal",
27+
"OptimizeOptions",
2628
"TealType",
2729
"TealInternalError",
2830
"TealTypeError",

pyteal/__init__.pyi

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ from .compiler import (
1111
DEFAULT_TEAL_VERSION,
1212
CompileOptions,
1313
compileTeal,
14+
OptimizeOptions,
1415
)
1516
from .types import TealType
1617
from .errors import TealInternalError, TealTypeError, TealInputError, TealCompileError
@@ -122,6 +123,7 @@ __all__ = [
122123
"Not",
123124
"OnComplete",
124125
"Op",
126+
"OptimizeOptions",
125127
"Or",
126128
"Pop",
127129
"Reject",

pyteal/compiler/__init__.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,13 @@
66
compileTeal,
77
)
88

9+
from .optimizer import OptimizeOptions
10+
911
__all__ = [
1012
"MAX_TEAL_VERSION",
1113
"MIN_TEAL_VERSION",
1214
"DEFAULT_TEAL_VERSION",
1315
"CompileOptions",
1416
"compileTeal",
17+
"OptimizeOptions",
1518
]

pyteal/compiler/compiler.py

Lines changed: 52 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
from typing import List, Tuple, Set, Dict, Optional, cast
22

3+
from .optimizer import OptimizeOptions, apply_global_optimizations
4+
35
from ..types import TealType
46
from ..ast import (
57
Expr,
@@ -13,7 +15,10 @@
1315

1416
from .sort import sortBlocks
1517
from .flatten import flattenBlocks, flattenSubroutines
16-
from .scratchslots import assignScratchSlotsToSubroutines
18+
from .scratchslots import (
19+
assignScratchSlotsToSubroutines,
20+
collect_unoptimized_slots,
21+
)
1722
from .subroutines import (
1823
spillLocalSlotsDuringRecursion,
1924
resolveSubroutines,
@@ -31,9 +36,11 @@ def __init__(
3136
*,
3237
mode: Mode = Mode.Signature,
3338
version: int = DEFAULT_TEAL_VERSION,
39+
optimize: OptimizeOptions = None,
3440
) -> None:
3541
self.mode = mode
3642
self.version = version
43+
self.optimize = optimize if optimize is not None else OptimizeOptions()
3744

3845
self.currentSubroutine: Optional[SubroutineDefinition] = None
3946

@@ -109,9 +116,9 @@ def verifyOpsForMode(teal: List[TealComponent], mode: Mode):
109116
def compileSubroutine(
110117
ast: Expr,
111118
options: CompileOptions,
112-
subroutineMapping: Dict[Optional[SubroutineDefinition], List[TealComponent]],
113119
subroutineGraph: Dict[SubroutineDefinition, Set[SubroutineDefinition]],
114-
subroutineBlocks: Dict[Optional[SubroutineDefinition], TealBlock],
120+
subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
121+
subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
115122
) -> None:
116123
currentSubroutine = (
117124
cast(SubroutineDeclaration, ast).subroutine
@@ -139,8 +146,8 @@ def compileSubroutine(
139146
verifyOpsForVersion(teal, options.version)
140147
verifyOpsForMode(teal, options.mode)
141148

142-
subroutineMapping[currentSubroutine] = teal
143-
subroutineBlocks[currentSubroutine] = start
149+
subroutine_start_blocks[currentSubroutine] = start
150+
subroutine_end_blocks[currentSubroutine] = end
144151

145152
referencedSubroutines: Set[SubroutineDefinition] = set()
146153
for stmt in teal:
@@ -150,23 +157,38 @@ def compileSubroutine(
150157
if currentSubroutine is not None:
151158
subroutineGraph[currentSubroutine] = referencedSubroutines
152159

153-
newSubroutines = referencedSubroutines - subroutineMapping.keys()
160+
newSubroutines = referencedSubroutines - subroutine_start_blocks.keys()
154161
for subroutine in sorted(newSubroutines, key=lambda subroutine: subroutine.id):
155162
compileSubroutine(
156163
subroutine.getDeclaration(),
157164
options,
158-
subroutineMapping,
159165
subroutineGraph,
160-
subroutineBlocks,
166+
subroutine_start_blocks,
167+
subroutine_end_blocks,
161168
)
162169

163170

171+
def sort_subroutine_blocks(
172+
subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
173+
subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock],
174+
) -> Dict[Optional[SubroutineDefinition], List[TealComponent]]:
175+
subroutine_mapping: Dict[
176+
Optional[SubroutineDefinition], List[TealComponent]
177+
] = dict()
178+
for subroutine, start in subroutine_start_blocks.items():
179+
order = sortBlocks(start, subroutine_end_blocks[subroutine])
180+
subroutine_mapping[subroutine] = flattenBlocks(order)
181+
182+
return subroutine_mapping
183+
184+
164185
def compileTeal(
165186
ast: Expr,
166187
mode: Mode,
167188
*,
168189
version: int = DEFAULT_TEAL_VERSION,
169190
assembleConstants: bool = False,
191+
optimize: OptimizeOptions = None,
170192
) -> str:
171193
"""Compile a PyTeal expression into TEAL assembly.
172194
@@ -181,6 +203,7 @@ def compileTeal(
181203
constants will be assembled in the most space-efficient way, so enabling this may reduce
182204
the compiled program's size. Enabling this option requires a minimum TEAL version of 3.
183205
Defaults to false.
206+
optimize (optional): OptimizeOptions that determine which optimizations will be applied.
184207
185208
Returns:
186209
A TEAL assembly program compiled from the input expression.
@@ -199,20 +222,32 @@ def compileTeal(
199222
)
200223
)
201224

202-
options = CompileOptions(mode=mode, version=version)
225+
options = CompileOptions(mode=mode, version=version, optimize=optimize)
203226

204-
subroutineMapping: Dict[
205-
Optional[SubroutineDefinition], List[TealComponent]
206-
] = dict()
207227
subroutineGraph: Dict[SubroutineDefinition, Set[SubroutineDefinition]] = dict()
208-
subroutineBlocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
228+
subroutine_start_blocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
229+
subroutine_end_blocks: Dict[Optional[SubroutineDefinition], TealBlock] = dict()
209230
compileSubroutine(
210-
ast, options, subroutineMapping, subroutineGraph, subroutineBlocks
231+
ast, options, subroutineGraph, subroutine_start_blocks, subroutine_end_blocks
211232
)
212233

213-
localSlotAssignments = assignScratchSlotsToSubroutines(
214-
subroutineMapping, subroutineBlocks
215-
)
234+
# note: optimizations are off by default, in which case, apply_global_optimizations
235+
# won't make any changes. Because the optimizer is invoked on a subroutine's
236+
# control flow graph, the optimizer requires context across block boundaries. This
237+
# is necessary for the dependency checking of local slots. Global slots, slots
238+
# used by DynamicScratchVar, and reserved slots are not optimized.
239+
if options.optimize.scratch_slots:
240+
options.optimize._skip_slots = collect_unoptimized_slots(
241+
subroutine_start_blocks
242+
)
243+
for start in subroutine_start_blocks.values():
244+
apply_global_optimizations(start, options.optimize)
245+
246+
localSlotAssignments = assignScratchSlotsToSubroutines(subroutine_start_blocks)
247+
248+
subroutineMapping: Dict[
249+
Optional[SubroutineDefinition], List[TealComponent]
250+
] = sort_subroutine_blocks(subroutine_start_blocks, subroutine_end_blocks)
216251

217252
spillLocalSlotsDuringRecursion(
218253
version, subroutineMapping, subroutineGraph, localSlotAssignments

pyteal/compiler/optimizer/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from .optimizer import OptimizeOptions, apply_global_optimizations
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
from typing import Set
2+
from ...ast import ScratchSlot
3+
from ...ir import TealBlock, TealOp, Op
4+
from ...errors import TealInternalError
5+
6+
7+
class OptimizeOptions:
8+
"""An object which specifies the optimizations to be performed and relevant context.
9+
10+
_skip_slots: the slots that should be skipped during optimization. At the moment this includes:
11+
1. reserved slots because they may have dependencies outside
12+
the current application. For example, the 'gloads' opcode can
13+
access the slots of other applications in the tx group.
14+
2. global slots because they're outside the scope of global
15+
optimizations, which only apply to the control flow graph of
16+
a single subroutine.
17+
3. slots used with dynamic scratch vars. These slots use
18+
indirection by means of the 'stores' opcode and dependencies
19+
can only be determined at runtime.
20+
21+
Args:
22+
23+
scratch_slots (optional): cancel contiguous store/load operations
24+
that have no load dependencies elsewhere.
25+
"""
26+
27+
def __init__(self, *, scratch_slots: bool = False):
28+
self.scratch_slots = scratch_slots
29+
self._skip_slots: Set[ScratchSlot] = set()
30+
31+
32+
def _remove_extraneous_slot_access(start: TealBlock, remove: Set[ScratchSlot]):
33+
def keep_op(op: TealOp) -> bool:
34+
if type(op) != TealOp or (op.op != Op.store and op.op != Op.load):
35+
return True
36+
37+
return not set(op.getSlots()).issubset(remove)
38+
39+
for block in TealBlock.Iterate(start):
40+
block.ops = list(filter(keep_op, block.ops))
41+
42+
43+
# Very dumb, overly eager dependency checking. A "dependency" is considered
44+
# any time the slot is loaded from in the entire control flow graph. This
45+
# can definitely be improved in the future.
46+
def _has_load_dependencies(
47+
cur_block: TealBlock, start: TealBlock, slot: ScratchSlot, pos: int
48+
) -> bool:
49+
for block in TealBlock.Iterate(start):
50+
for i, op in enumerate(block.ops):
51+
if block == cur_block and i == pos:
52+
continue
53+
54+
if type(op) == TealOp and op.op == Op.load and slot in set(op.getSlots()):
55+
return True
56+
57+
return False
58+
59+
60+
def _apply_slot_to_stack(
61+
cur_block: TealBlock, start: TealBlock, skip_slots: Set[ScratchSlot]
62+
):
63+
slots_to_remove = set()
64+
# surprisingly, this slicing is totally safe - even if the list is empty.
65+
for i, op in enumerate(cur_block.ops[:-1]):
66+
if type(op) != TealOp or op.op != Op.store:
67+
continue
68+
69+
if set(op.getSlots()).issubset(skip_slots):
70+
continue
71+
72+
next_op = cur_block.ops[i + 1]
73+
if type(next_op) != TealOp or next_op.op != Op.load:
74+
continue
75+
76+
cur_slots, next_slots = op.getSlots(), next_op.getSlots()
77+
if len(cur_slots) != 1 or len(next_slots) != 1:
78+
raise TealInternalError(
79+
"load/store op does not have exactly one slot argument"
80+
)
81+
if cur_slots[0] != next_slots[0]:
82+
continue
83+
84+
if not _has_load_dependencies(cur_block, start, cur_slots[0], i + 1):
85+
slots_to_remove.add(cur_slots[0])
86+
87+
_remove_extraneous_slot_access(start, slots_to_remove)
88+
89+
90+
def apply_global_optimizations(start: TealBlock, options: OptimizeOptions) -> TealBlock:
91+
# limit number of iterations to length of teal program to avoid potential
92+
# infinite loops.
93+
for block in TealBlock.Iterate(start):
94+
for _ in range(len(block.ops)):
95+
prev_ops = block.ops.copy()
96+
if options.scratch_slots:
97+
_apply_slot_to_stack(block, start, options._skip_slots)
98+
99+
if prev_ops == block.ops:
100+
break
101+
102+
return start

0 commit comments

Comments
 (0)