Method addConcatConv

core/graph.cpp:158–310

  }

  Ref<Op> Graph::addConcatConv(const std::string& name,
                               const Ref<Op>& src0Op,
                               const Ref<Op>& src1Op,
                               Activation activation,
                               Fusion fusion)
  {
    Device* device = engine->getDevice();

    if (!engine->isConcatConvSupported(fusion) && !engine->isConcatConv2Supported(fusion))
    {
      // Need to split into separate upsample and concat+conv
      auto src0Upsample = addUpsample(name + "_upsample0", src0Op);
      return addConcatConv(name, src0Upsample, src1Op, activation);
    }

    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");

    const int blockC = device->getTensorBlockC();

    auto src0Alloc = tensorAllocs[src0Op.get()];
    auto src1Alloc = tensorAllocs[src1Op.get()];
    const TensorDesc src0Desc = src0Alloc->desc;
    const TensorDesc src1Desc = src1Alloc->desc;

    TensorDims finalWeightDims{round_up(weight->getO(), blockC),
                               src0Desc.getPaddedC() + src1Desc.getPaddedC(),
                               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()};

    ConcatConvDesc concatConvDesc{src0Desc, src1Desc, finalWeightDesc, finalBiasDesc, activation, fusion, fastMath};

    if (engine->isConcatConvSupported(fusion))
    {
      auto concatConv = engine->newConcatConv(concatConvDesc);
      concatConv->setName(name);
      auto dstAlloc = addOp(concatConv, {src0Op, src1Op}, concatConv->getDstDesc(),
                            concatConv->isPreConcatConv());

      lazyInits.push_back([=]()
      {
        concatConv->setSrc(src0Alloc->tensor, src1Alloc->tensor);
        concatConv->setDst(dstAlloc->tensor);