| 566 | } |
| 567 | |
| 568 | int32_t FCPluginDynamic::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, |
| 569 | void const* const* inputs, void* const* outputs, void* workSpace, cudaStream_t stream) noexcept |
| 570 | { |
| 571 | try |
| 572 | { |
| 573 | size_t const workspaceSize = getWorkspaceSize(inputDesc, 1, outputDesc, 1); |
| 574 | |
| 575 | int32_t const S = inputDesc->dims.d[SDIM]; |
| 576 | int32_t const B = inputDesc->dims.d[BDIM]; |
| 577 | int32_t const n = S * B; |
| 578 | PLUGIN_VALIDATE(n >= 0); |
| 579 | mLtContext.setN(static_cast<uint64_t>(n)); |
| 580 | |
| 581 | if (mType == DataType::kFLOAT) |
| 582 | { |
| 583 | auto const* const input = static_cast<float const*>(inputs[0]); |
| 584 | auto* output = static_cast<float*>(outputs[0]); |
| 585 | |
| 586 | Gemm<float> g(mOutDim, n, mK, false, false); |
| 587 | if (mWdev == nullptr) |
| 588 | { |
| 589 | return STATUS_FAILURE; |
| 590 | } |
| 591 | g.A = static_cast<float*>(mWdev.get()); |
| 592 | g.B = const_cast<float*>(input); |
| 593 | g.C = output; |
| 594 | |
| 595 | return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream); |
| 596 | } |
| 597 | if (mType == DataType::kHALF) |
| 598 | { |
| 599 | auto const* const input = static_cast<half const*>(inputs[0]); |
| 600 | auto* output = static_cast<half*>(outputs[0]); |
| 601 | |
| 602 | Gemm<half> g(mOutDim, n, mK, false, false); |
| 603 | if (mWdev == nullptr) |
| 604 | { |
| 605 | return STATUS_FAILURE; |
| 606 | } |
| 607 | g.A = static_cast<half*>(mWdev.get()); |
| 608 | g.B = const_cast<half*>(input); |
| 609 | g.C = output; |
| 610 | return cublasLtMatmul(mLtContext, g, mAlgo, workSpace, workspaceSize, stream); |
| 611 | } |
| 612 | else |
| 613 | { |
| 614 | gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType) |
| 615 | << std::endl; |
| 616 | return STATUS_FAILURE; |
| 617 | } |
| 618 | } |
| 619 | catch (std::exception const& e) |
| 620 | { |
| 621 | caughtError(e); |
| 622 | } |
| 623 | return STATUS_FAILURE; |
| 624 | } |
| 625 |
no test coverage detected