| 3301 | } |
| 3302 | |
| 3303 | port::StatusOr<UniqueLayoutDesc> CreateCublasLtLayoutDesc( |
| 3304 | blas::DataType data_type, uint64 rows, uint64 cols, int64 ld, int64 stride, |
| 3305 | int batch_count) { |
| 3306 | cublasLtMatrixLayout_t desc; |
| 3307 | cublasStatus_t status = cublasLtMatrixLayoutCreate( |
| 3308 | &desc, GetCUDADataType(data_type), rows, cols, ld); |
| 3309 | if (status != CUBLAS_STATUS_SUCCESS) { |
| 3310 | return port::Status( |
| 3311 | port::error::INTERNAL, |
| 3312 | absl::StrCat("cublasLtMatrixLayoutCreate failed: ", ToString(status))); |
| 3313 | } |
| 3314 | UniqueLayoutDesc unique_desc(desc); |
| 3315 | SE_RETURN_IF_ERROR( |
| 3316 | SetCublasLtAttr(desc, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, batch_count)); |
| 3317 | SE_RETURN_IF_ERROR(SetCublasLtAttr( |
| 3318 | desc, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, stride)); |
| 3319 | return unique_desc; |
| 3320 | } |
| 3321 | |
| 3322 | // Helper function to allocate workspace. |
| 3323 | port::Status AllocateWorkspace(void** workspace, |
no test coverage detected