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

#define BOOST_TEST_MODULE StridedRegionsTest
#include "poplibs_support/StridedRegions.hpp"
#include "poplibs_support/TestDevice.hpp"
#include <boost/test/unit_test.hpp>

#include <poplar/Graph.hpp>

using namespace poplar;
using namespace poplibs_support;

// A regular strided input sequence.
static const std::vector<Interval> input1 = {
    {0, 10},  {20, 30}, {40, 50}, {60, 70}, {80, 90},
    {10, 20}, {30, 40}, {50, 60}, {70, 80}, {90, 100}};

// Input with broadcasting.
static const std::vector<Interval> input2 = {{0, 1}, {0, 1}, {0, 1}};

// Input without an inner-most dimension.
static const std::vector<Interval> input3 = {
    {3, 4}, {7, 8}, {11, 12}, {15, 16}, {5, 6}, {9, 10}, {13, 14}, {17, 18}};

// Input where the resulting offset + stride causes us to hit
// an edge case in the code that slices a strided region from a tensor because
// the underlying variable has just enough elements for the given intervals
// but the intermediate shape we apply to the tensor would reference more
// elements than this
static const std::vector<Interval> input4 = {
    {2, 3}, {6, 7}, {10, 11}, {14, 15}, {4, 5}, {8, 9}, {12, 13}, {16, 17}};

