Skip to content

Extract shared parts of duplicates first in SubexpressionOptimizer #1406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Extract shared parts of duplicates first in SubexpressionOptimizer #1406

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
161 changes: 108 additions & 53 deletions optimizer/src/main/java/dev/cel/optimizer/optimizers/SubexpressionOptimizer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,11 +26,11 @@
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Strings;
import com.google.common.base.Verify;
import com.google.common.collect.HashMultiset;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableListMultimap;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Multimaps;
import com.google.common.collect.Multiset;
import com.google.common.collect.Sets;
import com.google.common.collect.Streams;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
Expand Down Expand Up @@ -85,7 +85,6 @@
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;
import java.util.stream.Stream;

/**
Expand Down Expand Up @@ -412,55 +411,118 @@ private ImmutableList<CelMutableExpr> getCseCandidates(
Set<CelMutableExpr> ineligibleExprs = getIneligibleExprsFromComprehensionBranches(navAst);
ImmutableList<CelNavigableMutableExpr> allNodes =
navAst.getRoot().allNodes(TraversalOrder.PRE_ORDER).collect(toImmutableList());
ImmutableListMultimap<Integer, CelNavigableMutableExpr> nodesByHeight =
Multimaps.index(allNodes, CelNavigableMutableExpr::height);
Set<Integer> seenVns = new HashSet<>();
// Loop-invariant subexpressions are extracted even if they occur once: block slots are lazily
// memoized, so they then run at most once and only if reached, instead of once per iteration.
@SuppressWarnings("ValueClassIdentity") // Integer values are only unboxed, never compared.
IdentityHashMap<CelMutableExpr, Integer> loopDepths = new IdentityHashMap<>();
for (CelNavigableMutableExpr node : allNodes) {
ImmutableList<CelNavigableMutableExpr> sameHeightNodes = nodesByHeight.get(node.height());
if (sameHeightNodes.size() < 2 || !canEliminate(node, ineligibleExprs)) {
continue;
if (node.getKind().equals(Kind.COMPREHENSION)) {
CelNavigableMutableExpr.fromExpr(node.expr().comprehension().loopStep())
.allNodes()
.forEach(step -> loopDepths.merge(step.expr(), 1, Integer::sum));
}
int vn = valueNumberer.vnOf(node.expr());
if (!seenVns.add(vn)) {
return sameHeightNodes.stream()
.filter(candidate -> canEliminate(candidate, ineligibleExprs))
}
Set<CelMutableExpr> eligibleExprs = getEligibleExprs(allNodes, ineligibleExprs);
Multiset<Integer> vnCounts = HashMultiset.create();
eligibleExprs.forEach(expr -> vnCounts.add(valueNumberer.vnOf(expr)));
Set<Integer> seenVns = new HashSet<>();
for (CelNavigableMutableExpr node : allNodes) {
CelMutableExpr expr = node.expr();
if (eligibleExprs.contains(expr)
&& (loopDepths.containsKey(expr) || !seenVns.add(valueNumberer.vnOf(expr)))) {
int vn =
valueNumberer.vnOf(
extractionTarget(node, loopDepths, vnCounts, valueNumberer, eligibleExprs).expr());
return allNodes.stream()
.map(CelNavigableMutableExpr::expr)
.filter(expr -> valueNumberer.vnOf(expr) == vn)
.filter(e -> eligibleExprs.contains(e) && valueNumberer.vnOf(e) == vn)
.collect(toImmutableList());
}
}

return ImmutableList.of();
}

private boolean canEliminate(
CelNavigableMutableExpr navigableExpr, Set<CelMutableExpr> ineligibleExprs) {
return !navigableExpr.getKind().equals(Kind.CONSTANT)
&& !navigableExpr.getKind().equals(Kind.IDENT)
// Exclude empty lists (cel.bind sets this for iterRange).
&& !(navigableExpr.getKind().equals(Kind.LIST)
&& navigableExpr.expr().list().elements().isEmpty())
&& !ineligibleExprs.contains(navigableExpr.expr())
&& containsEliminableFunctionOnly(navigableExpr)
&& containsValidComprehensionIdents(navigableExpr);
// Slots are never revisited, so before extracting `node` this first extracts any eligible
// descendant that is duplicated within or also occurs outside of it, or (for a loop-step node)
// that is invariant to a loop nested inside it.
private static CelNavigableMutableExpr extractionTarget(
CelNavigableMutableExpr node,
IdentityHashMap<CelMutableExpr, Integer> loopDepths,
Multiset<Integer> vnCounts,
ValueNumberer valueNumberer,
Set<CelMutableExpr> eligibleExprs) {
// A descendant occurring more often than `node` is duplicated within it or shared outside it.
int copies = vnCounts.count(valueNumberer.vnOf(node.expr()));
Integer depth = loopDepths.get(node.expr());
return node.descendants(TraversalOrder.PRE_ORDER)
.filter(
d ->
eligibleExprs.contains(d.expr())
&& (vnCounts.count(valueNumberer.vnOf(d.expr())) > copies
|| (depth != null && loopDepths.get(d.expr()) > depth)))
.findFirst()
.map(d -> extractionTarget(d, loopDepths, vnCounts, valueNumberer, eligibleExprs))
.orElse(node);
}

private static boolean containsValidComprehensionIdents(CelNavigableMutableExpr navExpr) {
// For nested comprehensions of form [1].exists(x, [2].exists(y, x == y)), the inner
// comprehension [2].exists(y, x == y) should not be extracted out into a block index, as it
// causes issues with scoping. Non-comprehension subexpressions cannot reference any mangled
// comprehension variables.
Predicate<String> isInvalidIdent =
navExpr.getKind().equals(Kind.COMPREHENSION)
? CelNavigableExprUtil.getEnclosingComprehensionVariables(navExpr)::contains
: name ->
name.startsWith(MANGLED_COMPREHENSION_ITER_VAR_PREFIX)
|| name.startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX);
return navExpr
.descendants()
.filter(node -> node.getKind().equals(Kind.IDENT))
.map(node -> node.expr().ident().name())
.noneMatch(isInvalidIdent);
// The planner accumulates `@ac + [elem]` steps (map, filter) in place, so `[elem]` is not
// extracted.
private static boolean isAccumulation(CelMutableExpr expr) {
if (!expr.getKind().equals(Kind.CALL)
|| !expr.call().function().equals(Operator.ADD.getFunction())) {
return false;
}
CelMutableExpr accu = expr.call().args().get(0);
return accu.getKind().equals(Kind.IDENT)
&& accu.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)
&& expr.call().args().get(1).getKind().equals(Kind.LIST);
}

// Computed bottom-up, as checking every node's subtree separately is quadratic. Only
// comprehensions still walk their own subtree.
private Set<CelMutableExpr> getEligibleExprs(
ImmutableList<CelNavigableMutableExpr> preOrderNodes, Set<CelMutableExpr> ineligibleExprs) {
// Subtrees containing an ineliminable call, or any mangled comprehension variable.
Set<CelMutableExpr> impureExprs = Sets.newIdentityHashSet();
Set<CelMutableExpr> mangledExprs = Sets.newIdentityHashSet();
Set<CelMutableExpr> eligibleExprs = Sets.newIdentityHashSet();
// Reverse pre-order visits children before their parents.
for (CelNavigableMutableExpr node : preOrderNodes.reverse()) {
CelMutableExpr expr = node.expr();
Kind kind = node.getKind();
if ((kind.equals(Kind.CALL) && !cseEliminableFunctions.contains(expr.call().function()))
|| node.children().anyMatch(child -> impureExprs.contains(child.expr()))) {
impureExprs.add(expr);
}
if ((kind.equals(Kind.IDENT)
&& (expr.ident().name().startsWith(MANGLED_COMPREHENSION_ITER_VAR_PREFIX)
|| expr.ident().name().startsWith(MANGLED_COMPREHENSION_ACCU_VAR_PREFIX)))
|| node.children().anyMatch(child -> mangledExprs.contains(child.expr()))) {
mangledExprs.add(expr);
}
if (!kind.equals(Kind.CONSTANT)
&& !kind.equals(Kind.IDENT)
// Exclude empty lists (cel.bind sets this for iterRange).
&& !(kind.equals(Kind.LIST) && expr.list().elements().isEmpty())
&& !ineligibleExprs.contains(expr)
&& !impureExprs.contains(expr)
&& (kind.equals(Kind.COMPREHENSION)
? !referencesEnclosingComprehensionVars(node)
: !mangledExprs.contains(expr))) {
eligibleExprs.add(expr);
}
}
return eligibleExprs;
}

// Comprehensions may reference their own variables, but not those of enclosing ones (e.g.
// [2].exists(y, x == y) in [1].exists(x, [2].exists(y, x == y))), as that breaks scoping.
private static boolean referencesEnclosingComprehensionVars(CelNavigableMutableExpr node) {
return node.descendants()
.filter(d -> d.getKind().equals(Kind.IDENT))
.map(d -> d.expr().ident().name())
.anyMatch(CelNavigableExprUtil.getEnclosingComprehensionVariables(node)::contains);
}

/**
Expand Down Expand Up @@ -488,21 +550,14 @@ private static Set<CelMutableExpr> getIneligibleExprsFromComprehensionBranches(
.allNodes())
.map(CelNavigableMutableExpr::expr)
.forEach(ineligibleExprs::add));

return ineligibleExprs;
}

private boolean containsEliminableFunctionOnly(CelNavigableMutableExpr navigableExpr) {
return navigableExpr
navAst
.getRoot()
.allNodes()
.allMatch(
node -> {
if (node.getKind().equals(Kind.CALL)) {
return cseEliminableFunctions.contains(node.expr().call().function());
}
.map(CelNavigableMutableExpr::expr)
.filter(SubexpressionOptimizer::isAccumulation)
.forEach(expr -> ineligibleExprs.add(expr.call().args().get(1)));

return true;
});
return ineligibleExprs;
}

// Assigns equal value numbers to equivalent subexpressions, including commuted and converse calls
Expand Down
26 changes: 26 additions & 0 deletions optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerBaselineTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import static com.google.common.truth.Truth.assertThat;
import static dev.cel.common.CelOverloadDecl.newGlobalOverload;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.testing.junit.testparameterinjector.TestParameter;
import com.google.testing.junit.testparameterinjector.TestParameterInjector;
Expand Down Expand Up @@ -95,6 +96,8 @@ private static Cel setupCelEnv(CelBuilder celBuilder) {
TestAllTypes.newBuilder()
.setSingleInt64(3L)
.setSingleInt32(5)
.addAllRepeatedInt64(ImmutableList.of(1L, 2L, 3L))
.addRepeatedInt32(2)
.setOneofType(
NestedTestAllTypes.newBuilder()
.setPayload(
Expand Down Expand Up @@ -412,6 +415,29 @@ private enum CseTestCase {
"(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)"),
HOIST_NESTED_COMPREHENSION("msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == 2))"),
HOIST_INVARIANT_ARITHMETIC(
"msg.repeated_int64.exists(i, i == msg.single_int64 * 2 + msg.single_int32)"),
HOIST_ZERO_TRIP_ERROR_NOT_EVALUATED(
"msg.repeated_nested_message.exists(m, m.bb == 1 / (msg.single_int32 - 5))"),
HOIST_MAP_ACCUMULATION_PRESERVED("msg.repeated_int64.map(i, msg.single_int64 * 2)"),
HOIST_LIST_LITERAL("msg.repeated_int64.exists(i, i in [msg.single_int64, msg.single_int32])"),
HOIST_LIST_APPENDED_TO_ITER_VAR("[[1], [2]].map(l, l + [msg.single_int64])"),
HOIST_INVARIANT_INSIDE_DUPLICATE(
"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])"),
HOIST_INVARIANT_OF_NESTED_COMPREHENSION(
"msg.repeated_int64.all(i, msg.repeated_int32.exists(j, j == msg.single_int64 * 2))"),
HOIST_SHARED_TERMS_ONCE(
"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"),
HOIST_SHARED_WITH_ACCUMULATED_LIST(
"msg.repeated_int64.exists(i, i in [msg.single_int64]) &&"
+ " msg.repeated_int64.map(i, msg.single_int64).size() > 0"),
DUPLICATE_WITH_SHARED_COMMUTED_OPERANDS(
"[msg.single_int64 * msg.single_int32, msg.single_int32 * msg.single_int64,"
+ " msg.single_int64 + msg.single_int32, msg.single_int32 + msg.single_int64]"),
;

private final String source;
Expand Down
35 changes: 29 additions & 6 deletions optimizer/src/test/java/dev/cel/optimizer/optimizers/SubexpressionOptimizerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -302,7 +302,9 @@ public void cse_applyConstFoldingAfter() throws Exception {
CelAbstractSyntaxTree optimizedAst = optimizer.optimize(ast);

assertThat(CEL_UNPARSER.unparse(optimizedAst))
.isEqualTo("cel.@block([1, 2], @index0 + @index0 + @index1 + @index1 + x)");
.isEqualTo(
"cel.@block([3, size([@index0]), size([1, @index0])], @index1 + @index1 + @index2 +"
+ " @index2 + x)");
}

@Test
Expand Down Expand Up @@ -347,6 +349,27 @@ public void iterationLimitReached_throws() throws Exception {
assertThat(e).hasMessageThat().isEqualTo("Optimization failure: Max iteration count reached.");
}

@Test
@TestParameters(
"{source: 'msg.repeated_int64.exists(i, i == (msg.single_int64 + msg.single_int32) *"
+ " (msg.single_int32 + msg.single_int64))', expected: 'cel.@block([msg.single_int64 +"
+ " msg.single_int32, @index0 * @index0], msg.repeated_int64.exists(@it:0:0, @it:0:0 =="
+ " @index1))'}")
@TestParameters(
"{source: 'msg.repeated_int64.map(i, cel.bind(a, i, a + msg.single_int32 * 2))', expected:"
+ " 'cel.@block([msg.single_int32 * 2], msg.repeated_int64.map(@it:1:0, cel.bind(a,"
+ " @it:1:0, @ac:0:0 + @index0)))'}")
public void cse_hoistLoopInvariants(String source, String expected) throws Exception {
CelAbstractSyntaxTree ast = cel.compile(source).getAst();
CelOptimizer celOptimizer =
newCseOptimizer(
SubexpressionOptimizerOptions.newBuilder().populateMacroCalls(true).build());

CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);

assertThat(CEL_UNPARSER.unparse(optimizedAst)).isEqualTo(expected);
}

@Test
public void cse_comprehensionReferencingOuterBindVar_notExtractedOutOfBindScope()
throws Exception {
Expand Down Expand Up @@ -757,12 +780,12 @@ public void cse_nestedMacro_noOp_assertAstIdCorrectness() throws Exception {
.addAstOptimizers(SubexpressionOptimizer.getInstance())
.build();
CelAbstractSyntaxTree ast =
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))").getAst();
cel.compile("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))").getAst();

CelAbstractSyntaxTree optimizedAst = celOptimizer.optimize(ast);

assertThat(CEL_UNPARSER.unparse(optimizedAst))
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(x.a))");
.isEqualTo("[{}, {\"a\": 1}, {\"b\": 2}].filter(m, has(m.a))");
assertThat(optimizedAst).isSameInstanceAs(ast);
}

Expand Down Expand Up @@ -797,8 +820,8 @@ public void cse_nestedMacro_withOptimization_assertAstIdCorrectness() throws Exc

assertThat(CEL_UNPARSER.unparse(optimizedAst))
.isEqualTo(
"cel.@block([[{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, has(x.a))], @index0 =="
+ " @index0)");
"cel.@block([has(x.a), [{}, {\"a\": 1}, {\"b\": 2}].filter(@it:0:0, @index0)], @index1"
+ " == @index1)");
}

@Test
Expand Down Expand Up @@ -869,7 +892,7 @@ private enum ValueNumberingTestCase {
"cel.@block([msg.single_int64 <= msg.single_int32], @index0 && @index0)"),
NOT_EQUALS_WITHOUT_SHARED_EQUALS_PRESERVED(
"[msg.single_int64 != 1, msg.single_int64 != 1, msg.single_int64 == 1]",
"cel.@block([msg.single_int64 != 1], [@index0, @index0, msg.single_int64 == 1])"),
"cel.@block([msg.single_int64, @index0 != 1], [@index1, @index1, @index0 == 1])"),
OPTIONAL_LIST_AND_MAP_DISTINCT_VALUE_NUMBERS(
"([?optional.of(msg.single_int64)].size() == [optional.of(msg.single_int64)].size()) &&"
+ " ({?\"k\": optional.of(msg.single_int64)}.size() == {\"k\":"
Expand Down
Loading
Loading