-
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