/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file tirx/analysis/identify_memcpy.cc
 * \brief Check if a loop nest is equivalent to memcpy
 */

#include <tvm/arith/bound.h>
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/optional.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>

#include <optional>
#include <sstream>
#include <string>
#include <variant>

#include "../../arith/ir_visitor_with_analyzer.h"

namespace tvm {
namespace s_tir {
using namespace tvm::tirx;

std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop,
                                                            arith::AnalyzerObj* analyzer) {
  ffi::Map<Var, arith::IntSet> loop_intervals;
  ffi::Map<PrimVar, Range> loop_ranges;
  PrimExpr total_loop_iterations = 1;

  // Walk through the loop nest, stopping at the first loop whose body
  // is not a loop.
  Stmt stmt = loop;
  while (auto* for_node = stmt.as<ForNode>()) {
    loop_ranges.Set(for_node->loop_var, Range::FromMinExtent(for_node->min, for_node->extent));
    loop_intervals.Set(for_node->loop_var,
                       arith::IntSet::FromMinExtent(for_node->min, for_node->extent));
    total_loop_iterations = total_loop_iterations * for_node->extent;

    stmt = for_node->body;
  }

  BufferStore store;
  if (auto opt = stmt.as<BufferStore>()) {
    store = opt.value();
  } else {
    return static_cast<const std::stringstream&>(
               std::stringstream()
               << "Expected innermost loop to have BufferStore body, but instead found " << stmt)
        .str();
  }

  BufferLoad load;
  if (auto opt = store->value.as<BufferLoad>()) {
    load = opt.value();
  } else {
    return static_cast<const std::stringstream&>(
               std::stringstream()
               << "Expected BufferStore's value to be BufferLoad, but instead found "
               << store->value)
        .str();
  }

  // Now, we have a BufferStore whose value is a BufferLoad.  Because
  // non-flat physical indices are target-dependent, only handle cases
  // where the buffer will be flattened to a 1-d physical buffer.
  ffi::Array<PrimExpr> flattened_dst = store->buffer.OffsetOf(store->indices);
  ffi::Array<PrimExpr> flattened_src = load->buffer.OffsetOf(load->indices);

  if (flattened_dst.size() != 1 || flattened_src.size() != 1) {
    return static_cast<const std::stringstream&>(
               std::stringstream()
               << "Expected flattened dimension of src/dest to be 1, but found"
               << flattened_src.size() << "-d src and " << flattened_dst.size() << "-d dst")
        .str();
  }
  PrimExpr src_index = flattened_src[0];
  PrimExpr dst_index = flattened_dst[0];

  // First check, do the input/output form affine subsets of their
  // respective buffers?
  //
  // For example, should exclude the following, indices are not affine
  //
  // for i in T.serial(16):
  //     B[i] = A[T.abs(i-8)]

  arith::Analyzer analyzer_ref = ffi::GetRef<arith::Analyzer>(analyzer);
  auto src_iter_map = arith::DetectIterMap({src_index}, loop_ranges, IntImm::Bool(true),
                                           arith::IterMapLevel::Bijective, analyzer_ref);
  if (src_iter_map->errors.size()) {
    return static_cast<const std::stringstream&>(std::stringstream()
                                                 << "arith::DetectIterMap(src) returned "
                                                 << src_iter_map->errors.size() << " errors: ["
                                                 << src_iter_map->errors << "]"
                                                 << " for src_index = " << src_index)
        .str();
  }
  auto dst_iter_map = arith::DetectIterMap({dst_index}, loop_ranges, IntImm::Bool(true),
                                           arith::IterMapLevel::Bijective, analyzer_ref);
  if (dst_iter_map->errors.size()) {
    return static_cast<const std::stringstream&>(std::stringstream()
                                                 << "arith::DetectIterMap(dst) returned "
                                                 << dst_iter_map->errors.size() << " errors: ["
                                                 << dst_iter_map->errors << "]"
                                                 << " for dst_index = " << dst_index)
        .str();
  }

  // Second check, are those affine subsets contiguous?  If so, then
  // the index expressions will visit every location between the min
  // and the max.  This checks surjectivity over a linear region,
  // which may not be the same as DetectIterMap's check of
  // surjectivity over the affine subset.
  //
  // For example, should exclude the following, doesn't touch all
  // output locations within the output region touched.
  //
  // for i in T.serial(16):
  //     B[2*i] = A[i]
  //
  // Similarly, should exclude the following, doesn't touch all
  // input locations within the input region touched.
  //
  // for i in T.serial(16):
  //     B[i] = A[2*i]
  total_loop_iterations = analyzer->Simplify(total_loop_iterations);
  auto src_interval = analyzer->int_set(src_index, loop_intervals);
  auto dst_interval = analyzer->int_set(dst_index, loop_intervals);

  if (!src_interval.HasLowerBound() || !src_interval.HasUpperBound()) {
    return static_cast<const std::stringstream&>(std::stringstream()
                                                 << "Expected known bounds for src, but found "
                                                 << src_interval << " for expression " << src_index)
        .str();
  }
  if (!dst_interval.HasLowerBound() || !dst_interval.HasUpperBound()) {
    return static_cast<const std::stringstream&>(std::stringstream()
                                                 << "Expected known bounds for dst, but found "
                                                 << dst_interval << " for expression " << dst_index)
        .str();
  }

  {
    PrimExpr must_prove = total_loop_iterations == src_interval.max() - src_interval.min() + 1;
    PrimExpr simplified = analyzer->Simplify(must_prove);
    if (!analyzer->CanProve(simplified)) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Mismatch between loop iterations (" << total_loop_iterations
                 << ") and number of src indices touched (" << src_interval
                 << ".  Equality to prove simplified to " << simplified)
          .str();
    }
  }
  {
    PrimExpr must_prove = total_loop_iterations == dst_interval.max() - dst_interval.min() + 1;
    PrimExpr simplified = analyzer->Simplify(must_prove);
    if (!analyzer->CanProve(simplified)) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Mismatch between loop iterations (" << total_loop_iterations
                 << ") and number of dst indices touched (" << dst_interval
                 << ".  Equality to prove simplified to " << simplified)
          .str();
    }
  }

  // Third check, is there a transformation applied between the input
  // and output iterators?
  //
  // For example, the following would pass all checks so far, but
  // converts between row-major and column-major layouts, and could
  // not be specified as a memcpy.
  //
  // for i,j in T.grid(4,4):
  //     B[i,j] = A[j,i]

  auto src_iter_sum = src_iter_map->indices[0];
  auto dst_iter_sum = dst_iter_map->indices[0];

  if (src_iter_sum->args.size() != dst_iter_sum->args.size()) {
    return static_cast<const std::stringstream&>(
               std::stringstream()
               << "IterMap for src/dst unpacked to different number of IterSplitExpr: "
               << src_iter_sum->args.size() << " for src, " << dst_iter_sum->args.size()
               << " for dst.  "
               << "IterMaps were detected as src = " << src_iter_sum << ", dst = " << dst_iter_sum)
        .str();
  }
  std::vector<arith::IterSplitExpr> src_iter_terms(src_iter_sum->args.begin(),
                                                   src_iter_sum->args.end());
  std::vector<arith::IterSplitExpr> dst_iter_terms(dst_iter_sum->args.begin(),
                                                   dst_iter_sum->args.end());

  auto make_comparison_tuple = [](const arith::IterSplitExpr& expr) {
    auto as_int_or_zero = [](auto& val) -> int64_t {
      if (auto* as_int = val.template as<IntImmNode>()) {
        return as_int->value;
      } else {
        return 0;
      }
    };
    return std::tuple{
        static_cast<bool>(expr->scale.as<IntImmNode>()),        as_int_or_zero(expr->scale),
        static_cast<bool>(expr->extent.as<IntImmNode>()),       as_int_or_zero(expr->lower_factor),
        static_cast<bool>(expr->lower_factor.as<IntImmNode>()), as_int_or_zero(expr->lower_factor),
    };
  };
  auto sorting_function = [&make_comparison_tuple](const arith::IterSplitExpr& lhs,
                                                   const arith::IterSplitExpr& rhs) -> bool {
    return make_comparison_tuple(lhs) < make_comparison_tuple(rhs);
  };
  std::sort(src_iter_terms.begin(), src_iter_terms.end(), sorting_function);
  std::sort(dst_iter_terms.begin(), dst_iter_terms.end(), sorting_function);

  for (size_t i = 0; i < src_iter_terms.size(); i++) {
    const arith::IterSplitExpr& src_term = src_iter_terms[i];
    const arith::IterSplitExpr& dst_term = dst_iter_terms[i];

    if (!analyzer->CanProve(
            arith::NormalizeIterMapToExpr(src_term->source->source == dst_term->source->source))) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Term " << i << " had different source, src_term->source = " << src_term->source
                 << ", dst_term->source = " << dst_term->source)
          .str();
    }
    if (!analyzer->CanProve(src_term->lower_factor == dst_term->lower_factor)) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Term " << i << " had different lower_factor, src_term->lower_factor = "
                 << src_term->lower_factor
                 << ", dst_term->lower_factor = " << dst_term->lower_factor)
          .str();
    }
    if (!analyzer->CanProve(src_term->extent == dst_term->extent)) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Term " << i << " had different extent, src_term->extent = " << src_term->extent
                 << ", dst_term->extent = " << dst_term->extent)
          .str();
    }
    if (!analyzer->CanProve(src_term->scale == dst_term->scale)) {
      return static_cast<const std::stringstream&>(
                 std::stringstream()
                 << "Term " << i << " had different scale, src_term->scale = " << src_term->scale
                 << ", dst_term->scale = " << dst_term->scale)
          .str();
    }
  }

  BufferRegion src_region(load->buffer, arith::DomainTouched(loop, load->buffer, true, true));
  BufferRegion dst_region(store->buffer, arith::DomainTouched(loop, store->buffer, true, true));

  return MemCpyDetails{src_region, dst_region};
}

