| 3355 | class CUDABlasLtMatmulPlan final : public blas::IBlasLtMatmulPlan { |
| 3356 | public: |
| 3357 | port::Status init(const blas::BlasLtMatmulPlanParams& p) { |
| 3358 | params_ = p; |
| 3359 | scale_type_ = GetScaleType(p.c_type, p.computation_type); |
| 3360 | SE_ASSIGN_OR_RETURN( |
| 3361 | op_desc_, |
| 3362 | CreateCublasLtOperationDesc( |
| 3363 | p.computation_type, GetScaleType(p.c_type, p.computation_type), |
| 3364 | p.pointer_mode, p.epilogue, p.transa, p.transb)); |
| 3365 | uint64 rows_a = p.transa == blas::Transpose::kNoTranspose ? p.m : p.k; |
| 3366 | uint64 cols_a = p.transa == blas::Transpose::kNoTranspose ? p.k : p.m; |
| 3367 | uint64 rows_b = p.transb == blas::Transpose::kNoTranspose ? p.k : p.n; |
| 3368 | uint64 cols_b = p.transb == blas::Transpose::kNoTranspose ? p.n : p.k; |
| 3369 | SE_ASSIGN_OR_RETURN( |
| 3370 | a_desc_, CreateCublasLtLayoutDesc(p.ab_type, rows_a, cols_a, p.lda, |
| 3371 | p.stride_a, capped_batch_count())); |
| 3372 | SE_ASSIGN_OR_RETURN( |
| 3373 | b_desc_, CreateCublasLtLayoutDesc(p.ab_type, rows_b, cols_b, p.ldb, |
| 3374 | p.stride_b, capped_batch_count())); |
| 3375 | SE_ASSIGN_OR_RETURN( |
| 3376 | c_desc_, CreateCublasLtLayoutDesc(p.c_type, p.m, p.n, p.ldc, p.stride_c, |
| 3377 | capped_batch_count())); |
| 3378 | SE_ASSIGN_OR_RETURN( |
| 3379 | d_desc_, CreateCublasLtLayoutDesc(p.c_type, p.m, p.n, p.ldc, p.stride_c, |
| 3380 | capped_batch_count())); |
| 3381 | remainder_batch_count_ = |
| 3382 | p.batch_count > kMaxBatchCount ? p.batch_count % kMaxBatchCount : 0; |
| 3383 | if (remainder_batch_count_) { |
| 3384 | SE_ASSIGN_OR_RETURN( |
| 3385 | a_remainder_desc_, |
| 3386 | CreateCublasLtLayoutDesc(p.ab_type, rows_a, cols_a, p.lda, p.stride_a, |
| 3387 | remainder_batch_count_)); |
| 3388 | SE_ASSIGN_OR_RETURN( |
| 3389 | b_remainder_desc_, |
| 3390 | CreateCublasLtLayoutDesc(p.ab_type, rows_b, cols_b, p.ldb, p.stride_b, |
| 3391 | remainder_batch_count_)); |
| 3392 | SE_ASSIGN_OR_RETURN( |
| 3393 | c_remainder_desc_, |
| 3394 | CreateCublasLtLayoutDesc(p.c_type, p.m, p.n, p.ldc, p.stride_c, |
| 3395 | remainder_batch_count_)); |
| 3396 | SE_ASSIGN_OR_RETURN( |
| 3397 | d_remainder_desc_, |
| 3398 | CreateCublasLtLayoutDesc(p.c_type, p.m, p.n, p.ldc, p.stride_c, |
| 3399 | remainder_batch_count_)); |
| 3400 | } |
| 3401 | return port::Status::OK(); |
| 3402 | } |
| 3403 | |
| 3404 | cublasLtMatmulDesc_t op_desc() const { return op_desc_.get(); } |
| 3405 | cublasLtMatrixLayout_t a_desc() const { return a_desc_.get(); } |
no test coverage detected