1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ 17 #define TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ 18 19 #include "tensorflow/compiler/xla/service/hlo_instruction.h" 20 #include "tensorflow/compiler/xla/service/hlo_opcode.h" 21 22 namespace xla { 23 24 class BFloat16Support { 25 public: BFloat16Support()26 BFloat16Support() {} ~BFloat16Support()27 virtual ~BFloat16Support() {} 28 29 // Returns whether the backend supports BF16 operand for the HLO instruction 30 // at the given index. 31 virtual bool SupportsBF16Operand(const HloInstruction& hlo, 32 int64 operand_index) const; 33 34 // Returns whether the backend supports BF16 output for the HLO instruction. 35 virtual bool SupportsBF16Output(const HloInstruction& hlo) const; 36 37 // Returns whether the backend support mixed precision: the operands, output, 38 // and parameters/output of the called computations can have different 39 // precisions (BF16 and F32). 40 virtual bool SupportsMixedPrecisions(const HloInstruction& hlo) const; 41 42 // Returns whether the given HLO preserves its BF16 operand precision at the 43 // given index, so even if the output is F32, elements in the output that 44 // depend on the BF16 operand will still have BF16 effective precision even if 45 // they have F32 format. Similarly, this also means if the output is BF16 then 46 // increasing the operand precision from BF16 to F32 will not change the 47 // output. This typically includes HLOs that pass elements from the operand to 48 // the output without arithmetic operations. 49 static bool EffectiveOperandPrecisionIsOutputPrecision( 50 const HloInstruction& hlo, int64 operand_index); 51 52 // Returns if the backend only uses BF16 precision for the operand at the 53 // specified index, even if the operand is F32. 54 virtual bool EffectiveOperandPrecisionIsBF16(const HloInstruction& hlo, 55 int64 operand_index) const; 56 }; 57 58 } // namespace xla 59 60 #endif // TENSORFLOW_COMPILER_XLA_SERVICE_BFLOAT16_SUPPORT_H_ 61