std::optional<MemCpyDetails> IdentifyMemCpy(const For& loop, const arith::Analyzer& analyzer) {
  auto result = IdentifyMemCpyImpl(loop, analyzer.get());
  if (auto* ptr = std::get_if<MemCpyDetails>(&result)) {
    return *ptr;
  } else {
    return std::nullopt;
  }
}

// Expose the IdentifyMemCpy functionality to Python API for purpose of unit testing.
TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("s_tir.analysis._identify_memcpy", [](const Stmt& stmt) {
    ffi::Array<ffi::ObjectRef> output;

    struct Visitor : arith::IRVisitorWithAnalyzer {
      explicit Visitor(ffi::Array<ffi::ObjectRef>* output) : output(output) {}
      ffi::Array<ffi::ObjectRef>* output;

     private:
      using IRVisitorWithAnalyzer::VisitStmt_;
      void VisitStmt_(const ForNode* op) override {
        For loop = ffi::GetRef<For>(op);
        auto result = IdentifyMemCpyImpl(loop, Visitor::analyzer_.get());
        if (auto* ptr = std::get_if<MemCpyDetails>(&result)) {
          output->push_back(ffi::Array{ptr->source, ptr->dest});
        } else if (auto* ptr = std::get_if<std::string>(&result)) {
          output->push_back(StringImm(*ptr));
        } else {
          TVM_FFI_THROW(InternalError) << "Internal error, unhandled std::variant type";
        }

        IRVisitorWithAnalyzer::VisitStmt_(op);
      }
    };

    Visitor visitor(&output);
    visitor(stmt);

    return output;
  });
}

}  // namespace s_tir
}  // namespace tvm
