MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / DoBlasLtMatmulInternal

Method DoBlasLtMatmulInternal

tensorflow/stream_executor/cuda/cuda_blas.cc:3618–3776  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

3616
3617#if CUDA_VERSION >= 11000
3618bool CUDABlas::DoBlasLtMatmulInternal(
3619 Stream* stream, bool err_on_failure, const blas::IBlasLtMatmulPlan* plan,
3620 const HostOrDeviceScalar<void>& alpha, DeviceMemoryBase a,
3621 DeviceMemoryBase b, const HostOrDeviceScalar<void>& beta,
3622 DeviceMemoryBase c, DeviceMemoryBase d, ScratchAllocator* scratch_allocator,
3623 const blas::IBlasLtMatmulAlgorithm* algorithm, DeviceMemoryBase bias) {
3624 const auto& cuda_plan = *static_cast<const CUDABlasLtMatmulPlan*>(plan);
3625 const auto& cuda_algo =
3626 *static_cast<const CUDABlasLtMatmulAlgorithm*>(algorithm);
3627
3628 if (alpha.data_type() != cuda_plan.scale_type() ||
3629 beta.data_type() != cuda_plan.scale_type()) {
3630 VLOG(2) << "DoBlasLtMatmul returning false because alpha and beta types do "
3631 "not match plan: expected "
3632 << cuda_plan.c_type() << ", got alpha=" << alpha.data_type()
3633 << " beta=" << beta.data_type();
3634 return false;
3635 }
3636 if (alpha.is_pointer() != beta.is_pointer()) {
3637 VLOG(2) << "DoBlasLtMatmul returning false because one of `alpha` "
3638 "and `beta` is a pointer, but the other is not.";
3639 return false;
3640 }
3641 bool is_pointer_mode_host = !alpha.is_pointer();
3642 if ((cuda_plan.params().pointer_mode == blas::PointerMode::kHost) !=
3643 is_pointer_mode_host) {
3644 VLOG(2) << "DoBlasLtMatmul returning false because plan has wrong "
3645 "pointer_mode for the given alpha/beta.";
3646 return false;
3647 }
3648 if ((cuda_plan.params().epilogue == blas::Epilogue::kBias ||
3649 cuda_plan.params().epilogue == blas::Epilogue::kBiasThenReLU) !=
3650 (bias != nullptr)) {
3651 VLOG(2) << "DoBlasLtMatmul returning false because plan has wrong "
3652 "epilogue for the given bias pointer.";
3653 return false;
3654 }
3655 const void* alpha_ptr = alpha.is_pointer() ? alpha.opaque_pointer().opaque()
3656 : alpha.opaque_value();
3657 const void* beta_ptr =
3658 beta.is_pointer() ? beta.opaque_pointer().opaque() : beta.opaque_value();
3659
3660 void* workspace = nullptr;
3661 if (cuda_algo.workspace_size()) {
3662 port::Status allocation_status = AllocateWorkspace(
3663 &workspace, scratch_allocator, cuda_algo.workspace_size());
3664 if (!allocation_status.ok()) {
3665 if (err_on_failure || VLOG_IS_ON(3)) {
3666 LOG(ERROR)
3667 << "Failed to allocate workspace for cublasLtMatmul algo with id: "
3668 << cuda_algo.algo_id() << " requiring "
3669 << cuda_algo.workspace_size() << " bytes of workspace";
3670 }
3671 return false;
3672 }
3673 }
3674
3675 // This is only used when batch_count > kMaxBatchCount.

Callers

nothing calls this directly

Calls 15

AllocateWorkspaceFunction · 0.85
GetDataTypeSizeBytesFunction · 0.85
scale_typeMethod · 0.80
c_typeMethod · 0.80
paramsMethod · 0.80
opaqueMethod · 0.80
opaque_valueMethod · 0.80
workspace_sizeMethod · 0.80
remainder_batch_countMethod · 0.80
ConsumeValueOrDieMethod · 0.80
SetBiasPointerMethod · 0.80
capped_batch_countMethod · 0.80

Tested by

no test coverage detected