| 94 | } |
| 95 | |
| 96 | port::StatusOr<const blas::PlanAndAlgorithms*> GetPlanAndAlgorithms( |
| 97 | Stream* stream, BatchMatmulParameters matmul_parameters, int64_t batch_size, |
| 98 | tensorflow::DataType dtype, blas::MatrixDescriptor lhs_matrix, |
| 99 | blas::MatrixDescriptor rhs_matrix, blas::MatrixDescriptor output_matrix) { |
| 100 | static const int64_t max_scratch_size = |
| 101 | GetWorkspaceLimit(1LL << 32); // 4GB by default |
| 102 | static const int64_t max_autotune_algorithm_count = |
| 103 | MatmulMaxAutotuneAlgorithmCount(); |
| 104 | const blas::PlanAndAlgorithms* plan_and_algorithms = |
| 105 | BatchMatmulPlanMapSingleton::GetInstance()->Find(matmul_parameters); |
| 106 | if (!plan_and_algorithms) { |
| 107 | TF_ASSIGN_OR_RETURN( |
| 108 | blas::BlasLtMatmulPlanParams plan_params, |
| 109 | CreatePlanParams(batch_size, dtype, matmul_parameters.GetEpilogOp(), |
| 110 | lhs_matrix, rhs_matrix, output_matrix)); |
| 111 | |
| 112 | TF_ASSIGN_OR_RETURN(std::unique_ptr<blas::IBlasLtMatmulPlan> plan, |
| 113 | stream->parent()->CreateBlasLtMatmulPlan(plan_params)); |
| 114 | TF_ASSIGN_OR_RETURN( |
| 115 | std::vector<std::unique_ptr<blas::IBlasLtMatmulAlgorithm>> algorithms, |
| 116 | stream->parent()->GetBlasLtMatmulAlgorithms( |
| 117 | plan.get(), max_scratch_size, |
| 118 | /* max_algorithm_count */ max_autotune_algorithm_count)); |
| 119 | |
| 120 | plan_and_algorithms = BatchMatmulPlanMapSingleton::GetInstance()->Insert( |
| 121 | matmul_parameters, {std::move(plan), std::move(algorithms)}); |
| 122 | } |
| 123 | return plan_and_algorithms; |
| 124 | } |
| 125 | |
| 126 | port::StatusOr<blas::BlasLtMatmulPlanParams> CreatePlanParams( |
| 127 | int64_t batch_size, tensorflow::DataType dtype, blas::Epilogue epilog_op, |
no test coverage detected