// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2016 Igor Babuschkin <igor@babuschk.in>
//
// This Source Code Form is subject to the terms of the Mozilla
// Public License v. 2.0. If a copy of the MPL was not distributed
// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
// SPDX-License-Identifier: MPL-2.0

#ifndef EIGEN_TENSOR_TENSOR_SCAN_H
#define EIGEN_TENSOR_TENSOR_SCAN_H

// IWYU pragma: private
#include "./InternalHeaderCheck.h"

namespace Eigen {

namespace internal {

template <typename Op, typename XprType>
struct traits<TensorScanOp<Op, XprType> > : public traits<XprType> {
  typedef typename XprType::Scalar Scalar;
  typedef traits<XprType> XprTraits;
  typedef typename XprTraits::StorageKind StorageKind;
  static constexpr int NumDimensions = XprTraits::NumDimensions;
  static constexpr int Layout = XprTraits::Layout;
  typedef typename XprTraits::PointerType PointerType;
};

template <typename Op, typename XprType>
struct eval<TensorScanOp<Op, XprType>, Eigen::Dense> {
  typedef const TensorScanOp<Op, XprType>& type;
};

}  // end namespace internal

/**
 * \ingroup Tensor_Module
 *
 * \brief Tensor scan class.
 */
template <typename Op, typename XprType>
class TensorScanOp : public TensorBase<TensorScanOp<Op, XprType>, ReadOnlyAccessors> {
 public:
  typedef typename Eigen::internal::traits<TensorScanOp>::Scalar Scalar;
  typedef typename Eigen::NumTraits<Scalar>::Real RealScalar;
  typedef typename XprType::CoeffReturnType CoeffReturnType;
  typedef typename Eigen::internal::ref_selector<TensorScanOp>::type Nested;
  typedef typename Eigen::internal::traits<TensorScanOp>::StorageKind StorageKind;
  typedef typename Eigen::internal::traits<TensorScanOp>::Index Index;

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorScanOp(const XprType& expr, const Index& axis, bool exclusive = false,
                                                     const Op& op = Op())
      : m_expr(expr), m_axis(axis), m_accumulator(op), m_exclusive(exclusive) {}

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index axis() const { return m_axis; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const XprType& expression() const { return m_expr; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Op accumulator() const { return m_accumulator; }
  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool exclusive() const { return m_exclusive; }

 protected:
  typename XprType::Nested m_expr;
  const Index m_axis;
  const Op m_accumulator;
  const bool m_exclusive;
};

namespace internal {

template <typename Self>
EIGEN_STRONG_INLINE void ReduceScalar(Self& self, Index offset, typename Self::CoeffReturnType* data) {
  // Compute the scan along the axis, starting at the given offset
  typename Self::CoeffReturnType accum = self.accumulator().initialize();
  if (self.stride() == 1) {
    if (self.exclusive()) {
      for (Index curr = offset; curr < offset + self.size(); ++curr) {
        data[curr] = self.accumulator().finalize(accum);
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
      }
    } else {
      for (Index curr = offset; curr < offset + self.size(); ++curr) {
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
        data[curr] = self.accumulator().finalize(accum);
      }
    }
  } else {
    if (self.exclusive()) {
      for (Index idx3 = 0; idx3 < self.size(); idx3++) {
        Index curr = offset + idx3 * self.stride();
        data[curr] = self.accumulator().finalize(accum);
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
      }
    } else {
      for (Index idx3 = 0; idx3 < self.size(); idx3++) {
        Index curr = offset + idx3 * self.stride();
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
        data[curr] = self.accumulator().finalize(accum);
      }
    }
  }
}

template <typename Self>
EIGEN_STRONG_INLINE void ReducePacket(Self& self, Index offset, typename Self::CoeffReturnType* data) {
  using Scalar = typename Self::CoeffReturnType;
  using Packet = typename Self::PacketReturnType;
  // Compute the scan along the axis, starting at the calculated offset
  Packet accum = self.accumulator().template initializePacket<Packet>();
  if (self.stride() == 1) {
    if (self.exclusive()) {
      for (Index curr = offset; curr < offset + self.size(); ++curr) {
        internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
        self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
      }
    } else {
      for (Index curr = offset; curr < offset + self.size(); ++curr) {
        self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
        internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
      }
    }
  } else {
    if (self.exclusive()) {
      for (Index idx3 = 0; idx3 < self.size(); idx3++) {
        const Index curr = offset + idx3 * self.stride();
        internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
        self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
      }
    } else {
      for (Index idx3 = 0; idx3 < self.size(); idx3++) {
        const Index curr = offset + idx3 * self.stride();
        self.accumulator().reducePacket(self.inner().template packet<Unaligned>(curr), &accum);
        internal::pstoreu<Scalar, Packet>(data + curr, self.accumulator().finalizePacket(accum));
      }
    }
  }
}

template <typename Self, bool Vectorize, bool Parallel>
struct ReduceBlock {
  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1, typename Self::CoeffReturnType* data) const {
    for (Index idx2 = 0; idx2 < self.stride(); idx2++) {
      // Calculate the starting offset for the scan
      Index offset = idx1 + idx2;
      ReduceScalar(self, offset, data);
    }
  }
};

// Specialization for vectorized reduction.
template <typename Self>
struct ReduceBlock<Self, /*Vectorize=*/true, /*Parallel=*/false> {
  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1, typename Self::CoeffReturnType* data) const {
    using Packet = typename Self::PacketReturnType;
    const int PacketSize = internal::unpacket_traits<Packet>::size;
    Index idx2 = 0;
    for (; idx2 + PacketSize <= self.stride(); idx2 += PacketSize) {
      // Calculate the starting offset for the packet scan
      Index offset = idx1 + idx2;
      ReducePacket(self, offset, data);
    }
    for (; idx2 < self.stride(); idx2++) {
      // Calculate the starting offset for the scan
      Index offset = idx1 + idx2;
      ReduceScalar(self, offset, data);
    }
  }
};

// Single-threaded CPU implementation of scan
template <typename Self, typename Reducer, typename Device,
          bool Vectorize = (TensorEvaluator<typename Self::ChildTypeNoConst, Device>::PacketAccess &&
                            internal::reducer_traits<Reducer, Device>::PacketAccess)>
struct ScanLauncher {
  void operator()(Self& self, typename Self::CoeffReturnType* data) const {
    Index total_size = internal::array_prod(self.dimensions());

    // We fix the index along the scan axis to 0 and perform a
    // scan per remaining entry. The iteration is split into two nested
    // loops to avoid an integer division by keeping track of each idx1 and
    // idx2.
    for (Index idx1 = 0; idx1 < total_size; idx1 += self.stride() * self.size()) {
      ReduceBlock<Self, Vectorize, /*Parallel=*/false> block_reducer;
      block_reducer(self, idx1, data);
    }
  }
};

#ifdef EIGEN_USE_THREADS

// Adjust block_size to avoid false sharing of cachelines among
// threads. Currently set to twice the cache line size on Intel and ARM
// processors.
EIGEN_STRONG_INLINE Index AdjustBlockSize(Index item_size, Index block_size) {
  constexpr Index kBlockAlignment = 128;
  const Index items_per_cacheline = numext::maxi<Index>(1, kBlockAlignment / item_size);
  return items_per_cacheline * numext::div_ceil(block_size, items_per_cacheline);
}

template <typename Self>
struct ReduceBlock<Self, /*Vectorize=*/true, /*Parallel=*/true> {
  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1, typename Self::CoeffReturnType* data) const {
    using Scalar = typename Self::CoeffReturnType;
    using Packet = typename Self::PacketReturnType;
    const int PacketSize = internal::unpacket_traits<Packet>::size;
    Index num_scalars = self.stride();
    Index num_packets = 0;
    if (self.stride() >= PacketSize) {
      num_packets = self.stride() / PacketSize;
      self.device().parallelFor(
          num_packets,
          TensorOpCost(PacketSize * self.size(), PacketSize * self.size(), 16 * PacketSize * self.size(), true,
                       PacketSize),
          // Make the shard size large enough that two neighboring threads
          // won't write to the same cacheline of `data`.
          [=](Index blk_size) { return AdjustBlockSize(PacketSize * sizeof(Scalar), blk_size); },
          [&](Index first, Index last) {
            for (Index packet = first; packet < last; ++packet) {
              const Index idx2 = packet * PacketSize;
              ReducePacket(self, idx1 + idx2, data);
            }
          });
      num_scalars -= num_packets * PacketSize;
    }
    self.device().parallelFor(
        num_scalars, TensorOpCost(self.size(), self.size(), 16 * self.size()),
        // Make the shard size large enough that two neighboring threads
        // won't write to the same cacheline of `data`.
        [=](Index blk_size) { return AdjustBlockSize(sizeof(Scalar), blk_size); },
        [&](Index first, Index last) {
          for (Index scalar = first; scalar < last; ++scalar) {
            const Index idx2 = num_packets * PacketSize + scalar;
            ReduceScalar(self, idx1 + idx2, data);
          }
        });
  }
};

template <typename Self>
struct ReduceBlock<Self, /*Vectorize=*/false, /*Parallel=*/true> {
  EIGEN_STRONG_INLINE void operator()(Self& self, Index idx1, typename Self::CoeffReturnType* data) const {
    using Scalar = typename Self::CoeffReturnType;
    self.device().parallelFor(
        self.stride(), TensorOpCost(self.size(), self.size(), 16 * self.size()),
        // Make the shard size large enough that two neighboring threads
        // won't write to the same cacheline of `data`.
        [=](Index blk_size) { return AdjustBlockSize(sizeof(Scalar), blk_size); },
        [&](Index first, Index last) {
          for (Index idx2 = first; idx2 < last; ++idx2) {
            ReduceScalar(self, idx1 + idx2, data);
          }
        });
  }
};

// Specialization for multi-threaded execution.
template <typename Self, typename Reducer, bool Vectorize>
struct ScanLauncher<Self, Reducer, ThreadPoolDevice, Vectorize> {
  void operator()(Self& self, typename Self::CoeffReturnType* data) const {
    using Scalar = typename Self::CoeffReturnType;
    using Packet = typename Self::PacketReturnType;
    const int PacketSize = internal::unpacket_traits<Packet>::size;
    const Index total_size = internal::array_prod(self.dimensions());
    const Index inner_block_size = self.stride() * self.size();
    bool parallelize_by_outer_blocks = (total_size >= (self.stride() * inner_block_size));

    if ((parallelize_by_outer_blocks && total_size <= 4096) ||
        (!parallelize_by_outer_blocks && self.stride() < PacketSize)) {
      ScanLauncher<Self, Reducer, DefaultDevice, Vectorize> launcher;
      launcher(self, data);
      return;
    }

    if (parallelize_by_outer_blocks) {
      // Parallelize over outer blocks.
      const Index num_outer_blocks = total_size / inner_block_size;
      self.device().parallelFor(
          num_outer_blocks,
          TensorOpCost(inner_block_size, inner_block_size, 16 * PacketSize * inner_block_size, Vectorize, PacketSize),
          [=](Index blk_size) { return AdjustBlockSize(inner_block_size * sizeof(Scalar), blk_size); },
          [&](Index first, Index last) {
            for (Index idx1 = first; idx1 < last; ++idx1) {
              ReduceBlock<Self, Vectorize, /*Parallel=*/false> block_reducer;
              block_reducer(self, idx1 * inner_block_size, data);
            }
          });
    } else {
      // Parallelize over inner packets/scalars dimensions when the reduction
      // axis is not an inner dimension.
      ReduceBlock<Self, Vectorize, /*Parallel=*/true> block_reducer;
      for (Index idx1 = 0; idx1 < total_size; idx1 += self.stride() * self.size()) {
        block_reducer(self, idx1, data);
      }
    }
  }
};
#endif  // EIGEN_USE_THREADS

#if defined(EIGEN_USE_GPU) && (defined(EIGEN_GPUCC))

// GPU implementation of scan
// TODO(ibab): This placeholder implementation performs multiple scans in
// parallel, but it would be better to use a parallel scan algorithm and
// optimize memory access.
template <typename Self, typename Reducer>
__global__ EIGEN_HIP_LAUNCH_BOUNDS_1024 void ScanKernel(Self self, Index total_size,
                                                        typename Self::CoeffReturnType* data) {
  // Compute offset as in the CPU version
  Index val = threadIdx.x + blockIdx.x * blockDim.x;
  Index offset = (val / self.stride()) * self.stride() * self.size() + val % self.stride();

  if (offset + (self.size() - 1) * self.stride() < total_size) {
    // Compute the scan along the axis, starting at the calculated offset
    typename Self::CoeffReturnType accum = self.accumulator().initialize();
    for (Index idx = 0; idx < self.size(); idx++) {
      Index curr = offset + idx * self.stride();
      if (self.exclusive()) {
        data[curr] = self.accumulator().finalize(accum);
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
      } else {
        self.accumulator().reduce(self.inner().coeff(curr), &accum);
        data[curr] = self.accumulator().finalize(accum);
      }
    }
  }
  __syncthreads();
}

template <typename Self, typename Reducer, bool Vectorize>
struct ScanLauncher<Self, Reducer, GpuDevice, Vectorize> {
  void operator()(const Self& self, typename Self::CoeffReturnType* data) const {
    Index total_size = internal::array_prod(self.dimensions());
    Index num_blocks = (total_size / self.size() + 63) / 64;
    Index block_size = 64;

    LAUNCH_GPU_KERNEL((ScanKernel<Self, Reducer>), num_blocks, block_size, 0, self.device(), self, total_size, data);
  }
};
#endif  // EIGEN_USE_GPU && (EIGEN_GPUCC)

}  // namespace internal

