| 38 | { |
| 39 | template <typename T, typename std::enable_if<is_floating_point<T>::value, int>::type> |
| 40 | SimpleTensor<T> gemm(const SimpleTensor<T> &a, |
| 41 | const SimpleTensor<T> &b, |
| 42 | const SimpleTensor<T> &c, |
| 43 | float alpha, |
| 44 | float beta, |
| 45 | bool fast_math) |
| 46 | { |
| 47 | // Create reference |
| 48 | SimpleTensor<T> dst{c.shape(), c.data_type(), 1}; |
| 49 | |
| 50 | // Compute reference |
| 51 | const int M = a.shape().y(); |
| 52 | const int N = b.shape().x(); |
| 53 | const int K = a.shape().x(); |
| 54 | const int D = a.shape().z(); // Number of matrices in a batch |
| 55 | const int W = a.shape()[3]; // Number of batched-gemm (Winograd case) |
| 56 | |
| 57 | const int a_stride_z = K * M; |
| 58 | const int a_stride_w = K * M * D; |
| 59 | |
| 60 | const int b_stride_z = |
| 61 | b.shape().num_dimensions() > 2 |
| 62 | ? N * K |
| 63 | : 0; // Do not slide the matrix B along the 3th dimension in case matrix B has less than 3 dimensions |
| 64 | int b_stride_w = |
| 65 | b.shape().num_dimensions() > 3 |
| 66 | ? K * N * D |
| 67 | : 0; // Do not slide the matrix B along the 4th dimension in case matrix B has less than 4 dimensions |
| 68 | |
| 69 | // Note: There are 3 gemm types: batched-gemm, multi-gemm, and batched of multi-gemms. The third dimension of tensor b is overloaded when tensor b has exactly 3 dimensions: |
| 70 | // it can be either number of batches or multis. Batched-GEMM computation is detected only when the third dimension of "a" and "c" tensors is 1 and the number of dimensions is 4 |
| 71 | const bool is_batched_gemm = b.shape().num_dimensions() == 3 && a.shape().num_dimensions() == 4 && |
| 72 | c.shape().num_dimensions() == 4 && a.shape()[2] == 1 && c.shape()[2] == 1; |
| 73 | |
| 74 | // Batched-GEMM |
| 75 | if (is_batched_gemm) |
| 76 | { |
| 77 | b_stride_w = b_stride_z; |
| 78 | } |
| 79 | |
| 80 | const int c_stride_z = N * M; |
| 81 | const int c_stride_w = N * M * D; |
| 82 | |
| 83 | #if defined(_OPENMP) && !(defined(__arm__) && defined(__ANDROID__)) |
| 84 | #pragma omp parallel for collapse(2) |
| 85 | #endif /* _OPENMP */ |
| 86 | for (int w = 0; w < W; ++w) |
| 87 | { |
| 88 | for (int depth = 0; depth < D; ++depth) |
| 89 | { |
| 90 | const int base_addr_a = depth * a_stride_z + w * a_stride_w; |
| 91 | const int base_addr_b = depth * b_stride_z + w * b_stride_w; |
| 92 | const int base_addr_c = depth * c_stride_z + w * c_stride_w; |
| 93 | |
| 94 | for (int row = 0; row < M; ++row) |
| 95 | { |
| 96 | for (int col = 0; col < N; ++col) |
| 97 | { |
no test coverage detected