Skip to content
  • Taylor Cramer's avatar
    Modifications from profiling, including removal of `isinstance` casing · 91dd2a0e
    Taylor Cramer authored
    These are modifications I made in order to speed up the workload timing
    out in serialization + deserialization in b/156371202. I intend to follow
    it with a change to intern types, whose serialization, deserialization, and
    comparison is currently causing exponential blowups due to the nested
    structure of type signatures and the frequency with which they occur
    throughout the AST (one full copy per node).
    
    The biggest change in this CL comes from the move away from using
    `isinstance` to case over AST blocks and type nodes. This change was motivated
    by a few factors:
    - It's significantly faster (a single method lookup vs. type tree traversal).
    - It clarifies which `isinstance` checks are dynamic Python "downcasts" vs.
      "switch"es on a kind of sum-type. The downcast-like spots are in my opinion
      more suspect and should be subject to a decent bit of scrutiny when added,
      while switches on a sum-type should not be cause for concern.
    - I personally find it more ergonomic and easier to read:
      `my_structure.is_lambda()` vs.
      `isinstance(my_structure, computation_types.Lambda)`
    
    Some other performance-motivated changes included here are:
    - Caching the list-of-tuples representation of AnonymousTuples.
    - Moving formatting of error messages out of the success path.
    - Removing the `are_equivalent_types` check on building blocks
      which ensured that the deserialized type of a building block
      matched the type of the deserialized building block. This
      check isn't exactly redundant, but given correct serialization
      and deserialization it will never trigger, and it is extremely
      expensive (expontentially complex). This may be worth looking
      into reenabling once I've landed a change to intern Types.
    - Removing the `check_well_formed` on building blocks __init__.
      This is similar to the above in that it should be unnecessary
      given correct TFF code, and is extremely expensive. As above,
      it may be worth trying to enable when we start interning types.
    - I did *not* disable the `is_concrete_instance_of` check on
      the `Intrinsic` building block, which is also quite expensive.
      I think this one is rather higher-value than the two above,
      but this is totally a judgement call as they are all quite
      costly and scale exponentially in time with the depth of the
      type tree. This should also hopefully be much cheaper once
      we have type interning.
    - Moving the type tree visitor to internal iteration.
      I did not do the same for the building block visitor yet, since
      it has more visitors of varied shape and complexity, but I think
      this is worth investigating in the future.
    
    PiperOrigin-RevId: 317769828
    91dd2a0e