// A set of intervals from a poplibs norms test that caused multiple problems
static const std::vector<Interval> input5 = {
    {392, 394}, {6664, 6666}, {12936, 12938}, {19208, 19210},
    {394, 396}, {6666, 6668}, {12938, 12940}, {19210, 19212},
    {396, 398}, {6668, 6670}, {12940, 12942}, {19212, 19214},
    {398, 400}, {6670, 6672}, {12942, 12944}, {19214, 19216},
    {400, 402}, {6672, 6674}, {12944, 12946}, {19216, 19218},
    {402, 404}, {6674, 6676}, {12946, 12948}, {19218, 19220},
    {404, 406}, {6676, 6678}, {12948, 12950}, {19220, 19222},
    {406, 408}, {6678, 6680}, {12950, 12952}, {19222, 19224},
    {408, 410}, {6680, 6682}, {12952, 12954}, {19224, 19226},
    {410, 412}, {6682, 6684}, {12954, 12956}, {19226, 19228},
    {412, 414}, {6684, 6686}, {12956, 12958}, {19228, 19230},
    {414, 416}, {6686, 6688}, {12958, 12960}, {19230, 19232},
    {416, 418}, {6688, 6690}, {12960, 12962}, {19232, 19234},
    {418, 420}, {6690, 6692}, {12962, 12964}, {19234, 19236},
    {420, 422}, {6692, 6694}, {12964, 12966}, {19236, 19238},
    {422, 424}, {6694, 6696}, {12966, 12968}, {19238, 19240},
    {424, 426}, {6696, 6698}, {12968, 12970}, {19240, 19242},
    {426, 428}, {6698, 6700}, {12970, 12972}, {19242, 19244},
    {428, 430}, {6700, 6702}, {12972, 12974}, {19244, 19246},
    {430, 432}, {6702, 6704}, {12974, 12976}, {19246, 19248},
    {432, 434}, {6704, 6706}, {12976, 12978}, {19248, 19250},
    {434, 436}, {6706, 6708}, {12978, 12980}, {19250, 19252},
    {436, 438}, {6708, 6710}, {12980, 12982}, {19252, 19254},
    {438, 440}, {6710, 6712}, {12982, 12984}, {19254, 19256},
    {440, 442}, {6712, 6714}, {12984, 12986}, {19256, 19258},
    {442, 444}, {6714, 6716}, {12986, 12988}, {19258, 19260},
    {444, 446}, {6716, 6718}, {12988, 12990}, {19260, 19262},
    {446, 448}, {6718, 6720}, {12990, 12992}, {19262, 19264},
    {448, 450}, {6720, 6722}, {12992, 12994}, {19264, 19266},
    {450, 452}, {6722, 6724}, {12994, 12996}, {19266, 19268},
    {452, 454}, {6724, 6726}, {12996, 12998}, {19268, 19270},
    {454, 456}, {6726, 6728}, {12998, 13000}, {19270, 19272},
    {456, 458}, {6728, 6730}, {13000, 13002}, {19272, 19274},
    {458, 460}, {6730, 6732}, {13002, 13004}, {19274, 19276},
    {460, 462}, {6732, 6734}, {13004, 13006}, {19276, 19278},
    {462, 464}, {6734, 6736}, {13006, 13008}, {19278, 19280},
    {464, 466}, {6736, 6738}, {13008, 13010}, {19280, 19282},
    {466, 468}, {6738, 6740}, {13010, 13012}, {19282, 19284},
    {468, 470}, {6740, 6742}, {13012, 13014}, {19284, 19286},
    {470, 472}, {6742, 6744}, {13014, 13016}, {19286, 19288},
    {472, 474}, {6744, 6746}, {13016, 13018}, {19288, 19290},
    {474, 476}, {6746, 6748}, {13018, 13020}, {19290, 19292},
    {476, 478}, {6748, 6750}, {13020, 13022}, {19292, 19294},
    {478, 480}, {6750, 6752}, {13022, 13024}, {19294, 19296},
    {480, 482}, {6752, 6754}, {13024, 13026}, {19296, 19298},
    {482, 484}, {6754, 6756}, {13026, 13028}, {19298, 19300},
    {484, 486}, {6756, 6758}, {13028, 13030}, {19300, 19302},
    {486, 488}, {6758, 6760}, {13030, 13032}, {19302, 19304},
    {488, 490}, {6760, 6762}, {13032, 13034}, {19304, 19306},
    {490, 492}, {6762, 6764}, {13034, 13036}, {19306, 19308},
    {492, 494}, {6764, 6766}, {13036, 13038}, {19308, 19310},
    {494, 496}, {6766, 6768}, {13038, 13040}, {19310, 19312},
    {496, 498}, {6768, 6770}, {13040, 13042}, {19312, 19314},
    {498, 500}, {6770, 6772}, {13042, 13044}, {19314, 19316},
    {500, 502}, {6772, 6774}, {13044, 13046}, {19316, 19318},
    {502, 504}, {6774, 6776}, {13046, 13048}, {19318, 19320},
    {504, 506}, {6776, 6778}, {13048, 13050}, {19320, 19322},
    {506, 508}, {6778, 6780}, {13050, 13052}, {19322, 19324},
    {508, 510}, {6780, 6782}, {13052, 13054}, {19324, 19326},
    {510, 512}, {6782, 6784}, {13054, 13056}, {19326, 19328},
    {512, 514}, {6784, 6786}, {13056, 13058}, {19328, 19330},
    {514, 516}, {6786, 6788}, {13058, 13060}, {19330, 19332},
    {516, 518}, {6788, 6790}, {13060, 13062}, {19332, 19334},
    {518, 520}, {6790, 6792}, {13062, 13064}, {19334, 19336},
    {520, 522}, {6792, 6794}, {13064, 13066}, {19336, 19338},
    {522, 524}, {6794, 6796}, {13066, 13068}, {19338, 19340},
    {524, 526}, {6796, 6798}, {13068, 13070}, {19340, 19342},
    {526, 528}, {6798, 6800}, {13070, 13072}, {19342, 19344},
    {528, 530}, {6800, 6802}, {13072, 13074}, {19344, 19346},
    {530, 532}, {6802, 6804}, {13074, 13076}, {19346, 19348},
    {532, 534}, {6804, 6806}, {13076, 13078}, {19348, 19350},
    {534, 536}, {6806, 6808}, {13078, 13080}, {19350, 19352},
    {536, 538}, {6808, 6810}, {13080, 13082}, {19352, 19354},
    {538, 540}, {6810, 6812}, {13082, 13084}, {19354, 19356},
    {540, 542}, {6812, 6814}, {13084, 13086}, {19356, 19358},
    {542, 544}, {6814, 6816}, {13086, 13088}, {19358, 19360},
    {544, 546}, {6816, 6818}, {13088, 13090}, {19360, 19362},
    {546, 548}, {6818, 6820}, {13090, 13092}, {19362, 19364},
    {548, 550}, {6820, 6822}, {13092, 13094}, {19364, 19366},
    {550, 552}, {6822, 6824}, {13094, 13096}, {19366, 19368},
    {552, 554}, {6824, 6826}, {13096, 13098}, {19368, 19370},
    {554, 556}, {6826, 6828}, {13098, 13100}, {19370, 19372},
    {556, 558}, {6828, 6830}, {13100, 13102}, {19372, 19374},
    {558, 560}, {6830, 6832}, {13102, 13104}, {19374, 19376},
    {560, 562}, {6832, 6834}, {13104, 13106}, {19376, 19378},
    {562, 564}, {6834, 6836}, {13106, 13108}, {19378, 19380},
    {564, 566}, {6836, 6838}, {13108, 13110}, {19380, 19382},
    {566, 568}, {6838, 6840}, {13110, 13112}, {19382, 19384},
    {568, 570}, {6840, 6842}, {13112, 13114}, {19384, 19386},
    {570, 572}, {6842, 6844}, {13114, 13116}, {19386, 19388},
    {572, 574}, {6844, 6846}, {13116, 13118}, {19388, 19390},
    {574, 576}, {6846, 6848}, {13118, 13120}, {19390, 19392},
    {576, 578}, {6848, 6850}, {13120, 13122}, {19392, 19394},
    {578, 580}, {6850, 6852}, {13122, 13124}, {19394, 19396},
    {580, 582}, {6852, 6854}, {13124, 13126}, {19396, 19398},
    {582, 584}, {6854, 6856}, {13126, 13128}, {19398, 19400},
    {584, 586}, {6856, 6858}, {13128, 13130}, {19400, 19402},
    {586, 588}, {6858, 6860}, {13130, 13132}, {19402, 19404},
    {588, 590}, {6860, 6862}, {13132, 13134}, {19404, 19406},
    {590, 592}, {6862, 6864}, {13134, 13136}, {19406, 19408},
    {592, 594}, {6864, 6866}, {13136, 13138}, {19408, 19410},
    {594, 596}, {6866, 6868}, {13138, 13140}, {19410, 19412},
    {596, 598}, {6868, 6870}, {13140, 13142}, {19412, 19414},
    {598, 600}, {6870, 6872}, {13142, 13144}, {19414, 19416},
    {600, 602}, {6872, 6874}, {13144, 13146}, {19416, 19418},
    {602, 604}, {6874, 6876}, {13146, 13148}, {19418, 19420},
    {604, 606}, {6876, 6878}, {13148, 13150}, {19420, 19422},
    {606, 608}, {6878, 6880}, {13150, 13152}, {19422, 19424},
    {608, 610}, {6880, 6882}, {13152, 13154}, {19424, 19426},
    {610, 612}, {6882, 6884}, {13154, 13156}, {19426, 19428},
    {612, 614}, {6884, 6886}, {13156, 13158}, {19428, 19430},
    {614, 616}, {6886, 6888}, {13158, 13160}, {19430, 19432},
    {616, 618}, {6888, 6890}, {13160, 13162}, {19432, 19434},
    {618, 620}, {6890, 6892}, {13162, 13164}, {19434, 19436},
    {620, 622}, {6892, 6894}, {13164, 13166}, {19436, 19438},
    {622, 624}, {6894, 6896}, {13166, 13168}, {19438, 19440},
    {624, 626}, {6896, 6898}, {13168, 13170}, {19440, 19442},
    {626, 628}, {6898, 6900}, {13170, 13172}, {19442, 19444},
    {628, 630}, {6900, 6902}, {13172, 13174}, {19444, 19446},
    {630, 632}, {6902, 6904}, {13174, 13176}, {19446, 19448},
    {632, 634}, {6904, 6906}, {13176, 13178}, {19448, 19450},
    {634, 636}, {6906, 6908}, {13178, 13180}, {19450, 19452},
    {636, 638}, {6908, 6910}, {13180, 13182}, {19452, 19454},
    {638, 640}, {6910, 6912}, {13182, 13184}, {19454, 19456},
    {640, 642}, {6912, 6914}, {13184, 13186}, {19456, 19458},
    {642, 644}, {6914, 6916}, {13186, 13188}, {19458, 19460},
    {644, 646}, {6916, 6918}, {13188, 13190}, {19460, 19462},
    {646, 648}, {6918, 6920}, {13190, 13192}, {19462, 19464},
    {648, 650}, {6920, 6922}, {13192, 13194}, {19464, 19466},
    {650, 652}, {6922, 6924}, {13194, 13196}, {19466, 19468},
    {652, 654}, {6924, 6926}, {13196, 13198}, {19468, 19470},
    {654, 656}, {6926, 6928}, {13198, 13200}, {19470, 19472},
    {656, 658}, {6928, 6930}, {13200, 13202}, {19472, 19474},
    {658, 660}, {6930, 6932}, {13202, 13204}, {19474, 19476},
    {660, 662}, {6932, 6934}, {13204, 13206}, {19476, 19478},
    {662, 664}, {6934, 6936}, {13206, 13208}, {19478, 19480},
    {664, 666}, {6936, 6938}, {13208, 13210}, {19480, 19482},
    {666, 668}, {6938, 6940}, {13210, 13212}, {19482, 19484},
    {668, 670}, {6940, 6942}, {13212, 13214}, {19484, 19486},
    {670, 672}, {6942, 6944}, {13214, 13216}, {19486, 19488},
    {672, 674}, {6944, 6946}, {13216, 13218}, {19488, 19490},
    {674, 676}, {6946, 6948}, {13218, 13220}, {19490, 19492},
    {676, 678}, {6948, 6950}, {13220, 13222}, {19492, 19494},
    {678, 680}, {6950, 6952}, {13222, 13224}, {19494, 19496},
    {680, 682}, {6952, 6954}, {13224, 13226}, {19496, 19498},
    {682, 684}, {6954, 6956}, {13226, 13228}, {19498, 19500},
    {684, 686}, {6956, 6958}, {13228, 13230}, {19500, 19502},
    {686, 688}, {6958, 6960}, {13230, 13232}, {19502, 19504},
    {688, 690}, {6960, 6962}, {13232, 13234}, {19504, 19506},
    {690, 692}, {6962, 6964}, {13234, 13236}, {19506, 19508},
    {692, 694}, {6964, 6966}, {13236, 13238}, {19508, 19510},
    {694, 696}, {6966, 6968}, {13238, 13240}, {19510, 19512},
    {696, 698}, {6968, 6970}, {13240, 13242}, {19512, 19514},
    {698, 700}, {6970, 6972}, {13242, 13244}, {19514, 19516},
    {700, 702}, {6972, 6974}, {13244, 13246}, {19516, 19518},
    {702, 704}, {6974, 6976}, {13246, 13248}, {19518, 19520},
    {704, 706}, {6976, 6978}, {13248, 13250}, {19520, 19522},
    {706, 708}, {6978, 6980}, {13250, 13252}, {19522, 19524},
    {708, 710}, {6980, 6982}, {13252, 13254}, {19524, 19526},
    {710, 712}, {6982, 6984}, {13254, 13256}, {19526, 19528},
    {712, 714}, {6984, 6986}, {13256, 13258}, {19528, 19530},
    {714, 716}, {6986, 6988}, {13258, 13260}, {19530, 19532},
    {716, 718}, {6988, 6990}, {13260, 13262}, {19532, 19534},
    {718, 720}, {6990, 6992}, {13262, 13264}, {19534, 19536},
    {720, 722}, {6992, 6994}, {13264, 13266}, {19536, 19538},
    {722, 724}, {6994, 6996}, {13266, 13268}, {19538, 19540},
    {724, 726}, {6996, 6998}, {13268, 13270}, {19540, 19542},
    {726, 728}, {6998, 7000}, {13270, 13272}, {19542, 19544},
    {728, 730}, {7000, 7002}, {13272, 13274}, {19544, 19546},
    {730, 732}, {7002, 7004}, {13274, 13276}, {19546, 19548},
    {732, 734}, {7004, 7006}, {13276, 13278}, {19548, 19550},
    {734, 736}, {7006, 7008}, {13278, 13280}, {19550, 19552},
    {736, 738}, {7008, 7010}, {13280, 13282}, {19552, 19554},
    {738, 740}, {7010, 7012}, {13282, 13284}, {19554, 19556},
    {740, 742}, {7012, 7014}, {13284, 13286}, {19556, 19558},
    {742, 744}, {7014, 7016}, {13286, 13288}, {19558, 19560},
    {744, 746}, {7016, 7018}, {13288, 13290}, {19560, 19562},
    {746, 748}, {7018, 7020}, {13290, 13292}, {19562, 19564},
    {748, 750}, {7020, 7022}, {13292, 13294}, {19564, 19566},
    {750, 752}, {7022, 7024}, {13294, 13296}, {19566, 19568},
    {752, 754}, {7024, 7026}, {13296, 13298}, {19568, 19570},
    {754, 756}, {7026, 7028}, {13298, 13300}, {19570, 19572},
    {756, 758}, {7028, 7030}, {13300, 13302}, {19572, 19574},
    {758, 760}, {7030, 7032}, {13302, 13304}, {19574, 19576},
    {760, 762}, {7032, 7034}, {13304, 13306}, {19576, 19578},
    {762, 764}, {7034, 7036}, {13306, 13308}, {19578, 19580},
    {764, 766}, {7036, 7038}, {13308, 13310}, {19580, 19582},
    {766, 768}, {7038, 7040}, {13310, 13312}, {19582, 19584},
    {768, 770}, {7040, 7042}, {13312, 13314}, {19584, 19586},
    {770, 772}, {7042, 7044}, {13314, 13316}, {19586, 19588},
    {772, 774}, {7044, 7046}, {13316, 13318}, {19588, 19590},
    {774, 776}, {7046, 7048}, {13318, 13320}, {19590, 19592},
    {776, 778}, {7048, 7050}, {13320, 13322}, {19592, 19594},
    {778, 780}, {7050, 7052}, {13322, 13324}, {19594, 19596},
    {780, 782}, {7052, 7054}, {13324, 13326}, {19596, 19598},
    {782, 784}, {7054, 7056}, {13326, 13328}, {19598, 19600}};

