mapping_options.h

Go to the documentation of this file.

#pragma once

#include <mapping_options.pb.h>

#include <array>
#include <iostream>
#include <string>
#include <vector>

#include "tc/external/isl.h"

#include "tc/core/flags.h"
#include "tc/core/rtc.h"


namespace tc {

template <typename T>
class ValueAccessor {
 public:
  using Setter = std::function<void(T)>;
  using Getter = std::function<T()>;

  ValueAccessor(const Setter& s, const Getter& g) : setter_(s), getter_(g) {}
  ValueAccessor(const ValueAccessor&) = default;

  operator T() const {
    return getter_();
  }

  ValueAccessor& operator=(const T& t) {
    setter_(t);
    return *this;
  }

 private:
  Setter setter_;
  Getter getter_;
};

class CudaDimView {
 private:
  CudaDimView() = default;

 public:
  CudaDimView(const CudaDimView&) = default;
  explicit CudaDimView(CudaDimProto& buf) : proto(buf) {}

  inline size_t size() const;

  inline std::vector<uint64_t> extractVector() const;

  inline std::array<uint64_t, 3> extractDefaultedArray() const;

  inline ValueAccessor<uint64_t> operator[](size_t i);

  inline uint64_t operator[](size_t i) const;

  inline CudaDimView& operator=(const CudaDimView& view);

  inline bool operator==(const CudaDimView& view) const;
  inline bool operator!=(const CudaDimView& view) const;

  std::string toCommaSeparatedString() const;
  friend std::ostream& operator<<(std::ostream& os, const CudaDimView& view);

 public:
  CudaDimProto& proto;

  static const uint64_t defaultDim = 1;
};

class CudaDim : public CudaDimView {
 public:
  CudaDim() : CudaDimView(ownedProto_) {}
  CudaDim(const CudaDim& cudaDim)
      : CudaDimView(ownedProto_), ownedProto_(cudaDim.proto) {}
  CudaDim(const CudaDimProto& proto)
      : CudaDimView(ownedProto_), ownedProto_(proto) {}
  CudaDim(const CudaDimView& view)
      : CudaDimView(ownedProto_), ownedProto_(view.proto) {}
  inline CudaDim(std::initializer_list<uint64_t> il);
  inline CudaDim(std::vector<uint64_t> il);
  inline CudaDim(
      uint64_t x,
      uint64_t y = CudaDimView::defaultDim,
      uint64_t z = CudaDimView::defaultDim);

  using CudaDimView::operator=;

 private:
  CudaDimProto ownedProto_;
};

class Block : public CudaDim {
 public:
  Block() = default;
  Block(const CudaDimView& view) : CudaDim(view.proto) {}
  Block(const CudaDimProto& proto) : CudaDim(proto) {}
  Block(std::initializer_list<uint64_t> il) : CudaDim(il) {}
  Block(std::vector<uint64_t> il) : CudaDim(il) {}

  using CudaDimView::operator=;
};

class Grid : public CudaDim {
 public:
  Grid() = default;
  Grid(const CudaDimView& view) : CudaDim(view.proto) {}
  Grid(const CudaDimProto& proto) : CudaDim(proto) {}
  Grid(std::initializer_list<uint64_t> il) : CudaDim(il) {}
  Grid(std::vector<uint64_t> il) : CudaDim(il) {}

  using CudaDimView::operator=;
};

class TilingView {
 private:
  TilingView() = default;

 public:
  TilingView(const TilingView&) = default;
  explicit TilingView(TilingProto& p) : proto(p) {}

  inline std::vector<uint64_t> extractVector() const;

  inline size_t size() const;

  inline ValueAccessor<uint64_t> operator[](size_t i);

  inline uint64_t operator[](size_t i) const;

  inline TilingView& operator=(const TilingView& view);

  inline bool operator==(const TilingView& view) const;
  inline bool operator!=(const TilingView& view) const;

  std::string toCommaSeparatedString() const;
  friend std::ostream& operator<<(std::ostream& os, const TilingView& view);

 public:
  TilingProto& proto;
};

class Tiling : public TilingView {
 public:
  Tiling() : TilingView(ownedProto_) {}
  Tiling(const Tiling& t) : TilingView(ownedProto_), ownedProto_(t.proto) {}
  Tiling(const TilingProto& proto)
      : TilingView(ownedProto_), ownedProto_(proto) {}
  Tiling(const TilingView& view)
      : TilingView(ownedProto_), ownedProto_(view.proto) {}
  inline Tiling(std::initializer_list<uint64_t> il);
  inline Tiling(const std::vector<uint64_t>& sizes);

 private:
  TilingProto ownedProto_;
};

class SchedulerOptionsView {
 public:
  using MergeCallback = std::function<
      isl_bool(isl_union_map*, isl_union_map*, int, int, int, void*)>;
  using ConstraintsCallback = std::function<isl_basic_set*(
      isl_basic_set*,
      int,
      int,
      isl_id_list*,
      int*,
      int*,
      void*)>;

 private:
  SchedulerOptionsView() = default;

 public:
  SchedulerOptionsView(const SchedulerOptionsView&) = default;
  SchedulerOptionsView(SchedulerOptionsProto& buf) : proto(buf) {}

  inline SchedulerOptionsView& operator=(const SchedulerOptionsView&);

  inline bool operator==(const SchedulerOptionsView& view) const;
  inline bool operator!=(const SchedulerOptionsView& view) const;

  friend std::ostream& operator<<(
      std::ostream& os,
      const SchedulerOptionsView& options);

