// Copyright (c) 2016 Graphcore Ltd. All rights reserved.
#include "poplibs_test/TempDir.hpp"
#include "poputil/exceptions.hpp"
#include <algorithm>
#include <boost/multi_array.hpp>
#include <boost/optional.hpp>
#include <boost/optional/optional_io.hpp>
#include <boost/program_options.hpp>
#include <boost/random.hpp>
#include <chrono>
#include <exception>
#include <fstream>
#include <istream>
#include <optional>
#include <ostream>
#include <poplar/CycleCount.hpp>
#include <poplar/Engine.hpp>
#include <poplar/Graph.hpp>
#include <poplar/IPUModel.hpp>
#include <poplibs_support/TestDevice.hpp>
#include <poplibs_test/GeneralMatrixMultiply.hpp>
#include <poplibs_test/Pass.hpp>
#include <poplibs_test/SparseMatrix.hpp>
#include <poplibs_test/Util.hpp>
#include <poplin/codelets.hpp>
#include <popops/codelets.hpp>
#include <popsparse/SparsePartitioner.hpp>
#include <popsparse/codelets.hpp>
#include <queue>
#include <random>
#include <ratio>

#include "../lib/popsparse/FullyConnectedOnTile.hpp"
#include "../lib/popsparse/FullyConnectedOptions.hpp"
#include "../lib/popsparse/FullyConnectedPlan.hpp"
#include "../lib/popsparse/SparsePartitionerImpl.hpp"
#include "../lib/popsparse/SparseStorageInternal.hpp"
#include "poplibs_support/VectorUtils.hpp"
#include "poplibs_support/logging.hpp"
#include "popsparse/FullyConnected.hpp"
#include "popsparse/FullyConnectedParams.hpp"

#include <gccs/Algorithm.hpp>

using namespace poplar;
using namespace poplar::program;
using namespace poplibs_test::util;
using namespace poplar_test;
using namespace poputil;
using poplibs_test::Pass;
using namespace poplibs_support;

using namespace popsparse;
using namespace popsparse::dynamic;

// Tolerances used when data is not ignored
#define FLOAT_REL_TOL 0.01
#define HALF_REL_TOL 0.1

// TODO: This could be a program option if at all needed
using EType = float;

template <typename T>
static void logBucketStatistics(std::vector<PNBucket> &buckets,
                                const CSRMatrix<T> &csrMatrix) {
  if (buckets.empty()) {
    std::cerr << "   - No buckets found"
              << "\n";
    return;
  }

  std::size_t pnUsed = buckets.size();

  std::size_t maxNzElements = 0, maxMetaInfo = 0;
  std::size_t totalNzElements = 0, totalMetaInfo = 0;
  std::for_each(buckets.begin(), buckets.end(), [&](const PNBucket &b) {
    maxNzElements = std::max(maxNzElements, b.numNzElements);
    maxMetaInfo = std::max(maxMetaInfo, b.metaInfoElements);
    totalNzElements += b.numNzElements;
    totalMetaInfo += b.metaInfoElements;
  });

  std::cerr << "   - NZ entries " << csrMatrix.nzValues.size() << "\n";
  std::cerr << "   - NZ elements/PN max : " << maxNzElements;
  std::cerr << " avg : " << static_cast<double>(totalNzElements) / pnUsed;
  std::cerr << "\n";
  std::cerr << "   - Meta info elements/PN max : " << maxMetaInfo;
  std::cerr << " avg : " << static_cast<double>(totalMetaInfo) / pnUsed;
  std::cerr << "\n";
}

