Home
last modified time | relevance | path

Searched refs:reshard_cache (Results 1 – 4 of 4) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/service/spmd/
Dspmd_partitioner.h271 std::vector<std::pair<HloSharding, PartitionedHlo>> reshard_cache; member
290 ReshardCache* reshard_cache; member
490 state.reshard_cache = &reshard_cache_; in MakePartitioningState()
Dspmd_partitioner.cc349 auto& cache = state_.reshard_cache->per_hlo_cache[hlo()].reshard_cache; in Reshard()
360 state_.reshard_cache->per_hlo_cache[resharded.hlo()] in Reshard()
361 .reshard_cache.emplace_back(sharding(), *this); in Reshard()
550 auto& cache = state_.reshard_cache->per_hlo_cache[hlo()].window_reshard_cache; in ReshardAsWindowedInput()
824 auto& cache = state_.reshard_cache->per_hlo_cache[hlo()].reshard_cache; in Replicate()
845 state_.reshard_cache->per_hlo_cache[resharded.hlo()] in Replicate()
846 .reshard_cache.emplace_back(sharding, *this); in Replicate()
Dspmd_partitioner_util.cc1685 state.reshard_cache->groupd_caches[absl::StrJoin(per_group_strings, ";")]; in CreatePerGroupPartitioningState()
1689 result.reshard_cache = grouped_cache.get(); in CreatePerGroupPartitioningState()
Ddot_handler.cc830 state.reshard_cache->per_hlo_cache.erase(padded_slice_operand); in PartitionBaseCase()
1082 state.reshard_cache->per_hlo_cache.erase(slice_operand); in PartitionBaseCase()