compute/cker/include/cker/operation/reference/BatchMatMul.h

   1 /*
   2  * Copyright (c) 2020 Samsung Electronics Co., Ltd. All Rights Reserved
   3  * Copyright 2020 The TensorFlow Authors. All Rights Reserved.
   4  *
   5  * Licensed under the Apache License, Version 2.0 (the "License");
   6  * you may not use this file except in compliance with the License.
   7  * You may obtain a copy of the License at
   8  *
   9  *      http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 #ifndef __NNFW_CKER_REFERENCE_BATCH_MATMUL_H__
  19 #define __NNFW_CKER_REFERENCE_BATCH_MATMUL_H__
  20
  21 #include "cker/Types.h"
  22 #include "cker/Shape.h"
  23
  24 namespace nnfw
  25 {
  26 namespace cker
  27 {
  28 namespace reference
  29 {
  30
  31 inline void BatchMatMul(const Shape &lhs_shape, const float *lhs_data, const Shape &rhs_shape,
  32                         const float *rhs_data, const Shape &, float *output_data)
  33 {
  34   const Shape extended_lhs_shape = Shape::ExtendedShape(5, lhs_shape);
  35   const Shape extended_rhs_shape = Shape::ExtendedShape(5, rhs_shape);
  36
  37   // Determine which dimension is the broadcast dimension.
  38   auto broadcast_dim = [](int lhs_dim, int rhs_dim) {
  39     if (lhs_dim == rhs_dim)
  40       return lhs_dim;
  41     if (lhs_dim == 1)
  42       return rhs_dim;
  43     assert(rhs_dim == 1);
  44     return lhs_dim;
  45   };
  46
  47   // Compute the "extent" for iterating on this dimension.
  48   // If we are broadcasting, then don't advance (i.e return 0).
  49   auto extent = [](const Shape &shape, int x) {
  50     if (shape.Dims(x) == 1)
  51     {
  52       return 0;
  53     }
  54     int prod = 1;
  55     for (int i = x + 1; i < shape.DimensionsCount(); ++i)
  56     {
  57       prod *= shape.Dims(i);
  58     }
  59     return prod;
  60   };
  61
  62   const int batch_dim0 = broadcast_dim(extended_lhs_shape.Dims(0), extended_rhs_shape.Dims(0));
  63   const int batch_dim1 = broadcast_dim(extended_lhs_shape.Dims(1), extended_rhs_shape.Dims(1));
  64   const int batch_dim2 = broadcast_dim(extended_lhs_shape.Dims(2), extended_rhs_shape.Dims(2));
  65
  66   const int lhs_ext0 = extent(extended_lhs_shape, 0);
  67   const int lhs_ext1 = extent(extended_lhs_shape, 1);
  68   const int lhs_ext2 = extent(extended_lhs_shape, 2);
  69   const int rhs_ext0 = extent(extended_rhs_shape, 0);
  70   const int rhs_ext1 = extent(extended_rhs_shape, 1);
  71   const int rhs_ext2 = extent(extended_rhs_shape, 2);
  72
  73   // Set params for each matrix multiply.
  74   const int lhs_rows = extended_lhs_shape.Dims(3);
  75   const int rhs_cols = extended_rhs_shape.Dims(4);
  76   const int accum_depth = extended_lhs_shape.Dims(4);
  77
  78   for (int b0 = 0; b0 < batch_dim0; ++b0)
  79   {
  80     const float *lhs_ptr0 = lhs_data + (b0 * lhs_ext0);
  81     const float *rhs_ptr0 = rhs_data + (b0 * rhs_ext0);
  82     for (int b1 = 0; b1 < batch_dim1; ++b1)
  83     {
  84       const float *lhs_ptr1 = lhs_ptr0 + b1 * lhs_ext1;
  85       const float *rhs_ptr1 = rhs_ptr0 + b1 * rhs_ext1;
  86       for (int b2 = 0; b2 < batch_dim2; ++b2)
  87       {
  88         const float *lhs_ptr2 = lhs_ptr1 + b2 * lhs_ext2;
  89         const float *rhs_ptr2 = rhs_ptr1 + b2 * rhs_ext2;
  90         float *out_ptr =
  91             output_data +
  92             ((b0 * batch_dim1 * batch_dim2) + b1 * batch_dim2 + b2) * lhs_rows * rhs_cols;
  93         for (int j = 0; j < rhs_cols; ++j)
  94         {
  95           for (int i = 0; i < lhs_rows; ++i)
  96           {
  97             float total = 0.f;
  98             for (int k = 0; k < accum_depth; ++k)
  99             {
 100               total += lhs_ptr2[accum_depth * i + k] * rhs_ptr2[j * accum_depth + k];
 101             }
 102             int idx = lhs_rows * j + i;
 103             out_ptr[idx] = total;
 104           }
 105         }
 106       }
 107     }
 108   }
 109 }
 110
 111 } // namespace reference
 112 } // namespace cker
 113 } // namespace nnfw
 114
 115 #endif // __NNFW_CKER_REFERENCE_BATCH_MATMUL_H__