Skip to content

Commit 9dcc78e

Browse files
l46kokcopybara-github
l46kok
authored and
copybara-github
committed
Hoist loop-invariant subexpressions in SubexpressionOptimizer
PiperOrigin-RevId: 1005057100
1 parent b7096de commit 9dcc78e

8 files changed

Lines changed: 1863 additions & 598 deletions

File tree

‎optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java‎

Lines changed: 73 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,11 +26,13 @@
2626
import com.google.common.annotations.VisibleForTesting;
2727
import com.google.common.base.Strings;
2828
import com.google.common.base.Verify;
29+
import com.google.common.collect.HashMultiset;
2930
import com.google.common.collect.ImmutableList;
3031
import com.google.common.collect.ImmutableListMultimap;
3132
import com.google.common.collect.ImmutableMap;
3233
import com.google.common.collect.ImmutableSet;
3334
import com.google.common.collect.Multimaps;
35+
import com.google.common.collect.Multiset;
3436
import com.google.common.collect.Sets;
3537
import com.google.common.collect.Streams;
3638
import com.google.errorprone.annotations.CanIgnoreReturnValue;
@@ -414,15 +416,40 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
414416
navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).collect(toImmutableList());
415417
ImmutableListMultimap<Integer, CelNavigableMutableExpr> nodesByHeight =
416418
Multimaps.index(allNodes, CelNavigableMutableExpr::height);
419+
// Loop-invariant subexpressions are extracted even if they occur once: block slots are lazily
420+
// memoized, so they then run at most once and only if reached, instead of once per iteration.
421+
@SuppressWarnings("ValueClassIdentity") // Integer values are only unboxed, never compared.
422+
IdentityHashMap<CelMutableExpr, Integer> loopDepths = new IdentityHashMap<>();
423+
for (CelNavigableMutableExpr node : allNodes) {
424+
if (node.getKind().equals(Kind.COMPREHENSION)) {
425+
CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
426+
.allNodes()
427+
.forEach(step -> loopDepths.merge(step.expr(), 1, Integer::sum));
428+
}
429+
}
430+
// Only hoisting needs these, and computing them walks every node's subtree.
431+
Set<CelMutableExpr> eligibleExprs = Sets.newIdentityHashSet();
432+
Multiset<Integer> vnCounts = HashMultiset.create();
433+
if (!loopDepths.isEmpty()) {
434+
for (CelNavigableMutableExpr node : allNodes) {
435+
if (canEliminate(node, ineligibleExprs)) {
436+
eligibleExprs.add(node.expr());
437+
vnCounts.add(valueNumberer.vnOf(node.expr()));
438+
}
439+
}
440+
}
417441
Set<Integer> seenVns = new HashSet<>();
418442
for (CelNavigableMutableExpr node : allNodes) {
419-
ImmutableList<CelNavigableMutableExpr> sameHeightNodes = nodesByHeight.get(node.height());
420-
if (sameHeightNodes.size() < 2 || !canEliminate(node, ineligibleExprs)) {
443+
boolean inLoopStep = loopDepths.containsKey(node.expr());
444+
if ((!inLoopStep && nodesByHeight.get(node.height()).size() < 2)
445+
|| !canEliminate(node, ineligibleExprs)) {
421446
continue;
422447
}
423-
int vn = valueNumberer.vnOf(node.expr());
424-
if (!seenVns.add(vn)) {
425-
return sameHeightNodes.stream()
448+
CelNavigableMutableExpr target =
449+
inLoopStep ? hoistTarget(node, loopDepths, vnCounts, valueNumberer, eligibleExprs) : node;
450+
int vn = valueNumberer.vnOf(target.expr());
451+
if (!seenVns.add(vn) || inLoopStep) {
452+
return nodesByHeight.get(target.height()).stream()
426453
.filter(candidate -> canEliminate(candidate, ineligibleExprs))
427454
.map(CelNavigableMutableExpr::expr)
428455
.filter(expr -> valueNumberer.vnOf(expr) == vn)
@@ -433,6 +460,41 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
433460
return ImmutableList.of();
434461
}
435462

463+
// Slots are never revisited, so before hoisting `node` this first hoists any eligible descendant
464+
// that also occurs outside of it, or that is invariant to a loop nested inside it.
465+
private static CelNavigableMutableExpr hoistTarget(
466+
CelNavigableMutableExpr node,
467+
IdentityHashMap<CelMutableExpr, Integer> loopDepths,
468+
Multiset<Integer> vnCounts,
469+
ValueNumberer valueNumberer,
470+
Set<CelMutableExpr> eligibleExprs) {
471+
// A descendant occurring more often than `node` is duplicated within it or shared outside it.
472+
int copies = vnCounts.count(valueNumberer.vnOf(node.expr()));
473+
int depth = loopDepths.get(node.expr());
474+
return node.descendants(TraversalOrder.PRE_ORDER)
475+
.filter(
476+
d ->
477+
eligibleExprs.contains(d.expr())
478+
&& (vnCounts.count(valueNumberer.vnOf(d.expr())) > copies
479+
|| loopDepths.get(d.expr()) > depth))
480+
.findFirst()
481+
.map(d -> hoistTarget(d, loopDepths, vnCounts, valueNumberer, eligibleExprs))
482+
.orElse(node);
483+
}
484+
485+
// The planner accumulates `@ac + [elem]` steps (map, filter) in place, so `[elem]` is not
486+
// extracted.
487+
private static boolean isAccumulation(CelMutableExpr expr) {
488+
if (!expr.getKind().equals(Kind.CALL)
489+
|| !expr.call().function().equals(Operator.ADD.getFunction())) {
490+
return false;
491+
}
492+
CelMutableExpr accu = expr.call().args().get(0);
493+
return accu.getKind().equals(Kind.IDENT)
494+
&& accu.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)
495+
&& expr.call().args().get(1).getKind().equals(Kind.LIST);
496+
}
497+
436498
private boolean canEliminate(
437499
CelNavigableMutableExpr navigableExpr, Set<CelMutableExpr> ineligibleExprs) {
438500
return !navigableExpr.getKind().equals(Kind.CONSTANT)
@@ -488,6 +550,12 @@ private static Set<CelMutableExpr> getIneligibleExprsFromComprehensionBranches(
488550
.allNodes())
489551
.map(CelNavigableMutableExpr::expr)
490552
.forEach(ineligibleExprs::add));
553+
navAst
554+
.getRoot()
555+
.allNodes()
556+
.map(CelNavigableMutableExpr::expr)
557+
.filter(SubexpressionOptimizer::isAccumulation)
558+
.forEach(expr -> ineligibleExprs.add(expr.call().args().get(1)));
491559

492560
return ineligibleExprs;
493561
}

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerBaselineTest.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import static com.google.common.truth.Truth.assertThat;
1818
import static dev.cel.common.CelOverloadDecl.newGlobalOverload;
1919

20+
import com.google.common.collect.ImmutableList;
2021
import com.google.common.collect.ImmutableMap;
2122
import com.google.testing.junit.testparameterinjector.TestParameter;
2223
import com.google.testing.junit.testparameterinjector.TestParameterInjector;
@@ -95,6 +96,8 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
9596
TestAllTypes.newBuilder()
9697
.setSingleInt64(3L)
9798
.setSingleInt32(5)
99+
.addAllRepeatedInt64(ImmutableList.of(1L, 2L, 3L))
100+
.addRepeatedInt32(2)
98101
.setOneofType(
99102
NestedTestAllTypes.newBuilder()
100103
.setPayload(
@@ -412,6 +415,26 @@ private enum CseTestCase {
412415
"(msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 !="
413416
+ " msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 !="
414417
+ " msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)"),
418+
HOIST_NESTED_COMPREHENSION("msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))"),
419+
HOIST_INVARIANT_ARITHMETIC(
420+
"msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)"),
421+
HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED(
422+
"msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))"),
423+
HOIST_MAP_ACCUMULATION_PRESERVED("msg.repeated_int64.map(i, msg.single_int64 * 2)"),
424+
HOIST_LIST_LITERAL("msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
425+
HOIST_LIST_APPENDED_TO_ITER_VAR("[[1], [2]].map(l, l + [msg.single_int64])"),
426+
HOIST_INVARIANT_INSIDE_DUPLICATE(
427+
"msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) ||"
428+
+ " msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
429+
HOIST_INVARIANT_OF_NESTED_COMPREHENSION(
430+
"msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))"),
431+
HOIST_SHARED_TERMS_ONCE(
432+
"msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 +"
433+
+ " msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64]"
434+
+ ".size() > 0"),
435+
HOIST_SHARED_WITH_ACCUMULATED_LIST(
436+
"msg.repeated_int64.exists(i, i in [msg.single_int64]) &&"
437+
+ " msg.repeated_int64.map(i, msg.single_int64).size() > 0"),
415438
;
416439

417440
private final String source;

‎optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -347,6 +347,27 @@ public void iterationLimitReached_throws() throws Exception {
347347
assertThat(e).hasMessageThat().isEqualTo("Optimization failure: Max iteration count reached.");
348348
}
349349

350+
@Test
351+
@TestParameters(
352+
"{source: 'msg.repeated_int64.exists(i, i == (msg.single_int64 + msg.single_int32) *"
353+
+ " (msg.single_int32 + msg.single_int64))', expected: 'cel.@block([msg.single_int64 +"
354+
+ " msg.single_int32, @index0 * @index0], msg.repeated_int64.exists(@it:0:0, @it:0:0 =="
355+
+ " @index1))'}")
356+
@TestParameters(
357+
"{source: 'msg.repeated_int64.map(i, cel.bind(a, i, a + msg.single_int32 * 2))', expected:"
358+
+ " 'cel.@block([msg.single_int32 * 2], msg.repeated_int64.map(@it:1:0, cel.bind(a,"
359+
+ " @it:1:0, @ac:0:0 + @index0)))'}")
360+
public void cse_hoistLoopInvariants(String source, String expected) throws Exception {
361+
CelAbstractSyntaxTree ast = cel.compile(source).getAst();
362+
CelOptimizer celOptimizer =
363+
newCseOptimizer(
364+
SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());
365+
366+
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
367+
368+
assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(expected);
369+
}
370+
350371
@Test
351372
public void cse_comprehensionReferencingOuterBindVar_notExtractedOutOfBindScope()
352373
throws Exception {
@@ -757,12 +778,12 @@ public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
757778
.addAstOptimizers(SubexpressionOptimizer.getInstance())
758779
.build();
759780
CelAbstractSyntaxTree ast =
760-
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
781+
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))").getAst();
761782

762783
CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);
763784

764785
assertThat(CEL_UNPARSER.unparse(optimizedAst))
765-
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
786+
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))");
766787
assertThat(optimizedAst).isSameInstanceAs(ast);
767788
}
768789

