MCPcopy Create free account
hub / github.com/ARM-software/ComputeLibrary / gemm

Function gemm

tests/validation/reference/GEMM.cpp:40–121  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

38{
39template <typename T, typename std::enable_if<is_floating_point<T>::value, int>::type>
40SimpleTensor<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 {

Callers 4

gemm_accumulateFunction · 0.70
compute_referenceMethod · 0.50
compute_referenceMethod · 0.50

Calls 7

to_bf_precisionFunction · 0.85
shapeMethod · 0.45
data_typeMethod · 0.45
yMethod · 0.45
xMethod · 0.45
zMethod · 0.45
num_dimensionsMethod · 0.45

Tested by

no test coverage detected