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

TPU layout method cleanup by removing one parameter.

Browse Source 
PiperOrigin-RevId: 360292316
Change-Id: Ib8332ba5b0d7576cfb02773f9974a48ead7ae47c
This commit is contained in:
A. Unique TensorFlower 2021-03-01 15:33:37 -08:00 committed by TensorFlower Gardener
parent 6441ec4b0c
commit 1038045628
2 changed files with 9 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
tensorflow/compiler/xla/service/bfloat16_propagation.cc
View File

@ -345,11 +345,8 @@ bool BFloat16Propagation::AllUsersConsumeBF16(const HloInstruction& hlo,
return true; return true;
} }
namespace { bool BFloat16Propagation::ShouldKeepPrecisionUnchanged(
const HloInstruction* inst) {
// Returns whether we should avoid changing the precision of inst regardless of
// the producers and users.
bool ShouldKeepPrecisionUnchanged(const HloInstruction* inst) {
if (inst->opcode() == HloOpcode::kFusion && if (inst->opcode() == HloOpcode::kFusion &&
inst->fusion_kind() == HloInstruction::FusionKind::kCustom) { inst->fusion_kind() == HloInstruction::FusionKind::kCustom) {
return ShouldKeepPrecisionUnchanged( return ShouldKeepPrecisionUnchanged(
@ -358,14 +355,12 @@ bool ShouldKeepPrecisionUnchanged(const HloInstruction* inst) {
// Do not change precision for side-effecting instructions, control flow, and // Do not change precision for side-effecting instructions, control flow, and
// bitcast-convert, because this pass might break the interfaces or // bitcast-convert, because this pass might break the interfaces or
// assumptions for them. // assumptions for them.
return inst->opcode() == HloOpcode::kCustomCall || // return inst->opcode() == HloOpcode::kCustomCall ||
inst->opcode() == HloOpcode::kCall || // inst->opcode() == HloOpcode::kCall ||
inst->opcode() == HloOpcode::kBitcastConvert || // inst->opcode() == HloOpcode::kBitcastConvert ||
inst->HasSideEffectNoRecurse(); inst->HasSideEffectNoRecurse();
} }
} // namespace
void BFloat16Propagation::DetermineInstructionPrecision(HloInstruction* hlo, void BFloat16Propagation::DetermineInstructionPrecision(HloInstruction* hlo,
bool skip_parameters) { bool skip_parameters) {
// We handle any fusion computation, while body/condition or conditional // We handle any fusion computation, while body/condition or conditional

4
tensorflow/compiler/xla/service/bfloat16_propagation.h
View File

@ -72,6 +72,10 @@ class BFloat16Propagation : public HloModulePass {
// (precision reductions were added). // (precision reductions were added).
StatusOr<bool> Run(HloModule* module) override; StatusOr<bool> Run(HloModule* module) override;
// Returns whether we should avoid changing the precision of inst regardless
// of the producers and users.
virtual bool ShouldKeepPrecisionUnchanged(const HloInstruction* inst);
private: private:
// *************************** // ***************************
// Function called and state produced by the forward analysis pass (from // Function called and state produced by the forward analysis pass (from