MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / configurePlugin

Method configurePlugin

plugin/fcPlugin/fcPlugin.cpp:477–560  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

475}
476
477void FCPluginDynamic::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
478 DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
479{
480 try
481 {
482 // Validate input arguments
483 PLUGIN_VALIDATE(nbOutputs == 1);
484 PLUGIN_VALIDATE(nbInputs == 1);
485 PLUGIN_VALIDATE(inputs != nullptr);
486 PLUGIN_VALIDATE(outputs != nullptr);
487 PLUGIN_VALIDATE(mType == inputs[0].desc.type);
488 auto const& inDims0 = inputs[0].desc.dims;
489
490 PLUGIN_VALIDATE(inDims0.nbDims == 5);
491 mK = inDims0.d[HDIM]; // hiddensize
492 // PLUGIN_ASSERT(hiddenSize * mOutDim == mNumParams);
493 PLUGIN_VALIDATE(inDims0.d[3] == 1);
494 PLUGIN_VALIDATE(inDims0.d[4] == 1);
495
496 // m and k are mOutDim
497 // n is B*S
498 int32_t const S = inputs->max.d[SDIM];
499 int32_t const B = inputs->max.d[BDIM];
500
501 mNmax = S * B;
502
503 // Cleanup LtContext descriptors before creating new ones.
504 mLtContext.destroy();
505
506 if (mType == DataType::kFLOAT)
507 {
508 Gemm<float> g(mOutDim, mNmax, mK, false, false);
509 mLtContext.create(g, kMAX_WORKSPACE_BYTES);
510 }
511 else if (mType == DataType::kHALF)
512 {
513 Gemm<half> g(mOutDim, mNmax, mK, false, false);
514 mLtContext.create(g, kMAX_WORKSPACE_BYTES);
515 }
516 else
517 {
518 std::string const msg = "Unsupported type error, expected [kHALF,kFLOAT], but received ";
519 PLUGIN_VALIDATE(false, (msg + std::to_string(static_cast<int32_t>(mType))).c_str());
520 }
521
522 gLogVerbose << "FCPluginDynamic configurePlugin m=" << mOutDim << ", n=" << mNmax << ", k=" << mK << std::endl;
523
524 size_t actualWorkspace = 0;
525 if (mAlgo.data[0] == 0 && memcmp(mAlgo.data, mAlgo.data + 1, sizeof(mAlgo.data) - sizeof(mAlgo.data[0])) == 0)
526 {
527 gLogVerbose << "FCPluginDynamic gemmSearch\n";
528 if (mType == DataType::kFLOAT)
529 {
530 mAlgo = gemmSearch<float>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace);
531 }
532 else if (mType == DataType::kHALF)
533 {
534 mAlgo = gemmSearch<half>(mOutDim, mNmax, mK, kMAX_WORKSPACE_BYTES, actualWorkspace);

Callers 1

pyPlugin.cppFile · 0.45

Calls 5

caughtErrorFunction · 0.85
createMethod · 0.80
c_strMethod · 0.80
populateMethod · 0.80
destroyMethod · 0.45

Tested by

no test coverage detected