 public:
  SchedulerOptionsProto& proto;
};

class SchedulerOptions : public SchedulerOptionsView {
 public:
  SchedulerOptions() : SchedulerOptionsView(ownedProto_) {}
  SchedulerOptions(const SchedulerOptions& options)
      : SchedulerOptionsView(ownedProto_), ownedProto_(options.proto) {}
  explicit SchedulerOptions(const SchedulerOptionsProto& proto)
      : SchedulerOptionsView(ownedProto_), ownedProto_(proto) {}
  explicit SchedulerOptions(const SchedulerOptionsView& view)
      : SchedulerOptionsView(ownedProto_), ownedProto_(view.proto) {}

 private:
  SchedulerOptionsProto ownedProto_;
};

class MappingOptions {
 private:
  inline MappingOptions();
  static MappingOptions makeUnmappedMappingOptions();

 public:
  inline MappingOptions(const MappingOptions& options);
  inline explicit MappingOptions(const MappingOptionsProto& buf);

  inline explicit MappingOptions(const std::string& str);

  MappingOptions& operator=(const MappingOptions&) = default;

  inline bool operator==(const MappingOptions& options) const;
  inline bool operator!=(const MappingOptions& options) const;

  inline std::string toProtobufSerializedString() const;

  inline MappingOptions& tile(const std::vector<uint64_t>& sizes);
  inline MappingOptions& tile(std::initializer_list<uint64_t> sizes);
  MappingOptions& tile(const std::string& commaSeparatedSizes);
  inline MappingOptions& tile(const char* commaSeparatedSizes);
  template <typename... Args>
  MappingOptions& tile(Args...);

  inline MappingOptions& mapToThreads(std::initializer_list<uint64_t> threads);
  inline MappingOptions& mapToThreads(
      uint64_t x,
      uint64_t y = CudaDimView::defaultDim,
      uint64_t z = CudaDimView::defaultDim);
  inline MappingOptions& mapToThreads(const std::vector<uint64_t>& threads);
  MappingOptions& mapToThreads(const std::string& commaSeparatedSizes);

  inline MappingOptions& mapToBlocks(std::initializer_list<uint64_t> blocks);
  inline MappingOptions& mapToBlocks(
      uint64_t x,
      uint64_t y = CudaDimView::defaultDim,
      uint64_t z = CudaDimView::defaultDim);
  inline MappingOptions& mapToBlocks(const std::vector<uint64_t>& blocks);
  MappingOptions& mapToBlocks(const std::string& commaSeparatedSizes);

  inline MappingOptions& unroll(uint64_t size);

  inline MappingOptions& useSharedMemory(bool b);
  inline MappingOptions& usePrivateMemory(bool b);
  inline MappingOptions& maxSharedMemory(uint64_t size);
  inline MappingOptions& fixParametersBeforeScheduling(bool b);
  inline MappingOptions& unrollCopyShared(bool b);
  inline MappingOptions& tileImperfectlyNested(bool b);
  inline MappingOptions& matchLibraryCalls(bool b);

  inline MappingOptions& scheduleFusionStrategy(FusionStrategy fs);
  inline MappingOptions& scheduleFusionStrategy(const std::string& str);

  inline MappingOptions& outerScheduleFusionStrategy(FusionStrategy fs);
  inline MappingOptions& outerScheduleFusionStrategy(const std::string& str);
  inline MappingOptions& outerScheduleAllowSkewing(bool b);
  inline MappingOptions& outerSchedulePositiveOrthant(bool b);

  inline MappingOptions& intraTileScheduleFusionStrategy(FusionStrategy fs);
  inline MappingOptions& intraTileScheduleFusionStrategy(
      const std::string& str);
  inline MappingOptions& intraTileScheduleAllowSkewing(bool b);
  inline MappingOptions& intraTileSchedulePositiveOrthant(bool b);

  static MappingOptions makeNaiveMappingOptions();
  static MappingOptions makeSingleThreadMappingOptions();
  static MappingOptions makePointwiseMappingOptions();
  static MappingOptions makeMlpMappingOptions();
  static MappingOptions makeConvolutionMappingOptions();
  static MappingOptions makeGroupConvolutionMappingOptions();

  friend std::ostream& operator<<(
      std::ostream& os,
      const MappingOptions& options);

 public:
  MappingOptionsProto proto;

  // Views of sub-messages.
  CudaDimView block;
  CudaDimView grid;
  TilingView tiling;
  SchedulerOptionsView outerScheduleOptions;
  SchedulerOptionsView intraTileScheduleOptions;
};

namespace callbacks {
__isl_give isl_basic_set* AddPositiveCoefficientConstraints(
    __isl_take isl_basic_set* lp,
    int n_param,
    int dim,
    __isl_keep isl_id_list* stmt_ids,
    int* node_n_params,
    int* node_n_dims,
    void*);

isl_bool FuseAllPreserve3Coincident(
    __isl_take isl_union_map* original_schedule,
    __isl_take isl_union_map* updated_schedule,
    int n_updated_coincident,
    int n_original_coincident,
    int is_along_edge,
    void*);

isl_bool FuseAll(
    __isl_take isl_union_map* original_schedule,
    __isl_take isl_union_map* updated_schedule,
    int n_updated_coincident,
    int n_original_coincident,
    int is_along_edge,
    void*);

isl_bool FuseNone(
    __isl_take isl_union_map* original_schedule,
    __isl_take isl_union_map* updated_schedule,
    int n_updated_coincident,
    int n_original_coincident,
    int is_along_edge,
    void*);
} // namespace callbacks

} // namespace tc

#include "tc/core/mapping_options-inl.h"