message: [XLA/GPU] Declares that rsqrt is cheap and should be fused.
diff: @@ -29,13 +29,24 @@ limitations under the License.
  namespace xla {
  namespace gpu {
  +bool ElementIsF32OrF16(const Shape& shape) {
  + PrimitiveType type = shape.element_type();
  + return type == F32 || type == F16;
  +}
  +
  /*static*/ bool GpuInstructionFusion::IsExpensive(
  const HloInstruction& instruction) {
  - // We say that floating-point division is cheap on the GPU.
  - if (instruction.opcode() == HloOpcode::kDivide &&
  - ShapeUtil::ElementIsFloating(instruction.shape())) {
  +
  + // We say that some floating-point math ops are cheap on the GPU.
  + switch (instruction.opcode()) {
  + case HloOpcode::kDivide:
  + case HloOpcode::kRsqrt:
  + if (ElementIsF32OrF16(instruction.shape())) {
  return false;
  }
  + default:
  + break;
  + }
  return InstructionFusion::IsExpensive(instruction);
  }
  