| 84 | cublas_check(cublasLtMatrixLayoutDestroy(layout_trans_c)); |
| 85 | } |
| 86 | void CUBLASLTMatmulDesc::set(const SizeArgs& args, bool batched) { |
| 87 | cublasOperation_t trans_a, trans_b; |
| 88 | auto m = args.layout_c.shape[batched ? 1 : 0], |
| 89 | n = args.layout_c.shape[batched ? 2 : 1]; |
| 90 | auto k = batched ? args.layout_a.shape[args.transposeA ? 1 : 2] |
| 91 | : args.layout_a.shape[args.transposeA ? 0 : 1]; |
| 92 | int batch = (batched ? args.layout_a.shape[0] : 1); |
| 93 | uint32_t pm = CUBLAS_POINTER_MODE_DEVICE; |
| 94 | dt_b = to_cuda_dtype(args.layout_b.dtype); |
| 95 | dt_a = to_cuda_dtype(args.layout_a.dtype); |
| 96 | dt_c = to_cuda_dtype(args.layout_c.dtype); |
| 97 | |
| 98 | megdnn_assert(dt_a == dt_b, "matrix A and B should have same precision"); |
| 99 | #if CUDA_VERSION >= 11000 |
| 100 | dt_compute = to_cublas_compute_type(args.layout_c.dtype); |
| 101 | cublas_check(cublasLtMatmulDescCreate(&matmul_desc, dt_compute, dt_c)); |
| 102 | #else |
| 103 | dt_compute = dt_c; |
| 104 | cublas_check(cublasLtMatmulDescCreate(&matmul_desc, dt_compute)); |
| 105 | #endif |
| 106 | cublas_check(cublasLtMatmulDescSetAttribute( |
| 107 | matmul_desc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pm, sizeof(pm))); |
| 108 | |
| 109 | cublasLtOrder_t order_COL32 = CUBLASLT_ORDER_COL32; |
| 110 | cublasLtOrder_t order_COL4_4R2_8C = CUBLASLT_ORDER_COL4_4R2_8C; |
| 111 | /** |
| 112 | * \NOTE that cublas takes column-major matrices as inputs, |
| 113 | * but megdnn takes row-major ones. |
| 114 | * So we calculate C^t = B^t * A^t by cublas. Here the transpose symbol |
| 115 | * implies row-major to column-major conversion |
| 116 | */ |
| 117 | if (dt_c == CUDA_R_32I) { |
| 118 | /** |
| 119 | * \NOTE: To use IMMA kernels, use computeType = CUDA_R_32I and |
| 120 | * CUBLASLT_ORDER_COL32 for matrices A,C,D and |
| 121 | * CUBLASLT_ORDER_COL4_4R2_8C for matrix B. |
| 122 | */ |
| 123 | int ldbtransform, ldatransform, ldctransform; |
| 124 | size_t stride_b_trans, stride_a_trans, stride_c_trans; |
| 125 | ldbtransform = 32 * n; |
| 126 | ldatransform = 32 * round_up<int32_t>(m, 8); |
| 127 | ldctransform = 32 * n; |
| 128 | stride_b_trans = round_up<int32_t>(k, 32) / 32 * ldbtransform; |
| 129 | stride_a_trans = round_up<int32_t>(k, 32) / 32 * ldatransform; |
| 130 | stride_c_trans = round_up<int32_t>(m, 32) / 32 * ldctransform; |
| 131 | trans_b = CUBLAS_OP_T; |
| 132 | cublas_check(cublasLtMatmulDescSetAttribute( |
| 133 | matmul_desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b))); |
| 134 | // origin layout |
| 135 | cublas_check(cublasLtMatrixLayoutCreate( |
| 136 | &layout_b, dt_b, n, k, args.layout_b.stride[batched ? 1 : 0])); |
| 137 | cublas_check(cublasLtMatrixLayoutCreate( |
| 138 | &layout_a, dt_a, k, m, args.layout_a.stride[batched ? 1 : 0])); |
| 139 | cublas_check(cublasLtMatrixLayoutCreate( |
| 140 | &layout_c, dt_c, n, m, args.layout_c.stride[batched ? 1 : 0])); |
| 141 | // transformed layout |
| 142 | cublas_check( |
| 143 | cublasLtMatrixLayoutCreate(&layout_trans_b, dt_b, n, k, ldbtransform)); |
no test coverage detected