Branches in generated code?

[aaron-skydio]

Goal

Currently we support branching at the symbolic level by evaluating each path and doing a weighted sum of the branches given a conditional. This is fast for small branches with shared CSE, but cases with significant non-shared work will be slower than code that has a true branch.

The goal is to resolve this with a hard branching option in codegen that doesn't evaluate every branch.A small motivating example here was Rot3.from_rotation_matrix.

Piecewise functions

SymPy has Piecewise functions that print to inline ternaries that are quite relevant here. Side note: These were super slow for me at the symbolic level in the past dealing with splines, so I had implemented a simpler form in our codebase with some stricter assumptions that was much better. However, it generally works. The way it does printing is with nested ternary statements. This is something that basically should work within SymForce, but doesn't solve this problem because non-shared CSE terms can't be included inside the ternaries as separate statements.

Basic Expression Implementation

It's not hard to make a simple custom expression in SymPy (or symengine). I think derivatives, CSE, and numeric substitution work nominally with the below. The main challenge is on the printing side and how to handle CSE. We could also decide to augment Piecewise instead.

import sympy as sm

class Branch(sm.Function):
    nargs = 3

    @classmethod
    def eval(cls, conditional, x, y):
        # Simplify if conditional is not symbolic
        if conditional.is_number:
            if conditional.is_positive:
                return x
            else:
                return y

    def _eval_derivative(self, s):
        conditional = self.args[0]
        operands = self.args[1:]
        return self.__class__(conditional, *(v.diff(s) for v in operands))

CSE

CSE will pull out intermediate terms across everything. Let's say we then make a function for identifying which intermediates are needed where given the set of temporaries and the symbolic expression. We could somehow try to first do CSE on the inner scopes, but I think it might be easier to use the current CSE then post-process and pass that to printing somehow.

During printing SymPy traverses the expression tree as scalars and assumes you're printing expressions. This makes it hard to put statements inline. For example, how do we compute tmp1 only when needed in the example below? (ignoring tmp0 for now, but that is complex in its own right as it's needed in 2 out of 3 branches)

double Func(double x) {
	double tmp0 = pow(x, 2);
	double tmp1 = sin(x);
	
	return ((x < -1) ? (
		   0
		)
		: ((x <= 1) ? (
		   tmp0
		)
		: (
		   tmp0/tmp1 + tmp1
		)));
}

Handwritten, we'd hope for something like:

double Func2(double x) {
	double tmp0 = pow(x, 2);

    if (x < -1) {
        return 0;
    } else if (x <= 1) {
        return tmp0;
    } else {
	      double tmp1 = sin(x);
        return tmp0/tmp1 + tmp1;
    }
}

One idea is to use lambdas and let the complier optimize those out:

double Func3(double x) {
	double tmp0 = pow(x, 2);

	return ((x < -1) ? (
		   0
		)
		: ((x <= 1) ? (
		   tmp0
		)
		: [&](){
	     double tmp1 = sin(x);
		   return tmp0/tmp1 + tmp1;
      }()));
}

Well if we're doing that we could also do it on the whole thing and take control of the printing of that piecewise scalar ourselves.

double Func4(double x) {
	double tmp0 = pow(x, 2);

    return [&]() -> double{
        if (x < -1) {
            return 0;
        } else if (x <= 1) {
            return tmp0;
        } else {
            double tmp1 = sin(x);
            return tmp0/tmp1 + tmp1;
        }
    }();
}

I realized in this simple case the compiler is going to figure it out anyway and not compute sin unless needed. Can see this here, with all three versions being equivalent:

https://godbolt.org/z/G11rb1Yh8

So on one hand we could implement the lambda thing. On the other hand, one could argue what the point is if the compiler does it. Some things to think about here:

1. Will the compiler be able to do this in complicated scenarios where there are chains of branches on top of each other? Is our codegen from_rotation_matrix actually more instructions at the assembly level?
2. There is potentially some value in making the human-readable code easy to understand, pushing towards lambda over ternary and factoring the sub-expressions close to usage.
3. In many of our use cases we want to have a branch that applies to a bunch of outputs, not just one scalar output. Each output will share a conditional and likely share sub-expressions. In the current formulations, each scalar would generate its own if statements, and we might see that any sub-expressions needed for two scalars (which could be in the same logical branch) will still need to be at the top level.
   • SymPy pretty strongly draws a line that a function takes in multiple arguments and returns a scalar. Where could we try to inject the idea of a vector-valued output and effectively use that in printing?

Interested to discuss these points, particularly using Rot3.from_rotation_matrix as an example.

Weighted Sum

Might be better and more general to support a WeightedSum instead of a Branch, where you have a vector of weights and arguments and return the sum. A two-way only branch and relying on nesting seems prone to ugliness. Maybe the two-way case can be a wrapper helper.

Tangent: Naming

Where, Branch, If, Cond

[JzHuai0108]

I think this issue is related to the below example:

import symforce
symforce.set_epsilon_to_symbol()

import symforce.symbolic as sf

from symforce.codegen import Codegen
from symforce.codegen import CppConfig

