/* Copyright 2019 The TensorFlow Authors. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ #ifndef TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_ #define TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_ #include <cstdint> #include "ruy/profiler/instrumentation.h" // from @ruy #include "tensorflow/lite/kernels/cpu_backend_context.h" #include "tensorflow/lite/kernels/cpu_backend_gemm_custom_gemv.h" #include "tensorflow/lite/kernels/cpu_backend_gemm_params.h" #include "tensorflow/lite/kernels/cpu_backend_gemm_ruy.h" #ifndef TFLITE_WITH_RUY #include "tensorflow/lite/kernels/cpu_backend_gemm_eigen.h" #include "tensorflow/lite/kernels/cpu_backend_gemm_gemmlowp.h" #include "tensorflow/lite/kernels/cpu_backend_gemm_x86.h" #endif namespace tflite { namespace cpu_backend_gemm { // The main entry point for CpuBackendGemm::Gemm. // // If TFLITE_WITH_RUY is set, CpuBackendGemm::Gemm will always go to Ruy aka // GemmImplUsingRuy. Other cases are as follows: // // |Quantized (uint8)|Quantized (int8)| Float | // TFLITE_WITH_RUY | Ruy | Ruy | Ruy | // !TFLITE_WITH_RUY | gemmlowp | Ruy/gemmlowp* | eigen | // * - Ruy if NEON is not available. // On x86 platforms: // (default) | gemmlowp | Ruy | eigen | // TFLITE_X86_RUY_\ | Ruy | Ruy | Ruy | // ENABLED && (AVX // or above available) #if !defined(TFLITE_WITH_RUY) && defined(TFLITE_X86_PLATFORM) /* GEMM dispatch implementation for x86. */ template <typename LhsScalar, typename RhsScalar, typename AccumScalar, typename DstScalar, QuantizationFlavor quantization_flavor> struct GemmImpl : detail::GemmImplX86<LhsScalar, RhsScalar, AccumScalar, DstScalar, quantization_flavor> {}; #else /* Generic implementation using ruy. * Non-ruy implementation will be partial specializations of this template. */ template <typename LhsScalar, typename RhsScalar, typename AccumScalar, typename DstScalar, QuantizationFlavor quantization_flavor> struct GemmImpl : detail::GemmImplUsingRuy<LhsScalar, RhsScalar, AccumScalar, DstScalar, quantization_flavor> { … }; #if !defined(TFLITE_WITH_RUY) /* Specializations using gemmlowp */ template <typename SrcScalar, typename DstScalar, QuantizationFlavor quantization_flavor> struct GemmImpl<SrcScalar, SrcScalar, std::int32_t, DstScalar, quantization_flavor> : detail::GemmImplUsingGemmlowp<SrcScalar, SrcScalar, std::int32_t, DstScalar, quantization_flavor> {}; // When SrcScalar=int8 or DstScalar=int8, gemmlowp fails to compile // outside of NEON. We avoid the compilation failure by subspecializing these // cases, rerouting it back to ruy. #if !defined(GEMMLOWP_NEON) template <typename SrcScalar, QuantizationFlavor quantization_flavor> struct GemmImpl<SrcScalar, SrcScalar, std::int32_t, std::int8_t, quantization_flavor> : detail::GemmImplUsingRuy<SrcScalar, SrcScalar, std::int32_t, std::int8_t, quantization_flavor> {}; template <typename DstScalar, QuantizationFlavor quantization_flavor> struct GemmImpl<std::int8_t, std::int8_t, std::int32_t, DstScalar, quantization_flavor> : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t, DstScalar, quantization_flavor> {}; template <QuantizationFlavor quantization_flavor> struct GemmImpl<std::int8_t, std::int8_t, std::int32_t, std::int8_t, quantization_flavor> : detail::GemmImplUsingRuy<std::int8_t, std::int8_t, std::int32_t, std::int8_t, quantization_flavor> {}; #endif // not GEMMLOWP_NEON /* Specializations using Eigen */ template <> struct GemmImpl<float, float, float, float, QuantizationFlavor::kFloatingPoint> : detail::GemmImplUsingEigen {}; #endif // not TFLITE_WITH_RUY #endif // not TFLITE_WITH_RUY and TFLITE_X86_PLATFORM /* Public entry point */ template <typename LhsScalar, typename RhsScalar, typename AccumScalar, typename DstScalar, QuantizationFlavor quantization_flavor> void Gemm(const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data, const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data, const MatrixParams<DstScalar>& dst_params, DstScalar* dst_data, const GemmParams<AccumScalar, DstScalar, quantization_flavor>& params, CpuBackendContext* context) { … } // Special path for 16x8 quant gemm. template <QuantizationFlavor quantization_flavor> void Gemm(const MatrixParams<int8_t>& lhs_params, const int8_t* lhs_data, const MatrixParams<int16_t>& rhs_params, const int16_t* rhs_data, const MatrixParams<int16_t>& dst_params, int16_t* dst_data, const GemmParams<int32_t, int16_t, quantization_flavor>& params, CpuBackendContext* context) { … } // Special path for gemm with raw accumulator case. i.e. AccumScalar == // DstScalar == int32 case. template <typename LhsScalar, typename RhsScalar, QuantizationFlavor quantization_flavor> void Gemm(const MatrixParams<LhsScalar>& lhs_params, const LhsScalar* lhs_data, const MatrixParams<RhsScalar>& rhs_params, const RhsScalar* rhs_data, const MatrixParams<int32_t>& dst_params, int32_t* dst_data, const GemmParams<int32_t, int32_t, quantization_flavor>& params, CpuBackendContext* context) { … } } // namespace cpu_backend_gemm } // namespace tflite #endif // TENSORFLOW_LITE_KERNELS_CPU_BACKEND_GEMM_H_
Codebase Browser
chromium