// Check that the result of creating and merging strided regions
// for the given set of intervals matches an expected result.
static void checkMerge(const std::vector<Interval> &intervals,
                       const StridedRegionList &expected) {
  StridedRegionList merged;
  merged.reserve(intervals.size());
  for (const auto &i : intervals) {
    merged.emplace_back(i);
  }
  mergeStridedRegions(merged);

  BOOST_CHECK_EQUAL_COLLECTIONS(expected.begin(), expected.end(),
                                merged.begin(), merged.end());
}

// Check that the result of creating and merging strided regions
// for the given set of intervals and then retreiving the intervals
// again matches the input.
static void checkRoundTrip(const std::vector<Interval> &intervals) {
  StridedRegionList strided;
  strided.reserve(intervals.size());
  for (const auto &i : intervals) {
    strided.emplace_back(i);
  }

  {
    const auto is = intervalsForStridedRegions(strided);
    BOOST_CHECK_EQUAL_COLLECTIONS(intervals.begin(), intervals.end(),
                                  is.begin(), is.end());
  }

  mergeStridedRegions(strided);
  {
    const auto is = intervalsForStridedRegions(strided);
    BOOST_CHECK_EQUAL_COLLECTIONS(intervals.begin(), intervals.end(),
                                  is.begin(), is.end());
  }
}

