// Copyright (c) 2020 Graphcore Ltd. All rights reserved.

#include "popops/GatherStatistics.hpp"
#include "HistogramPerformanceEstimation.hpp"
#include "poplibs_support/Tracepoint.hpp"
#include "popops/Cast.hpp"
#include "popops/ElementWise.hpp"
#include "popops/Reduce.hpp"
#include "poputil/OptionParsing.hpp"
#include "poputil/Util.hpp"
#include "poputil/VertexTemplates.hpp"
#include <poplibs_support/logging.hpp>
#include <poputil/TileMapping.hpp>
#include <poputil/exceptions.hpp>

using namespace poplar;
using namespace poplar::program;
using namespace poputil;
using namespace poplibs_support;

namespace popops {

namespace {

struct HistogramOptions {
  bool useFloatArithmetic = false;
  bool useFloatArithmeticWithUnsignedIntOutput = false;
};

static void validateOptions(const HistogramOptions &opts,
                            bool useFloatSupported) {
  if (opts.useFloatArithmetic) {
    if (!useFloatSupported) {
      throw poplar::invalid_option(
          "Option `useFloatArithmetic` is not supported for this function,"
          " but option `useFloatArithmeticWithUnsignedIntOutput` is"
          " supported");
    } else if (opts.useFloatArithmeticWithUnsignedIntOutput) {
      throw poplar::invalid_option(
          "Both options `useFloatArithmetic`"
          " and `useFloatArithmeticWithUnsignedIntOutput` must not be switched"
          " ON at the same time");
    }
    logging::popops::warn(
        "The `useFloatArithmetic` option is deprecated."
        " Consider using the `useFloatArithmeticWithUnsignedIntOutput` option"
        " to use float arithmetic internally.");
  }
}

HistogramOptions parseOptionFlags(const OptionFlags &options,
                                  bool useFloatSupported) {
  HistogramOptions histogramOpts;
  const poplibs::OptionSpec histogramSpec{
      {"useFloatArithmetic", poplibs::OptionHandler::createWithBool(
                                 histogramOpts.useFloatArithmetic)},
      {"useFloatArithmeticWithUnsignedIntOutput",
       poplibs::OptionHandler::createWithBool(
           histogramOpts.useFloatArithmeticWithUnsignedIntOutput)},
  };
  for (const auto &entry : options) {
    histogramSpec.parse(entry.first, entry.second);
  }
  validateOptions(histogramOpts, useFloatSupported);
  return histogramOpts;
}

enum class VertexType {
  MULTIVERTEX_1D_BY_DATA,
  MULTIVERTEX_1D_BY_LIMIT,
  WORKER_2D
};

VertexType chooseVertexType(const std::vector<std::vector<Interval>> &intervals,
                            unsigned max1DElemsByLimit,
                            unsigned max1DElemsByData, bool isAbsolute,
                            bool isHalf, unsigned numWorkers,
                            unsigned numLimits, unsigned vectorWidth) {
  if (intervals.size() != 1 || intervals[0].size() > 1) {
    return VertexType::WORKER_2D;
  }
  const auto elements = std::accumulate(
      intervals[0].begin(), intervals[0].end(), 0u,
      [](std::size_t total, const Interval &i) { return total + i.size(); });

  // We can use one of the MultiVertex 1D types, which each have different
  // advantages:
  // MULTIVERTEX_1D_BY_LIMIT divides work using numLimits which can leave idle
  //                     workers, but there is no need to have 1 worker combine
  //                     partial results, so the vertex is simpler and faster.
  // MULTIVERTEX_1D_BY_DATA  divides work by data but requires each piece to be
  //                     recombined by one worker.  This makes it more complex.

  auto limitsCycles = histogram1DByLimitEstimate(
      elements, numLimits + 1, isAbsolute, isHalf, numWorkers, vectorWidth);
  auto dataCycles = histogram1DByDataEstimate(
      elements, numLimits + 1, isAbsolute, isHalf, numWorkers, vectorWidth);

  if (elements < max1DElemsByData && dataCycles < limitsCycles) {
    return VertexType::MULTIVERTEX_1D_BY_DATA;
  }
  if (elements < max1DElemsByLimit && dataCycles >= limitsCycles) {
    return VertexType::MULTIVERTEX_1D_BY_LIMIT;
  }

  return VertexType::WORKER_2D;
}

poplar::Tensor histogramImpl(poplar::Graph &graph, const poplar::Tensor &input,
                             const poplar::Tensor &levels, bool absoluteOfInput,
                             poplar::program::Sequence &prog,
                             const DebugNameAndId &dnai) {
  const auto &target = graph.getTarget();
  const auto numTiles = target.getNumTiles();
  const auto inType = input.elementType();
  const auto vectorWidth = target.getVectorWidth(inType);
  const auto numWorkers = target.getNumWorkerContexts();
  auto flattenedInput = input.flatten();

  const auto cs = graph.addComputeSet({dnai, "Histogram"});

  // As work is split by "limits" not data size, the rpt count is a severe
  // limitation on the input size when rptMax is small.  So in that case the
  // vertex overcomes the limitation.  In other cases we must split work by
  // producing more vertices
  const auto rptMax = target.getRptCountMax();
  const auto max1DElemsByLimit =
      (rptMax < 0xffff) ? UINT32_MAX : rptMax * vectorWidth;
  const auto max1DElemsByData = rptMax * numWorkers * vectorWidth;

  const auto codeletName2D =
      templateVertex("popops::Histogram2D", inType, absoluteOfInput);
  const auto max2DVectorElems =
      std::min<std::size_t>(graph.getMaxFieldDim(codeletName2D, "data", 1),
                            target.getRptCountMax() * vectorWidth);

  // Gather a vector of results from each vertex created
  std::vector<Tensor> results;
  graph.reorderToSimplify(&flattenedInput, {}, false);
  const auto mapping = graph.getTileMapping(flattenedInput);
  for (unsigned tile = 0; tile < numTiles; tile++) {
    if (mapping[tile].empty())
      continue;
    const auto tileContiguousRegions =
        graph.getSortedContiguousRegions(flattenedInput, mapping[tile]);

    if (tileContiguousRegions.size() == 0) {
      // No data on this tile
      continue;
    }
    const auto vertexType =
        chooseVertexType(tileContiguousRegions, max1DElemsByLimit,
                         max1DElemsByData, absoluteOfInput, inType == HALF,
                         numWorkers, levels.numElements(), vectorWidth);
    if (vertexType == VertexType::MULTIVERTEX_1D_BY_LIMIT ||
        vertexType == VertexType::MULTIVERTEX_1D_BY_DATA) {
      // 1 region of suitable size to use a MultiVertex
      const auto byLimit = (vertexType == VertexType::MULTIVERTEX_1D_BY_LIMIT);
      const auto codeletName1D = templateVertex("popops::Histogram1D", inType,
                                                absoluteOfInput, byLimit);
      auto v = graph.addVertex(cs, codeletName1D);
      graph.setTileMapping(v, tile);

      auto resultSize = levels.numElements() + 1;
      auto histogram = graph.addVariable(
          FLOAT, {byLimit ? resultSize : numWorkers * resultSize}, {dnai});
      results.push_back(histogram.slice(0, resultSize));
      graph.setTileMapping(histogram, tile);

      graph.setInitialValue(v["histogramCount"], results.back().numElements());
      graph.connect(v["limits"], levels);
      graph.connect(v["histogram"], histogram);
      graph.connect(v["data"],
                    flattenedInput.slice(tileContiguousRegions[0][0]));
    } else {
      // > 1 region or not suitable size so use 2D worker vertices
      auto vertexRegions = splitRegionsBetweenWorkers(
          target, tileContiguousRegions, vectorWidth, 2 * vectorWidth,
          UINT32_MAX, max2DVectorElems);
      for (const auto &regions : vertexRegions) {
        auto v = graph.addVertex(cs, codeletName2D);
        graph.setTileMapping(v, tile);
        results.push_back(
            graph.addVariable(FLOAT, {levels.numElements() + 1}, {dnai}));
        graph.setTileMapping(results.back(), tile);

        graph.setInitialValue(v["histogramCount"],
                              results.back().numElements());
        graph.connect(v["limits"], levels);
        graph.connect(v["histogram"], results.back());
        graph.connect(v["data"], flattenedInput.slices(regions));
      }
    }
  }
  prog.add(Execute(cs, {dnai}));

  return concat(results).reshape({results.size(), results[0].numElements()});
}

} // anonymous namespace

// The above function makes histograms of on-tile data using float format for
// speed. Float can represent exact integers in the range 0-16777216 (Note 1)
// but after that not every value can be represented.  This is not a problem
// on tile, as tile data < 16M * (size of half) = 32MBytes. But the data
// spread over the whole IPU can be more than 32MBytes.
// We need an accurate integer count in our histogram so we may need to use
// int data (Note 2), but at the cost of slower processing so cast only as
// necessary.
//
// Note 1: float can represent integers between +/-16777216 exactly but that
// won't totally solve the problem.  At best we could reduce with a start
// value of -16777216 and remove that bias after casting to int, doubling our
// range before int reduction is required.
//
// Note 2: int, not unsigned int as the reduction library doesn't support
// unsigned ints.  Additionally it doesn't have optimised vertices
// for int (whereas it does for float). int supports
// 2Giga entries (4GBytes of half data)
//
// The option is provided to override the above and do all calculation in
// float arithmetic and produce a float result.
// The functions below do reduction with the appropriate casting and
// data type to combine the vertex results.
//
// In short, if all elements are to be counted into a single histogram entry
// and we want an answer that is exact, we can only use float data for our
// arithmetic if there are <= `maxElementsForFloatReduction` values at any
// point in the calculation (So on 1 tile, float is OK).  Otherwise
// we must use unsigned or int values.
constexpr unsigned maxElementsForFloatReduction = 16777216u;

poplar::Tensor histogram(poplar::Graph &graph, const poplar::Tensor &input,
                         const poplar::Tensor &levels, bool absoluteOfInput,
                         poplar::program::Sequence &prog,
                         const poplar::DebugContext &debugContext,
                         const poplar::OptionFlags &options) {
  POPOPS_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext, DI_ARGS(input, levels, absoluteOfInput, options));

