Auto-cluster TensorFlow's resource operations when legal.
This CL teaches auto-clustering to cluster resource variable operations when
proved legal by a data-flow analysis. It is organized as follows:
- The core data flow analysis lives in jit/resource_operation_safety_analysis.
The output of this analysis is a set of pairs of nodes where two nodes that
belong to the same pair should not be put into the same cluster.
- jit/xla_cluster_util now has a routine that incorporates the results of the
data flow analysis into the cycle detection graph such that auto-clustering
automatically avoids putting incompatible nodes in the same cluster.
- jit/mark_for_compilation_pass and jit/xla_fusion_optimizer have some plumbing
changes to make all of this work. Changes to jit/xla_fusion_optimizer are a
bit simpler because, IIUC, it does not have to deal with the XLA_{CPU|GPU}
devices.
PiperOrigin-RevId: 210006397
Loading
Please sign in to comment