arrow/util/dispatch.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 #pragma once
0019
0020 #include <utility>
0021 #include <vector>
0022
0023 #include "arrow/status.h"
0024 #include "arrow/util/cpu_info.h"
0025
0026 namespace arrow {
0027 namespace internal {
0028
0029 enum class DispatchLevel : int {
0030   // These dispatch levels, corresponding to instruction set features,
0031   // are sorted in increasing order of preference.
0032   NONE = 0,
0033   SSE4_2,
0034   AVX2,
0035   AVX512,
0036   NEON,
0037   MAX
0038 };
0039
0040 /*
0041   A facility for dynamic dispatch according to available DispatchLevel.
0042
0043   Typical use:
0044
0045     static void my_function_default(...);
0046     static void my_function_avx2(...);
0047
0048     struct MyDynamicFunction {
0049       using FunctionType = decltype(&my_function_default);
0050
0051       static std::vector<std::pair<DispatchLevel, FunctionType>> implementations() {
0052         return {
0053           { DispatchLevel::NONE, my_function_default }
0054     #if defined(ARROW_HAVE_RUNTIME_AVX2)
0055           , { DispatchLevel::AVX2, my_function_avx2 }
0056     #endif
0057         };
0058       }
0059     };
0060
0061     void my_function(...) {
0062       static DynamicDispatch<MyDynamicFunction> dispatch;
0063       return dispatch.func(...);
0064     }
0065 */
0066 template <typename DynamicFunction>
0067 class DynamicDispatch {
0068  protected:
0069   using FunctionType = typename DynamicFunction::FunctionType;
0070   using Implementation = std::pair<DispatchLevel, FunctionType>;
0071
0072  public:
0073   DynamicDispatch() { Resolve(DynamicFunction::implementations()); }
0074
0075   FunctionType func = {};
0076
0077  protected:
0078   // Use the Implementation with the highest DispatchLevel
0079   void Resolve(const std::vector<Implementation>& implementations) {
0080     Implementation cur{DispatchLevel::NONE, {}};
0081
0082     for (const auto& impl : implementations) {
0083       if (impl.first >= cur.first && IsSupported(impl.first)) {
0084         // Higher (or same) level than current
0085         cur = impl;
0086       }
0087     }
0088
0089     if (!cur.second) {
0090       Status::Invalid("No appropriate implementation found").Abort();
0091     }
0092     func = cur.second;
0093   }
0094
0095  private:
0096   bool IsSupported(DispatchLevel level) const {
0097     static const auto cpu_info = arrow::internal::CpuInfo::GetInstance();
0098
0099     switch (level) {
0100       case DispatchLevel::NONE:
0101         return true;
0102       case DispatchLevel::SSE4_2:
0103         return cpu_info->IsSupported(CpuInfo::SSE4_2);
0104       case DispatchLevel::AVX2:
0105         return cpu_info->IsSupported(CpuInfo::AVX2);
0106       case DispatchLevel::AVX512:
0107         return cpu_info->IsSupported(CpuInfo::AVX512);
0108       default:
0109         return false;
0110     }
0111   }
0112 };
0113
0114 }  // namespace internal
0115 }  // namespace arrow