text: /**
   * \file dnn/src/common/basic_types.cpp
   * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
   *
   * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   */
  
  #include "megdnn/basic_types.h"
  #include "megdnn/tensor_format.h"
  
  #include "src/common/utils.h"
  
  #include <array>
  #include <cstdlib>
  #include <cstring>
  #include <mutex>
  #include <numeric>
  #include <tuple>
  #include <type_traits>
  
  using namespace megdnn;
  
  /* ===================== ErrorHandler =====================  */
  namespace {
  class DefaultErrorHandler final : public ErrorHandler {
      void do_on_megdnn_error(const std::string& msg) override {
          megdnn_ignore(msg);
  #if MEGDNN_ENABLE_EXCEPTIONS
          throw std::runtime_error{msg};
  #else
          megdnn_trap();
  #endif
      }
  };
  
  template <typename T>
  void serialize_pod(const T& val, std::string& result) {
      static_assert(std::is_standard_layout<T>::value, "invalid type");
      result.append(reinterpret_cast<const char*>(&val), sizeof(T));
  }
  
  template <typename T>
  void serialize_vec(const T* val, size_t size, std::string& result) {
      result.append(reinterpret_cast<const char*>(val), sizeof(T) * size);
  }
  
  template <typename T>
  T deserialize_pod(const std::string& data, size_t& offset) {
      T ret;
      memcpy(&ret, data.data() + offset, sizeof(T));
      offset += sizeof(T);
      return ret;
  }
  
  }  // namespace
  ErrorHandler* ErrorHandler::sm_inst;
  
  ErrorHandler* ErrorHandler::inst() {
      static DNN_MUTEX mtx;
      static DefaultErrorHandler default_handler;
      if (megdnn_unlikely(!sm_inst)) {
          MEGDNN_LOCK_GUARD(mtx);
          if (!sm_inst) {
              sm_inst = &default_handler;
          }
      }
      return sm_inst;
  }
  
  void ErrorHandler::on_megdnn_error(const std::string& msg) {
      inst()->do_on_megdnn_error(msg);
  
      // gcc seems to fail to recognize the noreturn attr of
      // do_on_tensor_reshape_error; explicitly mark this function as noreturn
      // here
      megdnn_trap();
  }
  
  void ErrorHandler::on_megdnn_error(const char* msg) {
      on_megdnn_error(std::string{msg});
  }
  
  void ErrorHandler::on_tensor_reshape_error(const std::string& msg) {
      inst()->do_on_tensor_reshape_error(msg);
      megdnn_trap();
  }
  
  void ErrorHandler::on_tensor_reshape_error(const char* msg) {
      on_tensor_reshape_error(std::string{msg});
  }
  
  void ErrorHandler::set_handler(ErrorHandler* handler) {
      sm_inst = handler;
  }
  
  /* ===================== logging =====================  */
  
  namespace {
  LogHandler g_log_handler = nullptr;
  }  // anonymous namespace
  
  #if MEGDNN_ENABLE_LOGGING
  void megdnn::__log__(
          LogLevel level, const char* file, const char* func, int line, const char* fmt,
          ...) {
      if (!g_log_handler)
          return;
      va_list ap;
      va_start(ap, fmt);
      g_log_handler(level, file, func, line, fmt, ap);
      va_end(ap);
  }
  #endif  // MEGDNN_ENABLE_LOGGING
  
  LogHandler megdnn::set_log_handler(LogHandler handler) {
      auto ret = g_log_handler;
      g_log_handler = handler;
      return ret;
  }
  
  /* ===================== TensorShape =====================  */
  
  TensorShape::TensorShape(const SmallVector<size_t>& init_shape) {
      megdnn_assert(
              init_shape.size() <= MAX_NDIM,
              "Illegal to construct a TensorShape with "
              "more than MAX_NDIM(%zu) axes; init_shape is %s",
              MAX_NDIM, vec2str(init_shape).c_str());
      ndim = init_shape.size();
      memcpy(this->shape, init_shape.data(), sizeof(size_t) * ndim);
  }
  
  TensorShape::TensorShape(std::initializer_list<size_t> init_shape)
          : TensorShape(SmallVector<size_t>{init_shape}) {}
  
  size_t TensorShape::total_nr_elems() const {
      if (!ndim)
          return 0;
      return std::accumulate(shape, shape + ndim, 1_z, SafeMultiplies<size_t>());
  }
  
  bool TensorShape::eq_shape(const TensorShape& rhs) const {
      MEGDNN_STATIC_ASSERT(MAX_NDIM == 7, "please update the code")
      if (ndim == rhs.ndim) {
          size_t eq = 0;
          switch (ndim) {
              case 7:
                  eq += shape[6] == rhs.shape[6];
                  MEGDNN_FALLTHRU
              case 6:
                  eq += shape[5] == rhs.shape[5];
                  MEGDNN_FALLTHRU
              case 5:
                  eq += shape[4] == rhs.shape[4];
                  MEGDNN_FALLTHRU
              case 4:
                  eq += shape[3] == rhs.shape[3];
                  MEGDNN_FALLTHRU
              case 3:
                  eq += shape[2] == rhs.shape[2];
                  MEGDNN_FALLTHRU
              case 2:
                  eq += shape[1] == rhs.shape[1];
                  MEGDNN_FALLTHRU
              case 1:
                  eq += shape[0] == rhs.shape[0];
          }
          return eq == ndim;
      }
      return false;
  }
  
  std::string TensorShape::to_string() const {
      std::string rst("{");
      for (size_t i = 0; i < ndim; i++) {
          if (i)
              rst.append(",");
          rst.append(std::to_string(shape[i]));
      }
      rst.append("}");
      return rst;
  }
  
  bool TensorShape::is_empty() const {
      for (size_t i = 0; i < ndim; ++i) {
          if (!shape[i]) {
              return true;
          }
      }
      return ndim == 0;
  }
  
  /* ===================== TensorLayout =====================  */
  TensorLayout::TensorLayout() = default;
  
  TensorLayout::TensorLayout(DType dtype_) : dtype{dtype_}, format{Format(dtype)} {}
  
  TensorLayout::TensorLayout(DType dtype_, Format format_)
          : dtype{dtype_}, format{format_} {}
  
  TensorLayout::TensorLayout(const TensorShape& shape, DType dtype)
          : TensorLayout(shape, dtype, Format(dtype)) {}
  
  TensorLayout::TensorLayout(const TensorShape& shape, DType dtype, TensorFormat format_)
          : TensorShape(shape), dtype{dtype}, format{format_} {
      init_contiguous_stride();
  }
  
  TensorLayout::TensorLayout(
          const TensorShape& shape, const std::vector<ptrdiff_t>& stride, DType dtype)
          : TensorLayout(shape, stride, dtype, Format(dtype)) {}
  
  TensorLayout::TensorLayout(
          const TensorShape& shape, const std::vector<ptrdiff_t>& stride, DType dtype,
          TensorFormat format_)
          : TensorShape(shape), dtype{dtype}, format{format_} {
      megdnn_assert_eq_size_t(stride.size(), ndim);
      for (size_t i = 0; i < shape.ndim; ++i)
          this->stride[i] = stride[i];
  }
  
  size_t TensorLayout::init_contiguous_stride() {
      return format.impl()->init_contiguous_stride(*this);
  }
  
  size_t TensorLayout::init_contiguous_stride(const TensorShape& shape) {
      this->TensorShape::operator=(shape);
      return init_contiguous_stride();
  }
  
  size_t TensorLayout::init_contiguous_stride(
          const TensorShape& shape, TensorFormat format_) {
      this->TensorShape::operator=(shape);
      this->format = format_;
      return init_contiguous_stride();
  }
  
  TensorLayout TensorLayout::dimshuffle(const std::vector<size_t>& dims) const {
      TensorLayout res{dtype, format};
      res.ndim = this->ndim;
      megdnn_assert_eq_size_t(dims.size(), this->ndim);
      auto ndim = this->ndim;
      rep(i, ndim) {
          auto dest = dims[i];
          megdnn_assert(dest < ndim);
          res.shape[i] = this->shape[dest];
          res.stride[i] = this->stride[dest];
      }
      return res;
  }
  
  TensorLayout TensorLayout::remove_axis(size_t idx) const {
      TensorLayout res{*this};
      res.remove_axis_inplace(idx);
      return res;
  }
  
  void TensorLayout::remove_axis_inplace(size_t axis) {
      megdnn_assert(ndim >= 2 && axis < ndim);
      --ndim;
      for (size_t i = axis; i < ndim; ++i) {
          shape[i] = shape[i + 1];
          stride[i] = stride[i + 1];
      }
  }
  
  void TensorLayout::add_axis_inplace(size_t axis, size_t shape, ptrdiff_t stride) {
      megdnn_assert(
              ndim + 1 <= MAX_NDIM && axis <= ndim && shape,
              "can not add axis at %zu (current ndim %zu, MAX_NDIM %zu)", axis, ndim,
              MAX_NDIM);
      ndim++;
      for (size_t i = ndim - 1; i > axis; i--) {
          this->shape[i] = this->shape[i - 1];
          this->stride[i] = this->stride[i - 1];
      }
      this->shape[axis] = shape;
      this->stride[axis] = stride;
  }
  
  void TensorLayout::modify_dtype_inplace(DType dtype_) {
      dtype = dtype_;
      format = Format(dtype);
  }
  
  bool TensorLayout::is_contiguous() const {
      return format.impl()->is_contiguous_spec(*this);
  }
  
  bool TensorLayout::is_physical_contiguous() const {
      ptrdiff_t expected = 1;
      for (int i = ndim - 1; i >= 0; --i) {
          if (shape[i] != 1 && stride[i] != expected)
              return false;
          expected *= shape[i];
      }
      // empty tensors are not contiguous
      return expected != 0;
  }
  
  bool TensorLayout::is_abs_monotonous_allow_brdcst() const {
      if (!ndim)
          return false;
      if (ndim == 1)
          return true;
      ptrdiff_t last =
              std::abs(stride[ndim - 1]) * static_cast<ptrdiff_t>(shape[ndim - 1]);
      for (int i = ndim - 2; i >= 0; --i) {
          if (!stride[i] || shape[i] == 1)
              continue;
          if (std::abs(stride[i]) < last)
              return false;
          last = std::abs(stride[i]) * static_cast<ptrdiff_t>(shape[i]);
      }
      return true;
  }
  
  bool TensorLayout::is_contiguous_allow_brdcst() const {
      if (!ndim)
          return false;
      ptrdiff_t expected = 1;
      for (int i = ndim - 1; i >= 0; --i) {
          if (!stride[i])
              continue;
          if (shape[i] != 1 && stride[i] != expected)
              return false;
          expected *= shape[i];
      }
      // empty tensors are not contiguous
      return expected != 0;
  }
  
  /**
   * \brief The collapse_contiguous function will convert a contiguous image like
   * tensor layout into a 2-dimensional layout, shape[0] = height of the image,
   * shape[1] = width of the image, axis = 1, stride[0] = row_pitch_size_in_elem,
   * and stride[1] = 1.
   * So if the nhwcd4 format layout is transformed into a 2d tensor
   * layout after calling this function, the nhwcd4 format layout is contiguous.
   */
  TensorLayout TensorLayout::collapse_contiguous() const {
      return format.impl()->collapse_contiguous_spec(*this);
  }
  
  bool TensorLayout::is_non_overlapping_strong() const {
      // abs(stride), stride, shape
      std::array<std::tuple<ptrdiff_t, ptrdiff_t, size_t>, MAX_NDIM> vec;
      for (size_t i = 0; i < this->ndim; ++i) {
          vec[i] = std::make_tuple(std::abs(stride[i]), stride[i], shape[i]);
      }
      std::sort(vec.begin(), vec.begin() + this->ndim);
      ptrdiff_t lo = 0, hi = 0;
      for (size_t i = 0; i < this->ndim; ++i) {
          auto cur_stride = std::get<1>(vec[i]);
          auto cur_shape = std::get<2>(vec[i]);
          megdnn_assert(cur_shape > 0);
          if (cur_shape == 1)
              continue;
          if (cur_stride > 0) {
              if (cur_stride <= hi)
                  return false;
              hi += cur_stride * (cur_shape - 1);
          } else {
              // cur_stride == 0 is handled here, which causes returning false
              if (lo <= cur_stride)
                  return false;
              lo += cur_stride * (cur_shape - 1);
          }
      }
      return true;
  }
  
  bool TensorLayout::eq_layout(const TensorLayout& rhs) const {
      megdnn_assert(
              dtype == rhs.dtype,
              "could not compare layout on different dtypes: %s vs %s", dtype.name(),
              rhs.dtype.name());
      MEGDNN_STATIC_ASSERT(MAX_NDIM == 7, "please update the code")
  
      auto ax = [](size_t shape0, size_t shape1, ptrdiff_t stride0, ptrdiff_t stride1) {
          return (shape0 == shape1) & ((shape0 <= 1) | (stride0 == stride1));
      };
      if (ndim == rhs.ndim) {
          size_t eq = 0;
          switch (ndim) {
              case 7:
                  eq += ax(shape[6], rhs.shape[6], stride[6], rhs.stride[6]);
                  MEGDNN_FALLTHRU
              case 6:
                  eq += ax(shape[5], rhs.shape[5], stride[5], rhs.stride[5]);
                  MEGDNN_FALLTHRU
              case 5:
                  eq += ax(shape[4], rhs.shape[4], stride[4], rhs.stride[4]);
                  MEGDNN_FALLTHRU
              case 4:
                  eq += ax(shape[3], rhs.shape[3], stride[3], rhs.stride[3]);
                  MEGDNN_FALLTHRU
              case 3:
                  eq += ax(shape[2], rhs.shape[2], stride[2], rhs.stride[2]);
                  MEGDNN_FALLTHRU
              case 2:
                  eq += ax(shape[1], rhs.shape[1], stride[1], rhs.stride[1]);
                  MEGDNN_FALLTHRU
              case 1:
                  eq += ax(shape[0], rhs.shape[0], stride[0], rhs.stride[0]);
          }
          return eq == ndim;
      }
      return false;
  }
  
  TensorLayout::Span TensorLayout::span() const {
      return format.impl()->span_spec(*this);
  }
  
  size_t TensorLayout::access_bytes() const {
      megdnn_assert(dtype.valid());
      auto contig = collapse_contiguous();
      size_t ret = 0;
      if (dtype.is_low_bit()) {
          ret = 1;
          int align_size_in_elements = 8 / dtype.low_bit();
          auto min_stride = contig.stride[0];
          for (size_t i = 0; i < contig.ndim; ++i) {
              if (contig.stride[i] == 1) {
                  ret *= round_up((int)contig.shape[i], align_size_in_elements);
              } else {
                  ret *= contig.shape[i];
              }
              if (min_stride > contig.stride[i]) {
                  min_stride = contig.stride[i];
              }
          }
          if (min_stride != 1) {
              megdnn_assert(min_stride == align_size_in_elements);
              ret *= min_stride;
          }
          ret /= align_size_in_elements;
      } else {
          ret = dtype.size(total_nr_elems());
      }
      return ret;
  }
  
  TensorLayout TensorLayout::broadcast(const TensorShape& tshape) const {
      megdnn_throw_if(
              !ndim || !tshape.ndim, tensor_reshape_error,
              "broadcast involves empty tensor");
  
      if (is_scalar()) {
          TensorLayout result{dtype, format};
          result.ndim = tshape.ndim;
          for (size_t i = 0; i < tshape.ndim; i++) {
              result.shape[i] = tshape.shape[i];
              result.stride[i] = (tshape.shape[i] == 1);
          }
          return result;
      }
  
      megdnn_throw_if(
              tshape.ndim < ndim, tensor_reshape_error,
              ssprintf(
                      "dimension for broadcast less than "
                      "dst_shape: src_shape=%s dst_shape=%s",
                      to_string().c_str(), tshape.to_string().c_str()));
      TensorLayout result{dtype, format};
      for (size_t i = 0; i < tshape.ndim; ++i) {
          int target_idx = tshape.ndim - i - 1;
          int cur_idx = ndim - i - 1;
          size_t cur_shape = (cur_idx >= 0 ? shape[cur_idx] : 1),
                 cur_stride = (cur_idx >= 0 ? stride[cur_idx] : 0);
          if (tshape.shape[target_idx] != cur_shape) {
              megdnn_throw_if(
                      cur_shape != 1 && cur_stride != 0, tensor_reshape_error,
                      ssprintf(
                              "broadcast on dim with shape not equal to 1: "
                              "src_shape=%s dst_shape=%s",
                              to_string().c_str(), tshape.to_string().c_str()));
              result.shape[target_idx] = tshape.shape[target_idx];
              result.stride[target_idx] = 0;
          } else {
              result.shape[target_idx] = cur_shape;
              result.stride[target_idx] = cur_stride;
          }
      }
      result.ndim = tshape.ndim;
      return result;
  }
  
  bool TensorLayout::try_reshape(TensorLayout& result, const TensorShape& tshp) const {
      megdnn_assert(tshp.ndim);
  
      bool is_empty_shape = false;
      for (size_t i = 0; i < tshp.ndim; ++i) {
          if (!tshp.shape[i]) {
              megdnn_throw_if(
                      !format.is_default(), tensor_reshape_error,
                      ssprintf("bad target tshp: %s", tshp.to_string().c_str()));
              is_empty_shape = true;
              break;
          }
      }
  
      megdnn_throw_if(
              !tshp.ndim || total_nr_elems() != tshp.total_nr_elems(),
              tensor_reshape_error,
              ssprintf(
                      "number of elements do not match "
                      "in reshape: src=%s dest=%s",
                      static_cast<const TensorShape&>(*this).to_string().c_str(),
                      tshp.to_string().c_str()));
  
      auto cont = collapse_contiguous();
      result.dtype = this->dtype;
      result.format = this->format;
      result.TensorShape::operator=(tshp);
  
      if (is_empty_shape) {
          result.init_contiguous_stride();
          return true;
      }
  
      size_t sdim = 0, prod = 1, cont_sdim = 0;
      for (size_t i = 0; i < tshp.ndim; ++i) {
          megdnn_assert(cont_sdim < cont.ndim);
          prod *= result.shape[i];
          if (prod > cont.shape[cont_sdim])
              return false;
  
          if (prod == cont.shape[cont_sdim] &&
              (i + 1 >= tshp.ndim || tshp.shape[i + 1] != 1)) {
              auto s = cont.stride[cont_sdim];
              for (int j = i; j >= static_cast<int>(sdim); --j) {
                  result.stride[j] = s;
                  s *= result.shape[j];
              }
              ++cont_sdim;
              sdim = i + 1;
              prod = 1;
          }
      }
      megdnn_assert(cont_sdim == cont.ndim);
  
      return true;
  }
  
  TensorLayout TensorLayout::reshape(const TensorShape& shape) const {
      TensorLayout ret;
      auto succ = try_reshape(ret, shape);
      megdnn_throw_if(
              !succ, tensor_reshape_error,
              ssprintf(
                      "can not reshape from %s to %s", to_string().c_str(),
                      shape.to_string().c_str()));
      return ret;
  }
  
  std::string TensorLayout::to_string() const {
      std::string rst("{");
      for (size_t i = 0; i < ndim; i++) {
          if (i)
              rst.append(",");
          rst.append(std::to_string(shape[i]));
  
          rst.push_back('(');
          rst.append(std::to_string(stride[i]));
          rst.push_back(')');
      }
      if (format.type() != Format::Type::DEFAULT) {
          rst.append(" @ ");
          rst.append(format.impl()->to_string());
      }
      rst.append(std::string(" ") + dtype.name());
      rst.append("}");
      return rst;
  }
  
  std::string TensorLayout::serialize() const {
      std::string rst;
      serialize_pod<size_t>(ndim, rst);
      serialize_vec<size_t>(shape, ndim, rst);
      serialize_vec<ptrdiff_t>(stride, ndim, rst);
      rst.append(format.impl()->to_string());
  
      //! serialize dtype
      serialize_pod(dtype.enumv(), rst);
      if (dtype.has_param()) {
          switch (dtype.enumv()) {
  #define cb(_dt)                                                       \
      case DTypeTrait<dtype::_dt>::enumv:                               \
          serialize_pod(dtype::_dt::downcast_from(dtype).param(), rst); \
          break;
              MEGDNN_FOREACH_PARAMETERIZED_DTYPE(cb)
  #undef cb
              default:
                  megdnn_assert(false, "cannot serialize unknown parameterized DType");
                  break;
          }
      }
  
      return rst;
  }
  
  void RefPtr::reset(const void* ptr, size_t offset) {
      megdnn_assert(m_mutable, "this RefPtr can't change.");
      *m_ref = const_cast<void*>(ptr);
      m_offset = offset;
  }
  
  void TensorND::reset_ptr(void* ptr, size_t offset) {
      m_ref_ptr.reset(ptr, offset);
  }
  
  // vim: syntax=cpp.doxygen
  