Canonical way to lift loop invariants? (Top-down find all sub expressions independent of list of symbols) #28101

csp256 · 2025-05-30T21:29:00Z

csp256
May 30, 2025

I am using sympy to emit generated code for a somewhat complicated residual vector and its Jacobian for numerical optimization with Levenberg Marquardt in C++, Python, and Matlab. I am using cse and _forward_jacobian_cse to optimize these expressions. I have a classic block structure to my residual / Jacobian: there are groups of rows which are identical, except a subset of their inputs vary for each group/block.

This means that a portion of the code for each residual / Jacobian block is shared. I would like to explicitly lift the shared computation before I generate the code. My code thus-far looks something like:

generator = numbered_symbols(prefix='x')
residual_replacements, reduced_residual = cse(residual_block, symbols=generator, optimizations='basic', order='none')
replacements, jacobian, _ = _forward_jacobian_cse(residual_replacements, reduced_residual, opt_params)
additional_replacements, output = cse(reduced_residual + jacobian, symbols=generator, optimizations='basic', order='none')
replacements += additional_replacements

My first stab at lifting shared compute was simply to find a list of temporaries which did not depend on the list of varying symbols, taking care to add new temporary symbols to the list if they were dependent on symbols which varied:

varies_per_block = set((...)) # set of Symbols
def is_invariant(expr, varying_symbols):
    return 0 == len(expr.free_symbols.intersection(varying_symbols))
invariants = set()
for temporary, sub_expr in replacements:
    if is_invariant(sub_expr, varies_per_block):
        invariants.add(temporary)
    else:
        varies_per_block.add(temporary)

From there it is straightforward to split the list of replacements into two lists: invariant_replacements, and varying_replacements.

invariant_replacements = []
varying_replacements = []
for repl in replacements:
    if repl[0] in invariants:
        invariant_replacements.append( repl ) 
    else:
        varying_replacements.append( repl )

However, this does not work in two cases:

what if there is invariant computations in the residual or Jacobian, instead of one of the "replacements" factored out by cse?
what if an expression tree does vary block-to-block, but contains one or more sub expressions which are invariant?

The first is easily solved by ensuring the residuals and Jacobians do no computation themselves, by explicitly unconditionally pulling out the expression tree for each term into a temporary that is simply appended to the end of the list of temporaries.

replacements += [(Symbol(f"residual_{i}"), expr) for i,expr in enumerate(output[0])]
num_residuals = len(output[0])
output[0] = ImmutableDenseMatrix([replacements[-num_residuals+i][0] for i in range(num_residuals)])

num_opt_params = output[1].shape[1]
final_jac_replacements = [[(Symbol(f"jacobian_row{row}_col{col}")) for col in range(num_opt_params)] for row in range(num_residuals)]
for row
8000
 in range(num_residuals):
    for col in range(num_opt_params):
        replacements.append( (final_jac_replacements[row][col], output[1][row, col] ) )
output[1] = ImmutableDenseMatrix(final_jac_replacements)

The second however is a little trickier, since is requires treacherously inserting values into replacements as you traverse it:

i = 0
while i < len(replacements):
    temporary, expr = replacements[i]
    if is_invariant(expr, varying_symbols):
        # This whole expression is already invariant between blocks
        invariants.add(temporary)
    else:
        # check to see if this contains a sub expression that is itself invariant
        for sub_expr in preorder_traversal(expr):
            # We don't need to extract leaf nodes from our expression tree
            if len(sub_expr.args) == 0:
                continue
            if is_invariant(sub_expr, varying_symbols):
                # Replace this sub-expression with a new temporary, 
                # and add it to list of replacements just before this expression.
                new_temporary = next(generator)
                invariants.add(new_temporary)
                # Mind the ordering of insert() and this indexing: other order will break things!
                replacements[i] = (temporary, expr.subs(sub_expr, new_temporary))
                new_replacement = (new_temporary, sub_expr)
                replacements.insert(i, new_replacement)
                # We will re-examine this same expr next iteration, intentionally, as there may be multiple invariant sub-expressions
                break
    i += 1

In my case the initial attempt to lift loop invariants moved only 10% of the work, but hunting for all sub expressions allowed me to lift 90% of the compute out of the loop. However, this seems like a fair bit of work to go through. Is there capability in Sympy that can help me with these procedures? Are there functions that are meant to consume the output of cse(), returning another directed acyclic graph?

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

Canonical way to lift loop invariants? (Top-down find all sub expressions independent of list of symbols) #28101

Uh oh!

{{title}}

Uh oh!

Uh oh!

{{editor}}'s edit

{{editor}}'s edit

Uh oh!

Replies: 0 comments

Select a reply

Uh oh!

Uh oh!

Canonical way to lift loop invariants? (Top-down find all sub expressions independent of list of symbols) #28101

Uh oh!

Uh oh!

csp256 May 30, 2025

Replies: 0 comments

csp256
May 30, 2025