  const auto opts = parseOptionFlags(options, true);
  auto histogramResult =
      histogramImpl(graph, input, levels, absoluteOfInput, prog, {di});

  if (opts.useFloatArithmeticWithUnsignedIntOutput || opts.useFloatArithmetic) {
    // Override all concerns over inaccurate integer representation as float,
    // but tolerate possible inaccurate results
    auto output = reduce(graph, histogramResult, FLOAT, {0},
                         popops::Operation::ADD, prog, {di});
    auto output_ = output;
    if (opts.useFloatArithmeticWithUnsignedIntOutput) {
      output_ = cast(graph, output, UNSIGNED_INT, prog, {di});
    }
    di.addOutput(output_);
    return output_;
  }
  // See the above explanation on numeric limits for exact integer
  // representation using float vs unsigned
  if (input.numElements() > maxElementsForFloatReduction) {
    // Reduce as INT, cast to unsigned to return
    histogramResult = cast(graph, histogramResult, INT, prog, {di});
  }
  // Reduce, cast to unsigned to return
  auto output = reduce(graph, histogramResult, histogramResult.elementType(),
                       {0}, popops::Operation::ADD, prog, {di});
  // When casting int to unsigned this appears to generate nothing
  auto output_ = cast(graph, output, UNSIGNED_INT, prog, {di});
  di.addOutput(output_);
  return output_;
}