// Eval as rvalue
template <typename Op, typename ArgType, typename Device>
struct TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> {
  typedef TensorScanOp<Op, ArgType> XprType;
  typedef typename XprType::Index Index;
  typedef const ArgType ChildTypeNoConst;
  typedef const ArgType ChildType;
  static constexpr int NumDims = internal::array_size<typename TensorEvaluator<ArgType, Device>::Dimensions>::value;
  typedef DSizes<Index, NumDims> Dimensions;
  typedef std::remove_const_t<typename XprType::Scalar> Scalar;
  typedef typename XprType::CoeffReturnType CoeffReturnType;
  typedef typename PacketType<CoeffReturnType, Device>::type PacketReturnType;
  typedef TensorEvaluator<const TensorScanOp<Op, ArgType>, Device> Self;
  typedef StorageMemory<Scalar, Device> Storage;
  typedef typename Storage::Type EvaluatorPointerType;

  static constexpr int Layout = TensorEvaluator<ArgType, Device>::Layout;
  enum {
    IsAligned = false,
    PacketAccess = (PacketType<CoeffReturnType, Device>::size > 1),
    // Scan eagerly materializes its result into m_output; once that buffer
    // exists, exposing block access is just a wrapper around it. Leave
    // PreferBlockAccess false so the executor still uses the cheaper
    // raw/packet paths by default; the flag matters only when an outer
    // expression calls block() directly.
    BlockAccess = (NumDims > 0),
    PreferBlockAccess = false,
    CoordAccess = false,
    RawAccess = true
  };

