| 475 | } |
| 476 | |
| 477 | void 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); |
no test coverage detected