// Check that creating and merging strided regions and then slicing
// a tensor produces a tensor referencing the same underlying elements
// as just concatenating the slices given by the input intervals.
static void checkSlice(const std::vector<Interval> &intervals) {
  StridedRegionList strided;
  strided.reserve(intervals.size());
  for (const auto &i : intervals) {
    strided.emplace_back(i);
  }

  // Find an upper bound for the intervals given to use as a variable
  // size.
  std::size_t maxElem = 0;
  for (const auto &i : intervals) {
    maxElem = std::max(maxElem, i.end());
  }

  const auto device = createTestDevice(TEST_TARGET);
  Graph graph(device.getTarget());

  auto t = graph.addVariable(INT, {maxElem});
  const auto expected = concat(t.slices(intervals));
  {
    auto actual = sliceStridedRegions(t, strided);
    BOOST_CHECK_EQUAL(expected, actual);
  }

  mergeStridedRegions(strided);
  {
    auto actual = sliceStridedRegions(t, strided);
    BOOST_CHECK_EQUAL(expected, actual);
  }
}

BOOST_AUTO_TEST_CASE(Merge1) {
  StridedRegionList expected = {
      StridedRegion(0, StrideDesc{{1, 10}, {20, 5}, {10, 2}})};
  checkMerge(input1, expected);
}

