Method addConv

core/graph.cpp:60–149

  }

  Ref<Op> Graph::addConv(const std::string& name,
                         const Ref<Op>& srcOp,
                         Activation activation,
                         Fusion fusion)
  {
    if (fusion != Fusion::None && !engine->isConvSupported(fusion))
    {
      // If the engine does not support the specified fused convolution, split it into two ops
      switch (fusion)
      {
      case Fusion::UpsampleSrc0:
        {
           auto srcUpsample = addUpsample(name + "_upsample", srcOp);
           return addConv(name, srcUpsample, activation);
        }

      case Fusion::PoolDst:
        {
           auto conv = addConv(name, srcOp, activation);
           return addPool(name + "_pool", conv);
        }

      default:
        throw std::invalid_argument("unsupported convolution fusion");
      }
    }

    const std::string weightName = name + ".weight";
    const std::string biasName   = name + ".bias";
    Ref<Tensor> weight = (*constTensors)[weightName];
    Ref<Tensor> bias   = (*constTensors)[biasName];

    if (weight->getRank() != 4 || bias->getRank() != 1)
      throw std::invalid_argument("invalid convolution weight/bias");

    Device* device = engine->getDevice();
    const int blockC = device->getTensorBlockC();

    TensorDims finalWeightDims{round_up(weight->getO(), blockC),
                               round_up(weight->getI(), blockC),
                               weight->getH(),
                               weight->getW()};

    TensorDesc finalWeightDesc = {weight->getDims(),
                                  finalWeightDims,
                                  device->getWeightLayout(),
                                  device->getWeightDataType()};

    TensorDesc finalBiasDesc = {bias->getDims(),
                                {round_up(bias->getX(), blockC)},
                                TensorLayout::x,
                                device->getTensorDataType()};

    auto srcAlloc = tensorAllocs[srcOp.get()];
    auto conv = engine->newConv({srcAlloc->desc, finalWeightDesc, finalBiasDesc, activation, fusion, fastMath});
    conv->setName(name);
    auto dstAlloc = addOp(conv, {srcOp}, conv->getDstDesc());