VertexIntrospector.hpp

// Copyright (c) 2017 Graphcore Ltd. All rights reserved.
 
#ifndef poplar_VertexIntrospector_hpp
#define poplar_VertexIntrospector_hpp
 
#include <poplar/GraphElements.hpp>
#include <poplar/Target.hpp>
#include <poplar/TypeTraits.hpp>
#include <poplar/VectorLayout.hpp>
#include <poplar/exceptions.hpp>
 
#include <limits>
#include <memory>
#include <string>
#include <vector>
 
namespace poplar {
namespace core {
class VertexIntrospector;
class FieldData;
} // namespace core
 
class FieldData {
public:
  struct SizeT {
    std::size_t value;
    std::size_t size() const { return value; }
  };
 
  FieldData(const FieldData &) = delete;
  FieldData(FieldData &&) noexcept;
  virtual ~FieldData();
 
  unsigned rank() const;
 
  std::size_t size() const;
 
  std::size_t getSizeAtIndex(std::size_t i) const;
 
  layout::Vector getProfilerVectorLayout(std::size_t nestingLevel) const;
 
  layout::VectorList getProfilerVectorListLayout() const;
 
  SizeT operator[](std::size_t i) const { return SizeT{getSizeAtIndex(i)}; }
 
  std::string name() const;
 
  template <typename T> T getInitialValue(const Target &target) const {
    static_assert(TypeTraits::isSimpleType<T>(),
                  "getInitialValue() called with invalid type");
 
    if (rank() != 0) {
      throw poplar_error("You may not call getInitialValue() on a non-scalar "
                         "(use getInitalValues() - with an 's')");
    }
    T result;
    const auto traits = TypeTraits::make<T>();
    getInitialValues(target, reinterpret_cast<void *>(&result), traits);
    return result;
  }
 
  template <typename T>
  std::vector<T> getInitialValues(const Target &target) const {
    if (rank() == 0) {
      throw poplar_error("You may not call getInitialValues() on a scalar "
                         "(use getInitalValue() - no 's')");
    }
 
    std::vector<T> result(size());
 
    // Resolve whether this function was called
    // with `T := SCALAR` or `T := std::vector<SCALAR>`:
    getInitialValuesOverload(target, result);
 
    return result;
  }
 
  // Implementation
  FieldData(std::unique_ptr<core::FieldData> fd);
 
private:
  std::unique_ptr<core::FieldData> impl;
 
  // This overload retrieves the initial value of any field with one index.
  template <typename T>
  void getInitialValuesOverload(const Target &target,
                                std::vector<T> &result) const {
    static_assert(TypeTraits::isSimpleType<T>(),
                  "getInitialValues() called with invalid type");
 
    if (rank() != 1) {
      throw poplar_error("getInitialValues() must be called with a "
                         "std::vector<T> for 1D fields");
    }
 
    auto traits = TypeTraits::make<T>();
    getInitialValues(target, result.data(), traits);
  }
 
  // This overload retrieves the initial values of any field with
  // two indices, for example, Vector<Input<Vector<scalar>>>
  template <typename T>
  void getInitialValuesOverload(const Target &target,
                                std::vector<std::vector<T>> &result) const {
    static_assert(TypeTraits::isSimpleType<T>(),
                  "getInitialValues() called with invalid type");
 
    if (rank() != 2) {
      throw poplar_error("getInitialValues() must be called with a "
                         "std::vector<std::vector<T>> for 1D fields");
    }
 
    auto traits = TypeTraits::make<T>();
    for (auto i = 0u; i < result.size(); ++i) {
      result[i] = std::vector<T>(getSizeAtIndex(i));
      // It is not an error to return empty vectors here
      // (for example, some tests in poplin generate empty worklists).
      if (!result[i].empty()) {
        getInitialValues(target, result[i].data(), traits, i);
      }
    }
  }
 
  void getInitialValues(
      const Target &target, void *dst, const TypeTraits &traits,
      std::size_t index = std::numeric_limits<std::size_t>::max()) const;
};
 
class VertexIntrospector {
public:
  FieldData getFieldInfo(const std::string &name) const;
 
  ComputeSet getComputeSet() const;
 
  // Implementation
  VertexIntrospector(std::unique_ptr<core::VertexIntrospector> impl);
  VertexIntrospector(VertexIntrospector &&) noexcept;
  const core::VertexIntrospector &getImpl() const { return *impl; }
 
private:
  std::unique_ptr<core::VertexIntrospector> impl;
};
 
} // end namespace poplar
 
#endif // poplar_VertexIntrospector_hpp