BOOST_AUTO_TEST_CASE(Merge2) {
  StridedRegionList expected = {StridedRegion(0, StrideDesc{{0, 3}})};
  checkMerge(input2, expected);
}

BOOST_AUTO_TEST_CASE(Merge3) {
  StridedRegionList expected = {StridedRegion(3, StrideDesc{{4, 4}, {2, 2}})};
  checkMerge(input3, expected);
}

BOOST_AUTO_TEST_CASE(Merge4) {
  StridedRegionList expected = {StridedRegion(2, StrideDesc{{4, 4}, {2, 2}})};
  checkMerge(input4, expected);
}

BOOST_AUTO_TEST_CASE(RoundTrip1) { checkRoundTrip(input1); }
BOOST_AUTO_TEST_CASE(RoundTrip2) { checkRoundTrip(input2); }
BOOST_AUTO_TEST_CASE(RoundTrip3) { checkRoundTrip(input3); }
BOOST_AUTO_TEST_CASE(RoundTrip4) { checkRoundTrip(input4); }
BOOST_AUTO_TEST_CASE(RoundTrip5) { checkRoundTrip(input5); }

BOOST_AUTO_TEST_CASE(Slice1) { checkSlice(input1); }
BOOST_AUTO_TEST_CASE(Slice2) { checkSlice(input2); }
BOOST_AUTO_TEST_CASE(Slice3) { checkSlice(input3); }
BOOST_AUTO_TEST_CASE(Slice4) { checkSlice(input4); }
BOOST_AUTO_TEST_CASE(Slice5) { checkSlice(input5); }