| 3371 | } |
| 3372 | |
| 3373 | bool MIOpenSupport::DoMatMul(Stream* stream, |
| 3374 | const DeviceMemory<float>& input_data, |
| 3375 | const DeviceMemory<float>& weights, |
| 3376 | const dnn::BatchDescriptor& input_dimensions, |
| 3377 | const dnn::BatchDescriptor& output_dimensions, |
| 3378 | DeviceMemory<float>* output_data) { |
| 3379 | if (input_dimensions.count() != output_dimensions.count()) { |
| 3380 | LOG(ERROR) << "MatMul input and output dimensions are not compatible."; |
| 3381 | return false; |
| 3382 | } |
| 3383 | |
| 3384 | // We do not permute the input or output, instead we just |
| 3385 | // reinterpret the layout. We are working with row-major matrices |
| 3386 | // and the rows of the input and output correspond to batch, so |
| 3387 | // batch has to be outermost in both the input and output. |
| 3388 | // |
| 3389 | // By adding transposes to the BLAS gemm call we could perhaps make |
| 3390 | // the kYXDepthBatch layout work as well, but there has been no need |
| 3391 | // for that so far. |
| 3392 | if (input_dimensions.layout() != dnn::DataLayout::kBatchYXDepth && |
| 3393 | input_dimensions.layout() != dnn::DataLayout::kBatchDepthYX) { |
| 3394 | LOG(ERROR) << "Unsupported MatMul input layout."; |
| 3395 | return false; |
| 3396 | } |
| 3397 | if (output_dimensions.layout() != dnn::DataLayout::kBatchYXDepth && |
| 3398 | output_dimensions.layout() != dnn::DataLayout::kBatchDepthYX) { |
| 3399 | LOG(ERROR) << "Unsupported MatMul output layout."; |
| 3400 | return false; |
| 3401 | } |
| 3402 | |
| 3403 | if (output_dimensions.width() == 1 && output_dimensions.height() == 1) { |
| 3404 | // This is a fast path that also supports the kBatchYXDepth layout. |
| 3405 | |
| 3406 | // The matrices here are in row-major format while BLAS expects |
| 3407 | // column-major, i.e. our matrices are transposed as far as BLAS |
| 3408 | // is concerned. So we need to compute output^T = |
| 3409 | // input^T*weights^T. There is no parameter for transposing the |
| 3410 | // output in BLAS gemm, but instead we can transpose both sides of |
| 3411 | // the equality to see that this is equivalent to |
| 3412 | // output=weights*input. So we only need to swap the order of |
| 3413 | // weights and input in the matrix product to correct for the |
| 3414 | // row-major versus column-major difference. |
| 3415 | const float alpha = 1.0f; // Take the matrix product without scaling it. |
| 3416 | const float beta = 0.0f; // Ignore the original values in output_data. |
| 3417 | const int64 m = output_dimensions.NodesAcrossFeatureMaps(); |
| 3418 | const int64 n = input_dimensions.count(); |
| 3419 | const int64 k = input_dimensions.NodesAcrossFeatureMaps(); |
| 3420 | stream->ThenBlasGemm(blas::Transpose::kNoTranspose, |
| 3421 | blas::Transpose::kNoTranspose, m, n, k, alpha, weights, |
| 3422 | m, input_data, k, beta, output_data, m); |
| 3423 | } else { |
| 3424 | // This is a slower and more complex path that supports output |
| 3425 | // width() * height() > 1, though it only supports the |
| 3426 | // kBatchYXDepth layout. Does support kBatchDepthYX if output |
| 3427 | // feature_map_count() == 1, as then there is no difference |
| 3428 | // between the two layouts. |
| 3429 | // |
| 3430 | // The operation here is the same as above, except that we have to |
no test coverage detected