| 926 | |
| 927 | |
| 928 | def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr: |
| 929 | result = expr |
| 930 | if stage == "logical": |
| 931 | return result |
| 932 | |
| 933 | # Simplify |
| 934 | expr = result.simplify() |
| 935 | if stage == "simplified-logical": |
| 936 | return expr |
| 937 | |
| 938 | # Manipulate Expression to make it more efficient |
| 939 | expr = expr.rewrite(kind="tune", rewritten={}) |
| 940 | if stage == "tuned-logical": |
| 941 | return expr |
| 942 | |
| 943 | # Lower |
| 944 | expr = expr.lower_completely() |
| 945 | if stage == "physical": |
| 946 | return expr |
| 947 | |
| 948 | # Simplify again |
| 949 | expr = expr.simplify() |
| 950 | if stage == "simplified-physical": |
| 951 | return expr |
| 952 | |
| 953 | # Final graph-specific optimizations |
| 954 | expr = expr.fuse() |
| 955 | if stage == "fused": |
| 956 | return expr |
| 957 | |
| 958 | raise ValueError(f"Stage {stage!r} not supported.") |
| 959 | |
| 960 | |
| 961 | class LLGExpr(Expr): |