compute/cker/include/cker/eigen/eigen_convolution_helpers.h

   1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
   2
   3 Licensed under the Apache License, Version 2.0 (the "License");
   4 you may not use this file except in compliance with the License.
   5 You may obtain a copy of the License at
   6
   7     http://www.apache.org/licenses/LICENSE-2.0
   8
   9 Unless required by applicable law or agreed to in writing, software
  10 distributed under the License is distributed on an "AS IS" BASIS,
  11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 See the License for the specific language governing permissions and
  13 limitations under the License.
  14 ==============================================================================*/
  15
  16 #ifndef __NNFW_CKER_EIGEN_EIGEN_CONVOLUTION_HELPERS_H__
  17 #define __NNFW_CKER_EIGEN_EIGEN_CONVOLUTION_HELPERS_H__
  18
  19 namespace Eigen
  20 {
  21 namespace internal
  22 {
  23
  24 // TensorEvaluatorHasPartialPacket<TensorEvaluatorType, PacketType, IndexType>
  25 // provides `value` that is true if TensorEvaluatorType has `PacketType
  26 // partialPacket<PacketType>(IndexType, unpacket_traits<PacketType>::mask_t)
  27 // const` and if the PacketType supports masked load.
  28 //
  29 // Partial packets are used to:
  30 //
  31 // 1) Split the packet over two columns in eigen based spatial convolution and
  32 // use partial loads for each individual part before combining them to get the
  33 // required packet. This class is used to pick the correct implementation of
  34 // loadPacketStandard function.
  35 //
  36 // 2) Split the packet over two rows (within the same column) in eigen based
  37 // cuboid convolution and use partial loads for each individual part before
  38 // combining them to get the required packet. This class is used to pick the
  39 // correct implementation of loadPacketStandard function. This usage is similar
  40 // to the usage in eigen based spatial convolution described above.
  41 //
  42 // 3) Finalize packing of columns in gemm_pack_colmajor after processing
  43 //    vectorized part with full packets (see eigen_spatial_convolutions.h).
  44 template <typename TensorEvaluatorType, typename PacketType, typename IndexType>
  45 class TensorEvaluatorHasPartialPacket
  46 {
  47 public:
  48   template <typename TensorEvaluatorT, typename PacketT, typename IndexT>
  49   static auto functionExistsSfinae(
  50       typename std::enable_if<
  51           unpacket_traits<PacketT>::masked_load_available &&
  52           std::is_same<
  53               PacketT,
  54               decltype(std::declval<const TensorEvaluatorT>().template partialPacket<PacketT>(
  55                   std::declval<IndexT>(),
  56                   std::declval<typename unpacket_traits<PacketT>::mask_t>()))>::value>::type *)
  57       -> std::true_type;
  58
  59   template <typename TensorEvaluatorT, typename PacketT, typename IndexT>
  60   static auto functionExistsSfinae(...) -> std::false_type;
  61
  62   typedef decltype(
  63       functionExistsSfinae<TensorEvaluatorType, PacketType, IndexType>(nullptr)) status;
  64
  65   static constexpr bool value = status::value;
  66 };
  67
  68 // Compute a mask for loading/storing coefficients in/from a packet in a
  69 // [from, to) range. If the mask bit is 1, element will be loaded/stored.
  70 template <typename Packet>
  71 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE
  72     typename std::enable_if<unpacket_traits<Packet>::masked_load_available,
  73                             typename unpacket_traits<Packet>::mask_t>::type
  74     mask(int from, int to)
  75 {
  76   const Index packet_size = internal::unpacket_traits<Packet>::size;
  77   eigen_assert(0 <= from && to <= (packet_size + 1) && from < to);
  78
  79   using Mask = typename internal::unpacket_traits<Packet>::mask_t;
  80   const Mask mask_max = std::numeric_limits<Mask>::max();
  81
  82   return (mask_max >> (packet_size - to)) ^ (mask_max >> (packet_size - from));
  83 }
  84
  85 } // namespace internal
  86 } // namespace Eigen
  87
  88 #endif // __NNFW_CKER_EIGEN_EIGEN_CONVOLUTION_HELPERS_H__