bool HabanaBackend::isOpSupported()

in lib/Backends/Habana/Habana.cpp [1087:1170]


bool HabanaBackend::isOpSupported(const NodeInfo &NI) const {
  if (NI.getKind() == Kinded::Kind::SaveNodeKind) {
    return true;
  }

  if (isQuantizedType(NI.getOutElemTy(0))) {
    switch (NI.getKind()) {
    case Kinded::Kind::AddNodeKind:
    case Kinded::Kind::AvgPoolNodeKind:
    case Kinded::Kind::ConvolutionNodeKind:
    case Kinded::Kind::DivNodeKind:
    case Kinded::Kind::FullyConnectedNodeKind:
    case Kinded::Kind::HabanaConvolutionNodeKind:
    case Kinded::Kind::HabanaConvolutionAddNodeKind:
    case Kinded::Kind::HabanaFullyConnectedNodeKind:
    case Kinded::Kind::MatMulNodeKind:
    case Kinded::Kind::MaxPoolNodeKind:
    case Kinded::Kind::MulNodeKind:
    case Kinded::Kind::QuantizeNodeKind:
    case Kinded::Kind::ReluNodeKind:
    case Kinded::Kind::ReshapeNodeKind:
    case Kinded::Kind::SliceNodeKind:
    case Kinded::Kind::SplatNodeKind:
    case Kinded::Kind::SubNodeKind:
    case Kinded::Kind::TileNodeKind:
    case Kinded::Kind::ConcatNodeKind:
    case Kinded::Kind::TransposeNodeKind:
      return true;
    case Kinded::Kind::RescaleQuantizedNodeKind:
      return NI.allInputsAndOutputsHaveSameElemKind(
          {ElemKind::Int8QTy, ElemKind::Int16QTy});
    default:
      return false;
    }
  }

  switch (NI.getKind()) {
  case Kinded::Kind::AddNodeKind:
  case Kinded::Kind::AvgPoolNodeKind:
  case Kinded::Kind::BatchBoxCoxNodeKind:
  case Kinded::Kind::BatchMatMulNodeKind:
  case Kinded::Kind::BatchedAddNodeKind:
  case Kinded::Kind::BatchedReduceAddNodeKind:
  case Kinded::Kind::ConcatNodeKind:
  case Kinded::Kind::DequantizeNodeKind:
  case Kinded::Kind::DivNodeKind:
  case Kinded::Kind::FullyConnectedNodeKind:
  case Kinded::Kind::HabanaFullyConnectedNodeKind:
  case Kinded::Kind::LogNodeKind:
  case Kinded::Kind::MatMulNodeKind:
  case Kinded::Kind::MaxNodeKind:
  case Kinded::Kind::MaxPoolNodeKind:
  case Kinded::Kind::MinNodeKind:
  case Kinded::Kind::MulNodeKind:
  case Kinded::Kind::ReluNodeKind:
  case Kinded::Kind::ReshapeNodeKind:
  case Kinded::Kind::SaveNodeKind:
  case Kinded::Kind::SigmoidNodeKind:
  case Kinded::Kind::SliceNodeKind:
  case Kinded::Kind::SoftMaxNodeKind:
  case Kinded::Kind::SplatNodeKind:
  case Kinded::Kind::SubNodeKind:
  case Kinded::Kind::TanhNodeKind:
  case Kinded::Kind::TileNodeKind:
  case Kinded::Kind::TransposeNodeKind:
  case Kinded::Kind::LengthsRangeFillNodeKind:
  case Kinded::Kind::LengthsSumNodeKind:
  case Kinded::Kind::SparseLengthsSumNodeKind:
  case Kinded::Kind::SparseLengthsWeightedSumNodeKind:
  case Kinded::Kind::FusedRowwiseQuantizedSparseLengthsSumNodeKind:
  case Kinded::Kind::FusedRowwiseQuantizedSparseLengthsWeightedSumNodeKind:
  case Kinded::Kind::LocalResponseNormalizationNodeKind:
    return true;
  case Kinded::Kind::BatchOneHotNodeKind:
    return NI.allInputsAndOutputsHaveSameElemKind({ElemKind::Int32ITy});
  case Kinded::Kind::GatherNodeKind:
    // Gather is technically supported but currently appears to trigger bugs in
    // large models.
  default:
    return false;
  }

  return false;
}