Function FloorDivImpl

tensorflow/compiler/tf2xla/kernels/binary_ops.cc:114–140

Implementation of FloorDiv. For floating-point values, simply returns floor(x / y). For integers, does: if ((x < 0) != (y < 0)) { T abs_x = std::abs(x); T abs_y = std::abs(y); return -(abs_x + abs_y - 1) / abs_y; } else { return x / y; }

//   return x / y;
// }
static xla::XlaOp FloorDivImpl(xla::XlaBuilder* b, DataType dtype, xla::XlaOp x,
                               xla::XlaOp y, const BCast& broadcast_helper) {
  std::tie(x, y) = XlaBinaryOp::Broadcast(x, y, broadcast_helper);
  if (DataTypeIsFloating(dtype)) {
    if (dtype == DataType::DT_BFLOAT16) {
      // The result of a BF16 division may produce the Ceil of what was
      // computed by F32 division, so avoid end user confusion by doing the
      // intermediate divide in F32.
      return xla::ConvertElementType(
          xla::Floor(xla::Div(xla::ConvertElementType(x, xla::F32),
                              xla::ConvertElementType(y, xla::F32))),
          xla::BF16);
    } else {
      return xla::Floor(xla::Div(x, y));
    }
  }
  if (DataTypeIsUnsigned(dtype)) {
    return xla::Div(x, y);
  }
  auto zero = XlaHelpers::Zero(b, dtype);
  auto one = XlaHelpers::One(b, dtype);
  auto different_sign = xla::Ne(xla::Lt(x, zero), xla::Lt(y, zero));
  auto abs_x = xla::Abs(x);
  auto abs_y = xla::Abs(y);
  auto t = xla::Neg(xla::Sub(xla::Add(abs_x, abs_y), one));
  return xla::Select(different_sign, xla::Div(t, abs_y), xla::Div(x, y));
}
XLA_MAKE_BINARY(FloorDiv,
                FloorDivImpl(b, input_type(0), lhs, rhs, broadcast_helper));