  //===- Tensor block evaluation strategy (see TensorBlock.h) -------------===//
  typedef internal::TensorBlockDescriptor<NumDims, Index> TensorBlockDesc;
  typedef internal::TensorBlockScratchAllocator<Device> TensorBlockScratch;
  typedef typename internal::TensorMaterializedBlock<Scalar, NumDims, Layout, Index> TensorBlock;
  //===--------------------------------------------------------------------===//

  EIGEN_STRONG_INLINE TensorEvaluator(const XprType& op, const Device& device)
      : m_impl(op.expression(), device),
        m_device(device),
        m_exclusive(op.exclusive()),
        m_accumulator(op.accumulator()),
        m_size(m_impl.dimensions()[op.axis()]),
        m_stride(1),
        m_consume_dim(op.axis()),
        m_output(nullptr) {
    // Accumulating a scalar isn't supported.
    EIGEN_STATIC_ASSERT((NumDims > 0), YOU_MADE_A_PROGRAMMING_MISTAKE);
    eigen_assert(op.axis() >= 0 && op.axis() < NumDims);

    // Compute stride of scan axis
    const Dimensions& dims = m_impl.dimensions();
    EIGEN_IF_CONSTEXPR (static_cast<int>(Layout) == static_cast<int>(ColMajor)) {
      for (int i = 0; i < op.axis(); ++i) {
        m_stride = m_stride * dims[i];
      }
    } else {
      // dims can only be indexed through unsigned integers,
      // so use an unsigned type to let the compiler know.
      // This prevents spurious warnings: "'*((void*)(& evaluator)+64)[18446744073709551615]' may be used uninitialized
      // in this function"
      unsigned int axis = internal::convert_index<unsigned int>(op.axis());
      for (unsigned int i = NumDims - 1; i > axis; --i) {
        m_stride = m_stride * dims[i];
      }
    }
  }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Dimensions& dimensions() const { return m_impl.dimensions(); }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& stride() const { return m_stride; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& consume_dim() const { return m_consume_dim; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Index& size() const { return m_size; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Op& accumulator() const { return m_accumulator; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE bool exclusive() const { return m_exclusive; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const TensorEvaluator<ArgType, Device>& inner() const { return m_impl; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE const Device& device() const { return m_device; }

  EIGEN_STRONG_INLINE bool evalSubExprsIfNeeded(EvaluatorPointerType data) {
    m_impl.evalSubExprsIfNeeded(nullptr);
    internal::ScanLauncher<Self, Op, Device> launcher;
    if (data) {
      launcher(*this, data);
      return false;
    }

    const Index total_size = internal::array_prod(dimensions());
    m_output =
        static_cast<EvaluatorPointerType>(m_device.get((Scalar*)m_device.allocate_temp(total_size * sizeof(Scalar))));
    launcher(*this, m_output);
    return true;
  }

  template <int LoadMode>
  EIGEN_DEVICE_FUNC PacketReturnType packet(Index index) const {
    return internal::ploadt<PacketReturnType, LoadMode>(m_output + index);
  }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE internal::TensorBlockResourceRequirements getResourceRequirements() const {
    return internal::TensorBlockResourceRequirements::any();
  }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorBlock block(TensorBlockDesc& desc, TensorBlockScratch& scratch,
                                                          bool /*root_of_expr_ast*/ = false) const {
    eigen_assert(m_output != nullptr);
    return TensorBlock::materialize(m_output, m_impl.dimensions(), desc, scratch);
  }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE EvaluatorPointerType data() const { return m_output; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE CoeffReturnType coeff(Index index) const { return m_output[index]; }

  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE TensorOpCost costPerCoeff(bool) const {
    return TensorOpCost(sizeof(CoeffReturnType), 0, 0);
  }

  EIGEN_STRONG_INLINE void cleanup() {
    if (m_output) {
      m_device.deallocate_temp(m_output);
      m_output = nullptr;
    }
    m_impl.cleanup();
  }

 protected:
  TensorEvaluator<ArgType, Device> m_impl;
  const Device EIGEN_DEVICE_REF m_device;
  const bool m_exclusive;
  Op m_accumulator;
  const Index m_size;
  Index m_stride;
  Index m_consume_dim;
  EvaluatorPointerType m_output;
};

}  // end namespace Eigen

#endif  // EIGEN_TENSOR_TENSOR_SCAN_H