| 107 | } |
| 108 | |
| 109 | void MatrixMulForward::check_exec( |
| 110 | const TensorLayout& A, const TensorLayout& B, const TensorLayout& C, |
| 111 | size_t workspace_in_bytes) { |
| 112 | auto errmsg = [&]() { |
| 113 | std::string msg; |
| 114 | msg.append("A="); |
| 115 | msg.append(A.to_string()); |
| 116 | msg.append(", B="); |
| 117 | msg.append(B.to_string()); |
| 118 | msg.append(", C="); |
| 119 | msg.append(C.to_string()); |
| 120 | msg.append(", transposeA="); |
| 121 | msg.append(std::to_string(param().transposeA)); |
| 122 | msg.append(", transposeB="); |
| 123 | msg.append(std::to_string(param().transposeB)); |
| 124 | return msg; |
| 125 | }; |
| 126 | MEGDNN_MARK_USED_VAR(errmsg); |
| 127 | if (param().format == param::MatrixMul::Format::DEFAULT) { |
| 128 | megdnn_assert_eq_size_t(A.ndim, 2_z); |
| 129 | megdnn_assert_eq_size_t(B.ndim, 2_z); |
| 130 | megdnn_assert_eq_size_t(C.ndim, 2_z); |
| 131 | |
| 132 | megdnn_assert(A.stride[1] == 1); |
| 133 | megdnn_assert(A.stride[0] >= static_cast<ptrdiff_t>(A.shape[1])); |
| 134 | megdnn_assert(B.stride[1] == 1); |
| 135 | megdnn_assert(B.stride[0] >= static_cast<ptrdiff_t>(B.shape[1])); |
| 136 | megdnn_assert(C.stride[1] == 1); |
| 137 | megdnn_assert(C.stride[0] >= static_cast<ptrdiff_t>(C.shape[1])); |
| 138 | size_t A0, A1, B0, B1, C0, C1; |
| 139 | A0 = A.shape[0]; |
| 140 | A1 = A.shape[1]; |
| 141 | B0 = B.shape[0]; |
| 142 | B1 = B.shape[1]; |
| 143 | C0 = C.shape[0]; |
| 144 | C1 = C.shape[1]; |
| 145 | if (m_param.transposeA) |
| 146 | std::swap(A0, A1); |
| 147 | if (m_param.transposeB) |
| 148 | std::swap(B0, B1); |
| 149 | megdnn_assert(A0 == C0, "%s", errmsg().c_str()); |
| 150 | megdnn_assert(B1 == C1, "%s", errmsg().c_str()); |
| 151 | megdnn_assert(A1 == B0, "%s", errmsg().c_str()); |
| 152 | } else if (param().format == param::MatrixMul::Format::N32K4_DOT) { |
| 153 | size_t A0 = A.shape[0]; |
| 154 | size_t A1 = A.shape[1]; |
| 155 | size_t B2 = B.shape[2]; |
| 156 | size_t B3 = B.shape[3]; |
| 157 | megdnn_assert(!m_param.transposeA && !m_param.transposeB); |
| 158 | megdnn_assert(A0 == 1 && A1 % 4 == 0); |
| 159 | megdnn_assert(B.ndim == 4); |
| 160 | megdnn_assert(B2 == 32 && B3 == 4); |
| 161 | megdnn_assert_contiguous(A); |
| 162 | megdnn_assert_contiguous(B); |
| 163 | megdnn_assert_contiguous(C); |
| 164 | } else { |
| 165 | megdnn_assert_eq_size_t(A.ndim, 4_z); |
| 166 | megdnn_assert_eq_size_t(B.ndim, 3_z); |