Rewrite graph to combine similarly-shaped variables (faster startup).

def autostack(self):
    """Rewrite graph to combine similarly-shaped variables (faster startup)."""
    num_slices = 0
    for v in self.graph.all_variables:
      num_slices += self.mesh_to_impl[v.mesh].size
    if num_slices >= 2 ** 16:
      # Startup times are slow with lots of variable slices.
      # Perform more aggressive stacking
      max_combined_slice_size = 2 ** 27
    else:
      # Stacking hurts memory utilization - only stack small variables.
      max_combined_slice_size = 2 ** 16
    self.graph.rewrite_stack_variables(
        mesh_to_impl=self.mesh_to_impl,
        max_combined_slice_size=max_combined_slice_size)