def piecewise_factor(td: sf.Scalar, t0: sf.Scalar, x: sf.Scalar, knot: sf.Scalar) -> sf.Scalar:
    t = t0 + td
    if t > knot:
        return x ** 3 + td ** 2
    else:
        return x ** 2 + td ** 3

output_dir = "/home/pi/Desktop/temp"
print_code = True

codegen = Codegen.function(piecewise_factor, config=CppConfig())
metadata = codegen.generate_function(output_dir=output_dir, skip_directory_nesting=True)
if print_code:
    print(metadata.generated_files[0].read_text())

codegen_with_jacs = codegen.with_jacobians(which_args=["td", "x"])
metadata = codegen_with_jacs.generate_function(
    output_dir=output_dir, skip_directory_nesting=True
)
if print_code:
    print(metadata.generated_files[0].read_text())

The generated code looks like

// -----------------------------------------------------------------------------
// This file was autogenerated by symforce from template:
//     function/FUNCTION.h.jinja
// Do NOT modify by hand.
// -----------------------------------------------------------------------------

#pragma once

#include <Eigen/Dense>

namespace sym {

/**
 * This function was autogenerated from a symbolic function. Do not modify by hand.
 *
 * Symbolic function: piecewise_factor
 *
 * Args:
 *     td: Scalar
 *     t0: Scalar
 *     x: Scalar
 *     knot: Scalar
 *
 * Outputs:
 *     res: Scalar
 */
template <typename Scalar>
Scalar PiecewiseFactor(const Scalar td, const Scalar t0, const Scalar x, const Scalar knot) {
  // Total ops: 3

  // Unused inputs
  (void)t0;
  (void)knot;

  // Input arrays

  // Intermediate terms (0)

  // Output terms (1)
  Scalar _res;

  _res = std::pow(td, Scalar(2)) + [&]() {
    const Scalar base = x;
    return base * base * base;
  }();

  return _res;
}  // NOLINT(readability/fn_size)

// NOLINTNEXTLINE(readability/fn_size)
}  // namespace sym

// -----------------------------------------------------------------------------
// This file was autogenerated by symforce from template:
//     function/FUNCTION.h.jinja
// Do NOT modify by hand.
// -----------------------------------------------------------------------------

#pragma once

#include <Eigen/Dense>

namespace sym {

/**
 * This function was autogenerated from a symbolic function. Do not modify by hand.
 *
 * Symbolic function: piecewise_factor
 *
 * Args:
 *     td: Scalar
 *     t0: Scalar
 *     x: Scalar
 *     knot: Scalar
 *
 * Outputs:
 *     res: Scalar
 *     res_D_td: (1x1) jacobian of res (1) wrt arg td (1)
 *     res_D_x: (1x1) jacobian of res (1) wrt arg x (1)
 */
template <typename Scalar>
Scalar PiecewiseFactorWithJacobians02(const Scalar td, const Scalar t0, const Scalar x,
                                      const Scalar knot,
                                      Eigen::Matrix<Scalar, 1, 1>* const res_D_td = nullptr,
                                      Eigen::Matrix<Scalar, 1, 1>* const res_D_x = nullptr) {
  // Total ops: 6

  // Unused inputs
  (void)t0;
  (void)knot;

  // Input arrays

  // Intermediate terms (0)

  // Output terms (3)
  Scalar _res;

  _res = std::pow(td, Scalar(2)) + [&]() {
    const Scalar base = x;
    return base * base * base;
  }();

  if (res_D_td != nullptr) {
    Eigen::Matrix<Scalar, 1, 1>& _res_D_td = (*res_D_td);

    _res_D_td(0, 0) = 2 * td;
  }

  if (res_D_x != nullptr) {
    Eigen::Matrix<Scalar, 1, 1>& _res_D_x = (*res_D_x);

    _res_D_x(0, 0) = 3 * std::pow(x, Scalar(2));
  }

  return _res;
}  // NOLINT(readability/fn_size)

// NOLINTNEXTLINE(readability/fn_size)
}  // namespace sym

which essentially ignores the second branch.

[aaron-skydio]

It's somewhat related, yes. The way you'd currently do what you want in SymForce is something like this:

def piecewise_factor(td: sf.Scalar, t0: sf.Scalar, x: sf.Scalar, knot: sf.Scalar) -> sf.Scalar:
    t = t0 + td
    c = sf.greater(t, knot)
    not_c = sf.logical_not(c, unsafe=True)
    return c * (x ** 3 + td ** 2) + not_c * (x ** 2 + td ** 3)

This should represent the function you want correctly.

This issue is generally about some alternatives we could take to the current approach, mostly with the goal of not needing to compute intermediate results for both branches. Regardless, python control flow isn't going to end up in the generated function, branches will be expressed symbolically, e.g. with SymPy's Piecewise; this is because SymForce is calling your function once to get the expression to generate, so it can only ever take one path through your function when that happens.

[JzHuai0108]

I'd like to confirm this.

Do you mean for now the generated code needs to compute immediate results for both branches?

From the generated code, it seems this is the case.
If so, I would write the branching code myself, and in each branch call the code generated by symforce for efficiency.

[aaron-skydio]

Yes that's right, it'll compute intermediate results for both branches. If it's easy to evaluate the condition yourself like that and generate two functions, yeah that's a good way to do this