Class DispatchKeySet

paddle/phi/api/include/compat/c10/core/DispatchKeySet.h:54–394

}

class DispatchKeySet final {
 public:
  enum Full { FULL };
  enum FullAfter { FULL_AFTER };
  enum Raw { RAW };

  constexpr DispatchKeySet() = default;

  constexpr DispatchKeySet(Full /*unused*/)
      : repr_((1ULL << (num_backends + num_functionality_keys - 1)) - 1) {}

  constexpr DispatchKeySet(FullAfter /*unused*/, DispatchKey t)
      : repr_((1ULL << (num_backends +
                        static_cast<uint8_t>(toFunctionalityKey(t)) - 1)) -
              1) {
    *this = add(DispatchKey::PythonDispatcher);
  }

  constexpr DispatchKeySet(Raw /*unused*/, uint64_t x) : repr_(x) {}

  constexpr explicit DispatchKeySet(BackendComponent k) {
    if (k == BackendComponent::InvalidBit) {
      repr_ = 0;
    } else {
      repr_ = 1ULL << (static_cast<uint8_t>(k) - 1);
    }
  }

  constexpr explicit DispatchKeySet(DispatchKey k) {  // NOLINT
    if (k == DispatchKey::Undefined) {
      repr_ = 0;
    } else if (k <= DispatchKey::EndOfFunctionalityKeys) {
      uint64_t functionality_val =
          1ULL << (num_backends + static_cast<uint8_t>(k) - 1);
      repr_ = functionality_val;
    } else if (k <= DispatchKey::EndOfRuntimeBackendKeys) {
      auto functionality_k = toFunctionalityKey(k);
      uint64_t functionality_val =
          1ULL << (num_backends + static_cast<uint8_t>(functionality_k) - 1);

      auto backend_k = toBackendComponent(k);
      uint64_t backend_val = backend_k == BackendComponent::InvalidBit
                                 ? 0
                                 : 1ULL
                                       << (static_cast<uint8_t>(backend_k) - 1);
      repr_ = functionality_val + backend_val;
    } else {
      repr_ = 0;
    }
  }

  constexpr uint64_t keys_to_repr(std::initializer_list<DispatchKey> ks) {
    uint64_t repr = 0;
    for (auto k : ks) {
      repr |= DispatchKeySet(k).repr_;
    }
    return repr;
  }