@@ -797,8 +818,8 @@ public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exc
797818

798819
assertThat(CEL_UNPARSER.unparse(optimizedAst))
799820
.isEqualTo(
800-
"cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
801-
+ " @index0)");
821+
"cel.@block([has(x.a), [{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, @index0)], @index1"
822+
+ " == @index1)");
802823
}
803824

804825
@Test

‎optimizer/src/test/resources/constfold_before_subexpression_unparsed.baseline‎

Lines changed: 61 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -326,4 +326,64 @@ Test case: VALUE_NUMBERING_COMMUTATIVE_MULTIPLY_EQUALITY_LOGICAL_OR
326326
Source: (msg.oneof_type.payload.single_int32 * msg.single_int64 == 24 || msg.single_int32 != msg.oneof_type.payload.single_int64) && (msg.oneof_type.payload.single_int64 != msg.single_int32 || 24 == msg.single_int64 * msg.oneof_type.payload.single_int32)
327327
=====>
328328
Result: true
329-
Unparsed: cel.@block([msg.oneof_type.payload, @index0.single_int32 * msg.single_int64 == 24 || msg.single_int32 != @index0.single_int64], @index1 && @index1)
329+
Unparsed: cel.@block([msg.oneof_type.payload, @index0.single_int32 * msg.single_int64 == 24 || msg.single_int32 != @index0.single_int64], @index1 && @index1)
330+
331+
Test case: HOIST_NESTED_COMPREHENSION
332+
Source: msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))
333+
=====>
334+
Result: true
335+
Unparsed: cel.@block([msg.repeated_int32.exists(@it:0:0, @it:0:0 == 2)], msg.repeated_int64.all(@it:1:0, @index0))
336+
337+
Test case: HOIST_INVARIANT_ARITHMETIC
338+
Source: msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)
339+
=====>
340+
Result: false
341+
Unparsed: cel.@block([msg.single_int64 * 2 + msg.single_int32], msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index0))
342+
343+
Test case: HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED
344+
Source: msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))
345+
=====>
346+
Result: false
347+
Unparsed: cel.@block([1 / (msg.single_int32 - 5)], msg.repeated_nested_message.exists(@it:0:0, @it:0:0.bb == @index0))
348+
349+
Test case: HOIST_MAP_ACCUMULATION_PRESERVED
350+
Source: msg.repeated_int64.map(i, msg.single_int64 * 2)
351+
=====>
352+
Result: [6, 6, 6]
353+
Unparsed: cel.@block([msg.single_int64 * 2], msg.repeated_int64.map(@it:0:0, @index0))
354+
355+
Test case: HOIST_LIST_LITERAL
356+
Source: msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])
357+
=====>
358+
Result: true
359+
Unparsed: cel.@block([[msg.single_int64, msg.single_int32]], msg.repeated_int64.exists(@it:0:0, @it:0:0 in @index0))
360+
361+
Test case: HOIST_LIST_APPENDED_TO_ITER_VAR
362+
Source: [[1], [2]].map(l, l + [msg.single_int64])
363+
=====>
364+
Result: [[1, 3], [2, 3]]
365+
Unparsed: cel.@block([[msg.single_int64]], [[1], [2]].map(@it:0:0, @it:0:0 + @index0))
366+
367+
Test case: HOIST_INVARIANT_INSIDE_DUPLICATE
368+
Source: msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32]) || msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])
369+
=====>
370+
Result: true
371+
Unparsed: cel.@block([[msg.single_int64, msg.single_int32], msg.repeated_int64.exists(@it:0:0, @it:0:0 in @index0)], @index1 || @index1)
372+
373+
Test case: HOIST_INVARIANT_OF_NESTED_COMPREHENSION
374+
Source: msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))
375+
=====>
376+
Result: false
377+
Unparsed: cel.@block([msg.single_int64 * 2, msg.repeated_int32.exists(@it:0:0, @it:0:0 == @index0)], msg.repeated_int64.all(@it:1:0, @index1))
378+
379+
Test case: HOIST_SHARED_TERMS_ONCE
380+
Source: msg.repeated_int64.exists(i, i == msg.single_int64 + msg.single_int32 + msg.single_sint64) && [msg.single_int64, msg.single_int32, msg.single_sint64].size() > 0
381+
=====>
382+
Result: false
383+
Unparsed: cel.@block([msg.single_int64, msg.single_int32, msg.single_sint64, @index0 + @index1 + @index2], msg.repeated_int64.exists(@it:0:0, @it:0:0 == @index3) && [@index0, @index1, @index2].size() > 0)
384+
385+
Test case: HOIST_SHARED_WITH_ACCUMULATED_LIST
386+
Source: msg.repeated_int64.exists(i, i in [msg.single_int64]) && msg.repeated_int64.map(i, msg.single_int64).size() > 0
387+
=====>
388+
Result: true
389+
Unparsed: cel.@block([msg.single_int64, [@index0], msg.repeated_int64], @index2.exists(@it:0:0, @it:0:0 in @index1) && @index2.map(@it:0:1, @index0).size() > 0)

‎optimizer/src/test/resources/large_expressions_block_common_subexpr.baseline‎

Lines changed: 1 addition & 1 deletion
Large diffs are not rendered by default.

0 commit comments

Comments
 (0)