}
    *count = total_count;
    return true;
  }

  // For conv2d/depthwise_conv/fully_connected ops.
  // This algorithm actually comes from TOCO tooling_util.cc
  static bool GetArithmeticCountForConvAndFullyconnectedOp(mlir::Operation* op,
                                                           int64_t* count) {
    auto weight = op->getOperand(1);
    auto weight_type =

class CpuConvOp : public TargetHardwareOperation {
  double GetOpCost(mlir::Operation* op) const override {
    float cost = 0.0;
    int64_t arithmetic_count;
    if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
            op, &arithmetic_count)) {
      cost = arithmetic_count * kCPUArithmeticUnitCost;
    } else {
      cost = kCPUDefaultFixedValuedCost;
    }

class GpuConvOp : public TargetHardwareOperation {
  double GetOpCost(mlir::Operation* op) const override {
    int64_t arithmetic_count;
    if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
            op, &arithmetic_count)) {
      return arithmetic_count * kGPUArithmeticUnitCost;
    }
    return kGPUDefaultFixedValuedCost;
  }

  results.add<RemoveOptionalZeroBias<FullyConnectedOp>>(context);
}

int64_t FullyConnectedOp::GetArithmeticCount(Operation* op) {
  int64_t count;
  if (ArithmeticCountUtilHelper::GetArithmeticCountForConvAndFullyconnectedOp(
          op, &count))
    return count;

  return -1;
}

//===----------------------------------------------------------------------===//
// Conv2DOp