Reduce the memory usage of sharding domains
Previously the domain instructions inserted before and after an `n` element tuple required `O(n^2)` memory (and compute) because every operand and user had its own domain instruction with a tuple sharding and tuple shape for the exit domains what constructed `n` HloSharding and `n` Shape proto per domain. After this change we keep track of the domain instructions inserted and if we already have a domain instruction with the correct operand and metadata then we re-use it instead of creating a new one. Additionally we change HloInstruction and ShardingMetadata to store a std::shared_ptr to HloSharding so the same instance can be shared by many instructions. This CL doesn't update all uses to remove all of the duplicated HloShardings but handles the most wastful cases to reduce memory usage. PiperOrigin-RevId: 209924260
Loading
Please sign in to comment