experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Make IdenticalSlowPath more strict

Browse Source 
- Map did not check |dimensions|. This has not resulted in miscompliation
  because none of the backends/frontends support/generate map's with dimensions.
- DynamicSlice did not check |dynamic_slice_sizes|. This has not resulted in
  miscompliation because it is not possible for two DynamicSlice operations
  with different |dynamic_slice_sizes| to have the same shape.
- HloCollective did not check |constrain_layout|.

PiperOrigin-RevId: 316765004
Change-Id: Id2629fef71446842eeae18901142a502a634b010
This commit is contained in:
David Majnemer 2020-06-16 15:01:21 -07:00 committed by TensorFlower Gardener
parent dd789e7783
commit d4c2030375
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorflow/compiler/xla/service/hlo_instructions.cc
View File

@ -550,6 +550,7 @@ bool HloCollectiveInstruction::IdenticalSlowPath(
const auto& casted_other =
static_cast<const HloCollectiveInstruction&>(other);
return HloChannelInstruction::IdenticalSlowPath(other, eq_computations) &&
constrain_layout() == casted_other.constrain_layout() &&
absl::c_equal(replica_groups(), casted_other.replica_groups(),
[](const ReplicaGroup& a, const ReplicaGroup& b) {
return absl::c_equal(a.replica_ids(), b.replica_ids());
@ -1101,7 +1102,9 @@ bool HloMapInstruction::IdenticalSlowPath(
const HloInstruction& other,
const std::function<bool(const HloComputation*, const HloComputation*)>&
eq_computations) const {
return eq_computations(to_apply(), other.to_apply());
const auto& casted_other = static_cast<const HloMapInstruction&>(other);
return eq_computations(to_apply(), casted_other.to_apply()) &&
dimensions() == casted_other.dimensions();
}
std::unique_ptr<HloInstruction> HloMapInstruction::CloneWithNewOperandsImpl(
@ -2515,7 +2518,8 @@ bool HloDynamicSliceInstruction::IdenticalSlowPath(
const HloInstruction& other,
const std::function<bool(const HloComputation*, const HloComputation*)>&
eq_computations) const {
return true;
const auto& casted_other = static_cast<const HloMapInstruction&>(other);
return dynamic_slice_sizes() == casted_other.dynamic_slice_sizes();
}
std::unique_ptr<HloInstruction>