void histogram(poplar::Graph &graph, const poplar::Tensor &input,
               poplar::Tensor &output, bool updateOutput,
               const poplar::Tensor &levels, bool absoluteOfInput,
               poplar::program::Sequence &prog,
               const poplar::DebugContext &debugContext,
               const poplar::OptionFlags &options) {
  POPOPS_TRACEPOINT();
  poputil::PoplibsOpDebugInfo di(
      debugContext,
      DI_ARGS(input, output, levels, updateOutput, absoluteOfInput, options));

  const auto opts = parseOptionFlags(options, false);
  if (opts.useFloatArithmeticWithUnsignedIntOutput &&
      output.elementType() != poplar::UNSIGNED_INT) {
    throw poputil::poplibs_error(
        "histogram output tensor is expected to be of type"
        " UNSIGNED_INT but " +
        output.elementType().toString() + " found");
  }

  auto histogramResult =
      histogramImpl(graph, input, levels, absoluteOfInput, prog, {di});

  // This output type checking condition is deprecated in favour of determining
  // the arithmetic based on the option
  // `useFloatArithmeticWithUnsignedIntOutput`.
  if (output.elementType() == FLOAT) {
    // Override all concerns over inaccurate integer representation as float,
    // but tolerate possible inaccurate results
    logging::popops::warn(
        "Passing an output tensor of type FLOAT to the"
        " histogram function is deprecated! Consider using the"
        " `useFloatArithmeticWithUnsignedIntOutput` option to use float"
        " arithmetic internally.");
    reduceWithOutput(graph, histogramResult, output, {0},
                     {popops::Operation::ADD, updateOutput}, prog, {di});
    return;
  } else if (opts.useFloatArithmeticWithUnsignedIntOutput) {
    if (updateOutput) {
      auto outFloat = graph.clone(FLOAT, output);
      reduceWithOutput(graph, histogramResult, outFloat, {0},
                       popops::Operation::ADD, prog, {di});
      mapInPlace(graph, expr::Cast(expr::_2, UNSIGNED_INT) + expr::_1,
                 {output, outFloat}, prog, {di});
    } else {
      auto output_ = reduce(graph, histogramResult, FLOAT, {0},
                            popops::Operation::ADD, prog, {di});
      output = cast(graph, output_, UNSIGNED_INT, prog, {di});
    }
    return;
  }

  // See the above explanation on numeric limits for exact integer
  // representation using float vs unsigned
  if (input.numElements() > maxElementsForFloatReduction) {
    // Reduce as INT, cast to unsigned to return
    output = cast(graph, output, INT, prog, {di});
    histogramResult = cast(graph, histogramResult, INT, prog, {di});
    reduceWithOutput(graph, histogramResult, output, {0},
                     {popops::Operation::ADD, updateOutput}, prog, {di});
    output = cast(graph, output, UNSIGNED_INT, prog, {di});
  } else {
    // Reduce as float
    auto result = reduce(graph, histogramResult, FLOAT, {0},
                         popops::Operation::ADD, prog, {di});
    if (updateOutput) {
      // Cast to unsigned and add to the result to return
      result = cast(graph, result, UNSIGNED_INT, prog, {di});
      addInPlace(graph, output, result, prog, {di});
    } else {
      // Cast to unsigned to return
      output = cast(graph, result, UNSIGNED_INT, prog, {di});
    }
  }
}

} // namespace popops