| 24 | return desc.get_workspace_bundle(ltArgs, algo).total_size_in_bytes(); |
| 25 | } |
| 26 | void MatrixMulForwardImpl::AlgoCuBlasLt::exec(const ExecArgs& args) const { |
| 27 | CUBLASLTMatmulDesc::SizeArgs ltArgs(args); |
| 28 | cublasLtMatmulAlgo_t algo; |
| 29 | CUBLASLTMatmulDesc desc(ltArgs); |
| 30 | auto&& handle = ltArgs.handle; |
| 31 | auto&& stream = handle->stream(); |
| 32 | auto&& cublasLt_handle = handle->cublasLt_handle(); |
| 33 | desc.get_algorithm_heuristic(ltArgs, INT_MAX, algo); |
| 34 | auto&& ws_bundle = desc.get_workspace_bundle(ltArgs, algo); |
| 35 | ws_bundle.set(args.workspace.raw_ptr); |
| 36 | |
| 37 | auto sgemm = [&]() { |
| 38 | auto zero = handle->zero_device(); |
| 39 | auto one = handle->one_device(); |
| 40 | megdnn_assert( |
| 41 | ws_bundle.nr_workspace() == 1, |
| 42 | "workspace bundle size should be 1(ws_algo)"); |
| 43 | cublas_check(cublasLtMatmul( |
| 44 | cublasLt_handle, desc.matmul_desc, one, |
| 45 | static_cast<void*>(args.tensor_b.ptr<dt_float32>()), desc.layout_b, |
| 46 | static_cast<void*>(args.tensor_a.ptr<dt_float32>()), desc.layout_a, |
| 47 | zero, static_cast<void*>(args.tensor_c.ptr<dt_float32>()), |
| 48 | desc.layout_c, static_cast<void*>(args.tensor_c.ptr<dt_float32>()), |
| 49 | desc.layout_c, &algo, ws_bundle.get(0), ws_bundle.get_size(0), stream)); |
| 50 | }; |
| 51 | auto hgemm = [&]() { |
| 52 | auto zero_half = handle->zero_device_h(); |
| 53 | auto one_half = handle->one_device_h(); |
| 54 | megdnn_assert( |
| 55 | ws_bundle.nr_workspace() == 1, |
| 56 | "workspace bundle size should be 1(ws_algo)"); |
| 57 | cublas_check(cublasLtMatmul( |
| 58 | cublasLt_handle, desc.matmul_desc, one_half, |
| 59 | static_cast<const __half*>(args.tensor_b.raw_ptr()), desc.layout_b, |
| 60 | static_cast<const __half*>(args.tensor_a.raw_ptr()), desc.layout_a, |
| 61 | zero_half, static_cast<const __half*>(args.tensor_c.raw_ptr()), |
| 62 | desc.layout_c, static_cast<__half*>(args.tensor_c.raw_ptr()), |
| 63 | desc.layout_c, &algo, ws_bundle.get(0), ws_bundle.get_size(0), stream)); |
| 64 | }; |
| 65 | auto igemm = [&]() { |
| 66 | auto zero = handle->zero_device(); |
| 67 | auto one = handle->one_device(); |
| 68 | megdnn_assert( |
| 69 | ws_bundle.nr_workspace() == 4, |
| 70 | "workspace bundle size should be 4(ws_algo, ws_a, ws_b, ws_c)"); |
| 71 | void* ws_b = ws_bundle.get(1); |
| 72 | void* ws_a = ws_bundle.get(2); |
| 73 | void* ws_c = ws_bundle.get(3); |
| 74 | int32_t pm = CUBLAS_POINTER_MODE_DEVICE; |
| 75 | cublasOperation_t trans_a = CUBLAS_OP_T, trans_c = CUBLAS_OP_N; |
| 76 | cublasLtMatrixTransformDesc_t transform_desc = nullptr; |
| 77 | cublas_check(cublasLtMatrixTransformDescCreate(&transform_desc, CUDA_R_32F)); |
| 78 | cublas_check(cublasLtMatrixTransformDescSetAttribute( |
| 79 | transform_desc, CUBLASLT_MATRIX_TRANSFORM_DESC_POINTER_MODE, &pm, |
| 80 | sizeof(pm))); |
| 81 | cublas_check(cublasLtMatrixTransform( |
| 82 | cublasLt_handle, transform_desc, one, args.tensor_b.raw_ptr(), |
| 83 | desc.layout_b, zero, nullptr, nullptr, ws_b, desc.layout_trans_b, |
nothing calls this directly
no test coverage detected