| 221 | |
| 222 | template<typename Ti, typename To> |
| 223 | void gemm(Array<To> &out, af_mat_prop optLhs, af_mat_prop optRhs, |
| 224 | const To *alpha, const Array<Ti> &lhs, const Array<Ti> &rhs, |
| 225 | const To *beta) { |
| 226 | const CBLAS_TRANSPOSE lOpts = toCblasTranspose(optLhs); |
| 227 | const CBLAS_TRANSPOSE rOpts = toCblasTranspose(optRhs); |
| 228 | |
| 229 | const int aRowDim = (lOpts == CblasNoTrans) ? 0 : 1; |
| 230 | const int aColDim = (lOpts == CblasNoTrans) ? 1 : 0; |
| 231 | const int bColDim = (rOpts == CblasNoTrans) ? 1 : 0; |
| 232 | |
| 233 | const dim4 &lDims = lhs.dims(); |
| 234 | const dim4 &rDims = rhs.dims(); |
| 235 | const int M = lDims[aRowDim]; |
| 236 | const int N = rDims[bColDim]; |
| 237 | const int K = lDims[aColDim]; |
| 238 | const dim4 oDims = out.dims(); |
| 239 | |
| 240 | using BT = typename blas_base<Ti>::type; |
| 241 | using CBT = const typename blas_base<Ti>::type; |
| 242 | |
| 243 | auto alpha_ = scale_type<Ti, false>(alpha); |
| 244 | auto beta_ = scale_type<Ti, false>(beta); |
| 245 | #ifdef USE_MKL |
| 246 | auto alpha_batched = scale_type<Ti, true>(alpha); |
| 247 | auto beta_batched = scale_type<Ti, true>(beta); |
| 248 | #endif |
| 249 | |
| 250 | auto func = [=](Param<Ti> output, CParam<Ti> left, CParam<Ti> right) { |
| 251 | dim4 lStrides = left.strides(); |
| 252 | dim4 rStrides = right.strides(); |
| 253 | dim4 oStrides = output.strides(); |
| 254 | |
| 255 | if (output.dims().ndims() <= 2) { |
| 256 | if (right.dims()[bColDim] == 1) { |
| 257 | dim_t incr = |
| 258 | (optRhs == AF_MAT_NONE) ? rStrides[0] : rStrides[1]; |
| 259 | gemv_func<Ti>()( |
| 260 | CblasColMajor, lOpts, lDims[0], lDims[1], alpha_.getScale(), |
| 261 | reinterpret_cast<CBT *>(left.get()), lStrides[1], |
| 262 | reinterpret_cast<CBT *>(right.get()), incr, |
| 263 | beta_.getScale(), reinterpret_cast<BT *>(output.get()), |
| 264 | oStrides[0]); |
| 265 | } else { |
| 266 | gemm_func<Ti>()( |
| 267 | CblasColMajor, lOpts, rOpts, M, N, K, alpha_.getScale(), |
| 268 | reinterpret_cast<CBT *>(left.get()), lStrides[1], |
| 269 | reinterpret_cast<CBT *>(right.get()), rStrides[1], |
| 270 | beta_.getScale(), reinterpret_cast<BT *>(output.get()), |
| 271 | oStrides[1]); |
| 272 | } |
| 273 | } else { |
| 274 | int batchSize = static_cast<int>(oDims[2] * oDims[3]); |
| 275 | |
| 276 | const bool is_l_d2_batched = oDims[2] == lDims[2]; |
| 277 | const bool is_l_d3_batched = oDims[3] == lDims[3]; |
| 278 | const bool is_r_d2_batched = oDims[2] == rDims[2]; |
| 279 | const bool is_r_d3_batched = oDims[3] == rDims[3]; |
| 280 |
no test coverage detected