int main(int argc, char **argv) try {
  namespace po = boost::program_options;

  DeviceType deviceType = DeviceType::IpuModel2;
  unsigned numGroups = 1;
  unsigned inputSize;
  unsigned outputSize;
  unsigned batchSize;
  bool reportPlan;
  boost::optional<std::string> profileDir;
  Type dataType;
  Type partialsType;
  unsigned numIPUs = 1;
  boost::optional<unsigned> tilesPerIPU;
  Pass pass = Pass::ALL;
  std::string matmulOptionsString;
  std::string planConstraintsString;
  double sparsityFactor;
  ShapeOption<std::size_t> weightedAreaBegin;
  ShapeOption<std::size_t> weightedAreaEnd;
  ShapeOption<std::size_t> blockSize;
  weightedAreaBegin.val = weightedAreaEnd.val = {0, 0};
  double weightedAreaWeighting = 1.0;
  bool denseGradWSerialSplits = false;

  po::options_description desc("Options");
  // clang-format off
  desc.add_options()
    ("help", "Produce help message")
    ("compile-only", "Stop after compilation; don't run the program")
    ("device-type",
     po::value<DeviceType>(&deviceType)->default_value(deviceType),
     deviceTypeHelp)
    ("input-size", po::value<unsigned>(&inputSize)->required(),
     "Number of inputs")
    ("output-size", po::value<unsigned>(&outputSize)->required(),
     "Number of output channels")
    ("sparsity-factor", po::value<double>(&sparsityFactor)->required(),
     "Sparsity factor (ratio of number of non-zero values to total weight "
     "values")
    ("data-type",
     po::value<Type>(&dataType)->default_value(HALF),
     "Type of the input and output data")
    ("partials-type",
     po::value<Type>(&partialsType)->default_value(FLOAT),
     "Type of partials used during the operation")
    ("tiles-per-ipu",
     po::value(&tilesPerIPU),
     "Number of tiles per IPU")
    ("batch-size",
     po::value<unsigned>(&batchSize)->default_value(1),
     "Batch size")
    ("block-size",
     po::value<ShapeOption<std::size_t>>(&blockSize)->default_value(1),
     "Block size as rows and columns (only square blocks are supported)")
    ("single-phase",
     po::value<Pass>(&pass)->default_value(pass),
     "Run phase all | fwd | bwd | wu")
    ("ignore-data", "When set, no upload/download or verification of "
     "results is performed")
    ("plan-only", "Whether to perform planning only and skip creation "
     "and running of the program")
    ("profile", "Enable profiling and print profiling report")
    ("profile-dir",
     po::value<decltype(profileDir)>(&profileDir)
      ->default_value(boost::none),
     "Write profile files to the specified directory.")
    ("report-plan", po::value<bool>(&reportPlan)->default_value(false),
     "Display plan")
    ("report-total-cycle-counts", "Report total cycle count ignoring upload/download for "
     "each pass. Note not compatible with 'profile' option")
    ("variable-seed", "Use a variable seed based on clock, rather than a "
     "single fixed seed that does not change between runs of this tool")
    ("weighted-area-begin",
     po::value<ShapeOption<std::size_t>>(&weightedAreaBegin)->default_value(weightedAreaBegin),
     "Starting indices of an area of the sparse operand with a different "
     "level of sparsity to the rest")
    ("weighted-area-end",
     po::value<ShapeOption<std::size_t>>(&weightedAreaEnd)->default_value(weightedAreaEnd),
     "Ending indices of an area of the sparse operand with a different "
     "level of sparsity to the rest")
    ("weighted-area-weighting",
     po::value<double>(&weightedAreaWeighting)->default_value(weightedAreaWeighting),
     "Weighting for probability that a sparse element resides within the "
     "specified area")
    ("matmul-options", po::value<std::string>(&matmulOptionsString),
     "Options to use for the matrix multiplication, specified as a JSON "
     "string, e.g. {\"key\":\"value\"}")
    ("plan-constraints", po::value<std::string>(&planConstraintsString),
     "Plan constraints to use for the matrix multiplication, specified as "
     "a JSON string")
    ("report-dense-gradw-serial-splits",
      po::value<bool>(&denseGradWSerialSplits)->
        default_value(denseGradWSerialSplits),
     "Report dense GradW splits when GradW pass is enabled")
  ;
  // clang-format on
  po::variables_map vm;
  try {
    po::store(po::parse_command_line(argc, argv, desc), vm);
    if (vm.count("help")) {
      std::cout << desc << "\n";
      return 1;
    }
    po::notify(vm);
  } catch (std::exception &e) {
    std::cerr << "error: " << e.what() << "\n";
    return 1;
  }

  bool profile = vm.count("profile");
  bool profilingEnabled = profile || profileDir;
  bool reportTotalCycleCounts =
      vm.count("report-total-cycle-counts") && deviceType == DeviceType::Hw;
  bool ignoreData = vm.count("ignore-data");
  bool planOnly = vm.count("plan-only");
  bool variableSeed = vm.count("variable-seed");

  if (reportTotalCycleCounts && profilingEnabled) {
    throw poputil::poplibs_error(
        "--report-total-cycle-counts and --profile or --profile-dr specified "
        "at the same time. This is not allowed as one affects the other");
  }

  if (blockSize.val.size() > 2) {
    throw poputil::poplibs_error("Block size must be of dimension 2");
  }

  const std::size_t blockRows = blockSize[0];
  const std::size_t blockCols =
      blockSize.val.size() == 1 ? blockRows : blockSize[1];
  const auto blockArea = blockRows * blockCols;

  if (inputSize % blockRows) {
    throw poputil::poplibs_error("Input size must be an integer multiple of "
                                 "rows in a block");
  }

  if (outputSize % blockCols) {
    throw poputil::poplibs_error("output size must be an integer multiple of "
                                 "columns in a block");
  }

  if (weightedAreaBegin[0] > weightedAreaEnd[0] ||
      weightedAreaBegin[1] > weightedAreaEnd[1]) {
    std::stringstream ss;
    ss << "Invalid weighted area specified: " << weightedAreaBegin.val << ","
       << weightedAreaEnd.val;
    throw poputil::poplibs_error(ss.str());
  }

  if (weightedAreaEnd[0] > outputSize || weightedAreaEnd[1] > inputSize) {
    std::stringstream ss;
    ss << "Specified weighted area is out of bounds: Weighted area="
       << weightedAreaBegin.val << "," << weightedAreaEnd.val
       << " out of bounds {" << outputSize << "," << inputSize << "}";
    throw poputil::poplibs_error(ss.str());
  }

  // align weighted area to a block size grid
  weightedAreaBegin.val[0] =
      gccs::alignPrev(weightedAreaBegin.val[0], blockRows);
  weightedAreaBegin.val[1] =
      gccs::alignPrev(weightedAreaBegin.val[1], blockCols);
  weightedAreaEnd.val[0] = gccs::alignPrev(weightedAreaEnd.val[0], blockRows);
  weightedAreaEnd.val[1] = gccs::alignPrev(weightedAreaEnd.val[1], blockCols);

  PlanningCache cache;

  poplar::OptionFlags options;
  bool doBwdPass = pass == Pass::BWD || pass == Pass::ALL;
  bool doWuPass = pass == Pass::WU || pass == Pass::ALL;
  options.set("availableMemoryProportion", "1.0");
  options.set("doGradAPass", doBwdPass ? "true" : "false");
  options.set("doGradWPass", doWuPass ? "true" : "false");
  options.set("partialsType", partialsType.toString());

  // User options specified via --matmul-options override defaults
  if (!matmulOptionsString.empty()) {
    poplar::readJSON(matmulOptionsString, options);
  }
  if (!planConstraintsString.empty()) {
    options.set("planConstraints", planConstraintsString);
  }

  auto device = tilesPerIPU
                    ? createTestDevice(deviceType, numIPUs, *tilesPerIPU, true)
                    : createTestDeviceFullSize(deviceType, numIPUs, true);
  const auto &target = device.getTarget();

  const auto sparsityType =
      blockArea == 1 ? SparsityType::Element : SparsityType::Block;

  SparsityParams sparsityParams(sparsityType, SparsityStructure::Unstructured,
                                {blockRows, blockCols});

  const auto params = FullyConnectedParams::createWithNzRatio(
      std::move(sparsityParams), sparsityFactor, batchSize, numGroups,
      inputSize, outputSize);

  popsparse::fullyconnected::Plan plan;
  popsparse::fullyconnected::Cost planCost;

  // Always do forward
  std::tie(plan, planCost) = popsparse::fullyconnected::getPlan(
      target, dataType, params, options, &cache);

  if (reportPlan) {
    std::string str;
    if (doBwdPass || doWuPass) {
      str += "Joint : ";
      if (doBwdPass) {
        str += " GradA + ";
      }
      if (doWuPass) {
        str += " GradW + ";
      }
    }
    str += "Fwd Plan \n";
    std::cerr << str << plan << "\n" << planCost << "\n";
  }

  std::size_t fwdMetaInfoBucketSize = plan.fwdMetaInfoElemsPerBucket;
  std::size_t gradAMetaInfoBucketSize = plan.gradAMetaInfoElemsPerBucket;
  std::size_t nzElementBucketSize = plan.nzElemsPerBucket;

  std::cerr << "Using bucket sizes:"
            << "\n  meta info (forward): " << fwdMetaInfoBucketSize
            << "\n  meta info (grad-a): " << gradAMetaInfoBucketSize
            << "\n  nz element : " << nzElementBucketSize << "\n";

  Partitioner<EType> partitioner(params, dataType, target, options, &cache);
  std::mt19937 randomEngine;
  if (variableSeed) {
    using namespace std::chrono;
    using SeedDurationType = duration<std::mt19937::result_type, std::nano>;
    const auto now = high_resolution_clock::now();
    const auto seed =
        duration_cast<SeedDurationType>(now.time_since_epoch()).count();
    randomEngine.seed(seed);
  }

  const bool floatingPointCouldRepresentMaxAccum = [&] {
    const auto maxVal = maxContiguousInteger(dataType);

    double weightedThreshold, remainingThreshold;
    std::tie(weightedThreshold, remainingThreshold) =
        poplibs_test::sparse::calculateWeightedVsRemainingSparsityFactor(
            {outputSize / blockRows, inputSize / blockCols}, sparsityFactor,
            {weightedAreaBegin[0] / blockRows,
             weightedAreaBegin[1] / blockCols},
            {weightedAreaEnd[0] / blockRows, weightedAreaEnd[1] / blockCols},
            weightedAreaWeighting);
    const auto numWeightedInputChannels =
        (weightedAreaEnd[1] - weightedAreaBegin[1]);
    const auto numWeightedOutputChannels =
        (weightedAreaEnd[0] - weightedAreaBegin[0]);
    std::size_t maxInputChannels =
        numWeightedInputChannels * weightedThreshold +
        (params.getInputChannelsPerGroup() - numWeightedInputChannels) *
            remainingThreshold;
    std::size_t maxOutputChannels =
        numWeightedOutputChannels * weightedThreshold +
        (params.getOutputChannelsPerGroup() - numWeightedOutputChannels) *
            remainingThreshold;
    maxInputChannels = gccs::alignPrev(maxInputChannels, blockCols);
    maxOutputChannels = gccs::alignPrev(maxOutputChannels, blockRows);
    const auto getOpsPerOutputElementEstimate = [&](const Pass &pass) -> int {
      const auto numAccumulations = pass == Pass::FWD   ? maxInputChannels
                                    : pass == Pass::BWD ? maxOutputChannels
                                                        : params.getBatchSize();
      return numAccumulations;
    };
    // We use a modifier to account for the unlikeliness of picking all positive
    // or negative 1s which would actually get us to the max precisely
    // represented integer.
    constexpr int modifier = 10;
    // We use another modifier to account for the chance that sparsity is not
    // perfectly evenly spread in this instant.
    constexpr double wiggleRoom = 1.3;
    if (wiggleRoom * getOpsPerOutputElementEstimate(Pass::FWD) >
        maxVal * modifier) {
      return false;
    }
    if (doBwdPass && wiggleRoom * getOpsPerOutputElementEstimate(Pass::BWD) >
                         maxVal * modifier) {
      return false;
    }
    if (doWuPass && wiggleRoom * getOpsPerOutputElementEstimate(Pass::WU) >
                        maxVal * modifier) {
      return false;
    }
    return true;
  }();

  // create CSR matrix for the given sparsity factor
  const bool useBipolarDistribution = floatingPointCouldRepresentMaxAccum;
  std::array<std::size_t, 2> blockDims = {blockRows, blockCols};
  CSRMatrix<EType> csrMatrix(blockDims);
  std::tie(csrMatrix.nzValues, csrMatrix.columnIndices, csrMatrix.rowIndices) =
      poplibs_test::sparse::buildCSRMatrix<EType, std::size_t>(
          randomEngine, {outputSize, inputSize}, {blockRows, blockCols},
          sparsityFactor, weightedAreaBegin, weightedAreaEnd,
          weightedAreaWeighting, useBipolarDistribution);

  // Forward
  boost::multi_array<double, 2> hostInput(boost::extents[batchSize][inputSize]);
  if (useBipolarDistribution) {
    writeRandomBinaryValues(target, dataType, hostInput, -1.0, 1.0,
                            randomEngine);
  } else {
    writeRandomValues(target, dataType, hostInput, -3.0, 3.0, randomEngine);
  }
  boost::multi_array<double, 2> hostOutputActs(
      boost::extents[batchSize][outputSize]);

  // GradA
  boost::multi_array<double, 2> hostOutputGrad(
      boost::extents[batchSize][outputSize]);
  if (useBipolarDistribution) {
    writeRandomBinaryValues(target, dataType, hostOutputGrad, -1.0, 1.0,
                            randomEngine);
  } else {
    writeRandomValues(target, dataType, hostOutputGrad, -3.0, 3.0,
                      randomEngine);
  }
  boost::multi_array<double, 2> hostInputGrad(
      boost::extents[batchSize][inputSize]);

  if (!plan.useDense) {
    std::cerr << "Logging Forward pass bucket statistics: ";
    auto pnBucketsImpl = partitioner.getImpl().createBuckets(csrMatrix);

    logBucketStatistics(pnBucketsImpl.pnBuckets, csrMatrix);
  }

  if (planOnly) {
    return 0;
  }

  Graph graph(target);
  popops::addCodelets(graph);
  poplin::addCodelets(graph);
  popsparse::addCodelets(graph);
  Sequence fwdProg, bwdProg, wuProg, uploadProg, downloadProg;

  // Build the graph
  std::cerr << "Constructing graph...\n";
  const SparseTensor weights = createFullyConnectedWeights(
      graph, dataType, params, "weights", options, &cache);
  const Tensor input = createFullyConnectedInput(graph, dataType, params,
                                                 "input", options, &cache);
  const Tensor outputActs = fullyConnectedFwd(graph, weights, input, params,
                                              fwdProg, "fwd", options, &cache);

  // GradW
  boost::multi_array<double, 1> hostWeightGrad(
      boost::extents[weights.getNzValuesTensor().numElements()]);

  Tensor outputGrad;
  if (doBwdPass || doWuPass) {
    outputGrad = graph.clone(outputActs, "outputGrad");
  }
  Tensor inputGrad;
  if (doBwdPass) {
    inputGrad = fullyConnectedGradA(graph, weights, outputGrad, params, bwdProg,
                                    "grada", options, &cache);
  }

  Tensor weightGrad;
  if (doWuPass) {
    weightGrad = fullyConnectedSparseGradW(graph, weights.getMetaInfoTensor(),
                                           outputGrad, input, params, wuProg,
                                           "wu", options, &cache);
  }
  std::cerr << "Done\n";

  std::vector<std::pair<std::string, HostMemory>> tmap;
  auto rawMetaInfo =
      allocateHostMemoryForTensor(weights.getMetaInfoTensor(), "weights.meta",
                                  graph, uploadProg, downloadProg, tmap);

  auto rawNzInfo =
      allocateHostMemoryForTensor(weights.getNzValuesTensor(), "weights.nz",
                                  graph, uploadProg, downloadProg, tmap);
  auto rawInput = allocateHostMemoryForTensor(input, "input", graph, uploadProg,
                                              downloadProg, tmap);
  auto rawOutputActs = allocateHostMemoryForTensor(
      outputActs, "outputActs", graph, uploadProg, downloadProg, tmap);

  std::unique_ptr<char[]> rawOutputGrad;
  if (doBwdPass || doWuPass) {
    rawOutputGrad = allocateHostMemoryForTensor(outputGrad, "outputGrad", graph,
                                                uploadProg, downloadProg, tmap);
  }

  std::unique_ptr<char[]> rawInputGrad;
  if (!ignoreData && doBwdPass) {
    rawInputGrad = allocateHostMemoryForTensor(inputGrad, "inputGrad", graph,
                                               uploadProg, downloadProg, tmap);
  }

  std::unique_ptr<char[]> rawWeightGrad;
  if (!ignoreData && doWuPass) {
    rawWeightGrad = allocateHostMemoryForTensor(weightGrad, "weightGrad", graph,
                                                uploadProg, downloadProg, tmap);
  }

  Tensor fwdCycles, bwdCycles, wuCycles;
  if (reportTotalCycleCounts) {
    fwdCycles = cycleCount(graph, fwdProg, 0, SyncType::INTERNAL, "fwdCycles");
    graph.createHostRead("fwdCycles", fwdCycles);
    if (doBwdPass) {
      bwdCycles =
          cycleCount(graph, bwdProg, 0, SyncType::INTERNAL, "bwdCycles");
      graph.createHostRead("bwdCycles", bwdCycles);
    }
    if (doWuPass) {
      wuCycles = cycleCount(graph, wuProg, 0, SyncType::INTERNAL, "wuCycles");
      graph.createHostRead("wuCycles", wuCycles);
    }
  }
  Sequence controlProg({std::move(uploadProg), std::move(fwdProg),
                        std::move(bwdProg), std::move(wuProg)});
  if (!ignoreData) {
    controlProg.add(std::move(downloadProg));
  }

  std::cerr << "Creating engine...\n";
  std::optional<TempDir> tempDir;
  OptionFlags engineOptions;
  if (vm.count("profile") || profileDir) {
    engineOptions.set("autoReport.outputExecutionProfile", "true");
    if (profileDir) {
      engineOptions.set("autoReport.directory", *profileDir);
    } else {
      tempDir.emplace(TempDir::create());
      engineOptions.set("autoReport.directory", tempDir->getPath());
    }
  }
  Engine engine(graph, std::move(controlProg), engineOptions);

  if (vm.count("compile-only"))
    return 0;

  attachStreams(engine, tmap);

  std::cerr << "Done\n";

  std::cerr << "Running...\n";

  // Actual bucket info use by device graph
  auto buckets = partitioner.createSparsityDataImpl(csrMatrix);

  const auto &metaInfoFlat = buckets.metaInfo;
  const auto &nzValuesFlat = buckets.nzValues;
  // Overflow info is the same for all passes at time of writing.
  if (!plan.useDense) {
    std::cerr << "overflowInfo = {" << metaInfoFlat.at(0) << ","
              << metaInfoFlat.at(1) << "," << metaInfoFlat.at(2) << "}\n";
  }

  copy(target, hostInput, dataType, rawInput.get());
  if (rawMetaInfo) {
    copy(target, metaInfoFlat, UNSIGNED_SHORT, rawMetaInfo.get());
  }
  copy(target, nzValuesFlat, dataType, rawNzInfo.get());

  if (!ignoreData && (doBwdPass || doWuPass)) {
    copy(target, hostOutputGrad, outputGrad.elementType(), rawOutputGrad.get());
  }

  device.bind([&](const Device &d) {
    engine.loadAndRun(d);
    if (reportTotalCycleCounts) {
      std::uint64_t cyclesBuffer;
      engine.readTensor("fwdCycles", &cyclesBuffer, &cyclesBuffer + 1);
      std::cerr << "  Forward pass cycles: " << cyclesBuffer << "\n";
      if (doBwdPass) {
        engine.readTensor("bwdCycles", &cyclesBuffer, &cyclesBuffer + 1);
        std::cerr << "  GradA pass cycles: " << cyclesBuffer << "\n";
      }
      if (doWuPass) {
        engine.readTensor("wuCycles", &cyclesBuffer, &cyclesBuffer + 1);
        std::cerr << "  GradW pass cycles: " << cyclesBuffer << "\n";
      }
    }
  });

  bool matchesModel = true;
  if (!ignoreData) {
    const double relTolerance = dataType == HALF ? HALF_REL_TOL : FLOAT_REL_TOL;
    copy(target, outputActs.elementType(), rawOutputActs.get(), hostOutputActs);
    boost::multi_array<double, 2> hostDenseWeights(
        boost::extents[outputSize][inputSize]);
    boost::multi_array<double, 2> modelOutputActs(
        boost::extents[batchSize][outputSize]);
    hostDenseWeights = poplibs_test::sparse::csrToDenseMatrix(
        csrMatrix.nzValues.data(), csrMatrix.columnIndices.data(),
        csrMatrix.rowIndices.data(), csrMatrix.nzValues.size(), outputSize,
        inputSize, blockRows, blockCols);

    poplibs_test::gemm::generalMatrixMultiply(hostInput, hostDenseWeights,
                                              modelOutputActs, false, true);
    matchesModel &= checkIsClose("outputActs", hostOutputActs, modelOutputActs,
                                 relTolerance);
    if (doBwdPass) {
      copy(target, inputGrad.elementType(), rawInputGrad.get(), hostInputGrad);
      boost::multi_array<double, 2> modelInputGrad(
          boost::extents[batchSize][inputSize]);
      poplibs_test::gemm::generalMatrixMultiply(
          hostOutputGrad, hostDenseWeights, modelInputGrad, false, false);
      matchesModel &= checkIsClose("inputGrad", hostInputGrad, modelInputGrad,
                                   relTolerance);
    }
    if (doWuPass) {
      copy(target, weightGrad.elementType(), rawWeightGrad.get(),
           hostWeightGrad);
      boost::multi_array<double, 2> modelWeightGrad(
          boost::extents[outputSize][inputSize]);
      poplibs_test::gemm::generalMatrixMultiply(hostOutputGrad, hostInput,
                                                modelWeightGrad, true, false);

      std::vector<EType> modelNzValuesCSR;
      auto columnIdxIt = csrMatrix.columnIndices.begin();
      for (auto rowIt = std::next(csrMatrix.rowIndices.begin());
           rowIt != csrMatrix.rowIndices.end(); ++rowIt) {
        const auto nnzThisRow = *rowIt - *std::prev(rowIt);
        const auto rowIdx =
            (std::distance(csrMatrix.rowIndices.begin(), rowIt) - 1) *
            blockRows;
        for (std::size_t i = 0; i < nnzThisRow / blockArea; ++i) {
          const auto columnIdx = *columnIdxIt++;
          for (std::size_t r = 0; r != blockRows; ++r) {
            for (std::size_t c = 0; c != blockCols; ++c) {
              // row major order within a block
              modelNzValuesCSR.emplace_back(
                  modelWeightGrad[rowIdx + r][columnIdx + c]);
            }
          }
        }
      }

      assert(modelNzValuesCSR.size() ==
             csrMatrix.columnIndices.size() * blockArea);

      SparsityDataImpl<EType> actualBuckets;
      std::vector<EType> actualWeightGrads;
      actualWeightGrads.reserve(hostWeightGrad.size());
      std::copy(hostWeightGrad.begin(), hostWeightGrad.end(),
                std::back_inserter(actualWeightGrads));

      actualBuckets.nzValues = std::move(actualWeightGrads);
      actualBuckets.metaInfo = std::move(buckets.metaInfo);
      auto actualCSR = partitioner.sparsityDataImplToCSRMatrix(actualBuckets);

      auto ait = actualCSR.nzValues.begin();
      auto mit = modelNzValuesCSR.begin();
      for (; ait != actualCSR.nzValues.end(); ++ait, ++mit) {
        bool elemMatch = checkIsClose(*mit, *ait, relTolerance);
        if (!elemMatch) {
          std::cerr << "mismatch at  WeightsGrad.nz[";
          std::cerr << std::distance(actualCSR.nzValues.begin(), ait);
          std::cerr << "]:" << *mit << "!=" << *ait << "\n";
        }
        matchesModel &= elemMatch;
      }
      auto columnsMatch = std::equal(
          actualCSR.columnIndices.begin(), actualCSR.columnIndices.end(),
          csrMatrix.columnIndices.begin(), csrMatrix.columnIndices.end());
      if (!columnsMatch) {
        std::cerr << "CSR columns indices do not match\n";
      }
      matchesModel &= columnsMatch;
    }
  }

  if (denseGradWSerialSplits && doWuPass) {
    auto serialSplits =
        fullyConnectedDenseGradWSerialSplits(graph, dataType, params, options);
    std::cerr << "Dense GradW serial splits : "
              << "   groups " << std::get<0>(serialSplits)
              << "   input channel " << std::get<1>(serialSplits)
              << "   output channel " << std::get<2>(serialSplits) << "\n";
  }

  std::cerr << "Done\n";

  if (profile) {
    engine.printProfileSummary(std::cout, {{"showExecutionSteps", "true"}});
  }

  if (!matchesModel) {
    std::cerr << "Validation failed\n";
    return 1;
  }

  return 0;
} catch (const poplar::graph_memory_allocation_error &e) {
  if (!e.profilePath.empty()) {
    poplar::printGraphSummary(std::cerr, e.profilePath,
                              {{"showVarStorage", "true"}});
  }
  throw;
}