Method configurePlugin

plugin/embLayerNormPlugin/embLayerNormPlugin.cpp:250–302

}

void EmbLayerNormPluginDynamic::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
{
    BERT_DEBUG_MSG("EmbLayerNormPluginDynamic configurePlugin.");

    // Validate input arguments
    PLUGIN_ASSERT(nbOutputs == 2);
    PLUGIN_ASSERT(nbInputs == 3);

    PLUGIN_ASSERT(inputs[0].desc.dims.nbDims == 2);
    int32_t const S = inputs[0].desc.dims.d[SDIM];
    mS = S;
    int32_t const B = inputs[0].desc.dims.d[BDIM];
    TRT_UNUSED B;
    PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[1].desc.dims.d[SDIM]));
    PLUGIN_ASSERT(B == inputs[1].desc.dims.d[BDIM]);
    PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[2].desc.dims.d[SDIM]));
    PLUGIN_ASSERT(B == inputs[2].desc.dims.d[BDIM]);

    PLUGIN_ASSERT(outputs[0].desc.dims.nbDims == 5);
    PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[SDIM]) == mS);
    PLUGIN_ASSERT(outputs[0].desc.dims.d[BDIM] == B);
    PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[2]) == mLd);
    PLUGIN_ASSERT(outputs[0].desc.dims.d[3] == 1);
    PLUGIN_ASSERT(outputs[0].desc.dims.d[4] == 1);

    if (mUseFullMask)
    {
        // user force full_mask
        PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 2);
        PLUGIN_ASSERT(outputs[1].desc.dims.d[0] == B);
        PLUGIN_ASSERT((outputs[1].desc.dims.d[1] == -1) || (outputs[1].desc.dims.d[1] == packedMaskSize384)
            || (outputs[1].desc.dims.d[1] == packedMaskSize128));
    }
    else
    {
        // auto detect using mhatype
        if (S != -1 && B != -1)
        {
            PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 2);
            PLUGIN_ASSERT(outputs[1].desc.dims.d[0] == B);
            int32_t packedSize = getMHAMaskPackedSize(mSM, mMhaType, S);
            TRT_UNUSED packedSize;
            PLUGIN_ASSERT(outputs[1].desc.dims.d[1] == -1 || outputs[1].desc.dims.d[1] == packedSize);
        }
    }

    PLUGIN_ASSERT(inputs[0].desc.type == DataType::kINT32);
    PLUGIN_ASSERT(inputs[1].desc.type == DataType::kINT32);
    PLUGIN_ASSERT(inputs[2].desc.type == DataType::kINT32);
    PLUGIN_ASSERT(outputs[0].desc.type == mType);
    PLUGIN_ASSERT(outputs[1].desc.type == DataType::kINT32);
}

size_t EmbLayerNormPluginDynamic::getWorkspaceSize(
    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept