/*
 * 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 lower_match_buffer.cc
 * \brief The pass for lowering match_buffer.
 */

#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/function.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>

#include "../../tirx/ir/functor_common.h"
#include "../../tirx/transform/ir_utils.h"

namespace tvm {
namespace s_tir {
using namespace tvm::tirx;
class MatchBufferLower : public StmtExprMutator {
 public:
  explicit MatchBufferLower(const PrimFunc& func) {
    for (const Var& param : func->params) {
      // Mark input var as const variable.
      auto prim_type = param->ty.as<PrimType>();
      if (prim_type) {
        var_map_.Set(param, param.as_or_throw<PrimExpr>());
      }
    }
  }

 private:
  Stmt VisitStmt_(const SBlockNode* op) final {
    for (const MatchBufferRegion& match_buffer : op->match_buffers) {
      CheckAndUpdateVarMap(match_buffer);
    }
    // After CheckAndUpdateVarMap sets up match_buffers_ and var_map_,
    // the base class visit may remap buffer objects via VisitBufferDef
    // (e.g., remapping elem_offset). We need match_buffers_ to also
    // contain the remapped buffer keys for lookups in the body.
    // Snapshot current match_buffers_ keys before base visit.
    std::vector<Buffer> orig_buffers;
    for (const auto& kv : match_buffers_) {
      orig_buffers.push_back(kv.first);
    }
    Stmt stmt = StmtExprMutator ::VisitStmt_(op);
    // Add remapped buffer keys to match_buffers_
    for (const Buffer& orig_buf : orig_buffers) {
      if (auto remap_it = buffer_remap_.find(orig_buf); remap_it != buffer_remap_.end()) {
        Buffer remapped = (*remap_it).second;
        if (!match_buffers_.count(remapped)) {
          match_buffers_.Set(remapped, match_buffers_[orig_buf]);
        }
      }
    }
    op = stmt.as<SBlockNode>();
    TVM_FFI_ICHECK(op != nullptr);
    ffi::Array<BufferRegion> reads =
        op->reads.Map(std::bind(&MatchBufferLower::VisitBufferRegion, this, std::placeholders::_1));
    ffi::Array<BufferRegion> writes = op->writes.Map(
        std::bind(&MatchBufferLower::VisitBufferRegion, this, std::placeholders::_1));

    if (reads.same_as(op->reads) && writes.same_as(op->writes) && op->match_buffers.empty()) {
      return stmt;
    } else {
      auto n = CopyOnWrite(op);
      n->match_buffers = {};
      n->reads = std::move(reads);
      n->writes = std::move(writes);
      return Stmt(n);
    }
  }

  Stmt VisitStmt_(const ForNode* op) final {
    analyzer_->Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
    return StmtExprMutator::VisitStmt_(op);
  }

  Expr VisitExpr_(const VarNode* op) final {
    Var v = ffi::GetRef<Var>(op);
    auto it = var_map_.find(v);
    if (it != var_map_.end()) {
      return (*it).second;
    } else {
      return v;
    }
  }

  Stmt VisitStmt_(const BufferStoreNode* op) final {
    // Save the original buffer before base class mutation may remap it
    Buffer orig_buffer = op->buffer;
    Stmt stmt = StmtExprMutator::VisitStmt_(op);
    op = stmt.as<BufferStoreNode>();
    TVM_FFI_ICHECK(op != nullptr);

    // Look up using original buffer (before VisitBufferDef may have remapped it)
    auto it = match_buffers_.find(orig_buffer);
    if (it == match_buffers_.end()) {
      return stmt;
    } else {
      const Buffer& buffer = (*it).first;
      const BufferRegion& source = (*it).second;

      auto n = CopyOnWrite(op);
      n->indices = ConvertIndices(MatchBufferRegion(buffer, source), op->indices);
      n->buffer = source->buffer;
      TVM_FFI_ICHECK(!op->predicate.has_value())
          << "Predicated buffer store is not currently supported in lower match buffer pass.";
      return Stmt(n);
    }
  }

  Expr VisitExpr_(const BufferLoadNode* op) final {
    // Save the original buffer before base class mutation may remap it
    Buffer orig_buffer = op->buffer;
    PrimExpr expr = StmtExprMutator::VisitExpr_(op).as_or_throw<PrimExpr>();
    op = expr.as<BufferLoadNode>();
    TVM_FFI_ICHECK(op != nullptr);

    auto it = match_buffers_.find(orig_buffer);
    if (it == match_buffers_.end()) {
      return expr;
    } else {
      const Buffer& buffer = (*it).first;
      const BufferRegion& source = (*it).second;
      ffi::Array<PrimExpr> indices = ConvertIndices(MatchBufferRegion(buffer, source), op->indices);
      TVM_FFI_ICHECK(!op->predicate.has_value())
          << "Predicated buffer load is not currently supported in lower match buffer pass.";
      return BufferLoad(source->buffer, indices);
    }
  }

  BufferRegion VisitBufferRegion(const BufferRegion& buffer_region) {
    const Buffer& buffer = buffer_region->buffer;
    auto it = match_buffers_.find(buffer);
    if (it == match_buffers_.end()) {
      return buffer_region;
    } else {
      const BufferRegion& source = (*it).second;
      Region region = ConvertRegion(MatchBufferRegion(buffer, source), buffer_region->region);
      return BufferRegion(source->buffer, std::move(region));
    }
  }

 private:
  void CheckAndUpdateVarMap(const MatchBufferRegion& match_buffer) {
    // Step.1. Check
    const Buffer& buffer = match_buffer->buffer;
    const BufferRegion& source = VisitBufferRegion(match_buffer->source);
    const Buffer& source_buffer = source->buffer;

    // Step.1.1. Check scope & dtype
    TVM_FFI_ICHECK_EQ(buffer.scope(), source_buffer.scope())
        << "MatchBuffer " << buffer << " scope mismatch:" << buffer.scope() << "vs."
        << source_buffer.scope();
    TVM_FFI_ICHECK_EQ(buffer->dtype, source_buffer->dtype)
        << "MatchBuffer " << buffer << " data type mismatch:" << buffer->dtype << "vs."
        << source_buffer->dtype;

    // Step.1.2. Check data alignment
    if (source_buffer->data_alignment % buffer->data_alignment != 0) {
      LOG(WARNING) << "Trying to bind buffer to another one with lower alignment requirement "
                   << " required alignment=" << buffer->data_alignment
                   << ", provided alignment=" << source_buffer->data_alignment;
    }
    if (is_zero(buffer->elem_offset)) {
      TVM_FFI_ICHECK(is_zero(source_buffer->elem_offset))
          << "Trying to bind a Buffer with offset into one without offset "
          << " required elem_offset=" << buffer->elem_offset
          << ", provided elem_offset=" << source_buffer->elem_offset;
    }

    // Step.2. Update
    match_buffers_.Set(buffer, source);
    // Step.2.1. Update buffer data
    Bind(buffer->data, source_buffer->data, buffer->name + ".data");

    // Step.2.2. Update element offset
    // We use the ElemOffset method to avoid duplicating the index calculation.
    {
      ffi::Array<PrimExpr> indices;
      indices.reserve(source->region.size());
      for (const Range& range : source->region) {
        indices.push_back(range->min);
      }

      ffi::Array<PrimExpr> buffer_start_indices = source_buffer->ElemOffset(indices);
      if (buffer_start_indices.size() == 1) {
        Bind(buffer->elem_offset, buffer_start_indices[0], buffer->name + ".elem_offset");
        TVM_FFI_ICHECK(
            analyzer_->CanProve(truncmod(buffer->elem_offset, buffer->offset_factor) == 0))
            << "The source elem_offset " << buffer_start_indices[0]
            << " does not satisfy the offset_factor " << buffer->offset_factor << ".";
      } else {
        // Non-zero elem_offset is ill-defined for non-flat memory.
        // If needed in the future, will require `ffi::Array<PrimExpr>
        // elem_offsets`, with one offset for each flattened index.
        Bind(buffer->elem_offset, IntImm(buffer->elem_offset.ty(), 0));
      }
    }

    // Step 2.3. Check and update strides
    // Check if target buffer strides are defined
    TVM_FFI_ICHECK(source->region.size() >= buffer->shape.size());
    int offset = source->region.size() - buffer->shape.size();
    if (!buffer->strides.empty()) {
      TVM_FFI_ICHECK_EQ(buffer->strides.size(), buffer->shape.size());
      if (source_buffer->strides.empty()) {
        PrimExpr stride = MakeConst(buffer->strides.back().ty(), 1);
        for (size_t i = buffer->shape.size(); i > 0; --i) {
          const PrimExpr& shape = source_buffer->shape[i - 1 + offset];
          Bind(buffer->strides[i - 1], stride, buffer->name + ".strides_" + std::to_string(i - 1));
          stride *= shape;
        }
      } else {
        TVM_FFI_ICHECK_EQ(buffer->shape.size() + offset, source_buffer->strides.size());
        for (size_t i = buffer->shape.size(); i > 0; --i) {
          const PrimExpr& stride = source_buffer->strides[i - 1 + offset];
          Bind(buffer->strides[i - 1], stride, buffer->name + ".strides_" + std::to_string(i - 1));
        }
      }
    }

    // Step 2.4. Check and update shape
    for (size_t i = 0; i < buffer->shape.size(); ++i) {
      const Range& range = source->region[i + offset];
      Bind(buffer->shape[i], range->extent, buffer->name + ".shape_" + std::to_string(i));
    }
  }

  void Bind(const Expr& arg, Expr value, const std::string& arg_name = "argument") {
    auto arg_prim = arg.as<PrimExpr>();
    auto value_prim = value.as<PrimExpr>();
    if (arg_prim && value_prim && arg_prim.value().ty() != value_prim.value().ty()) {
      PrimType arg_ty = arg_prim.value().ty();
      PrimType value_ty = value_prim.value().ty();
      bool same_lanes = arg_ty.lanes() == value_ty.lanes();
      if (arg_ty.MatchesCode(DLDataTypeCode::kDLInt) &&
          value_ty.MatchesCode(DLDataTypeCode::kDLInt) && same_lanes) {
        value = cast(arg_ty, value_prim.value());
      } else {
        TVM_FFI_ICHECK_EQ(arg_ty->dtype, value_ty->dtype)
            << "The data type mismatched: " << arg_ty->dtype << " vs. " << value_ty->dtype;
      }
    }
    // Handle recursive case
    value = Substitute(std::move(value), var_map_);
    if (arg->IsInstance<VarNode>()) {
      Var v = arg.as_or_throw<Var>();
      auto it = var_map_.find(v);
      if (it == var_map_.end()) {
        var_map_.Set(v, value);
        if (auto prim_value = value.as<PrimExpr>()) {
          analyzer_->Bind(v, prim_value.value());
        }
      } else {
        AssertBinding((*it).second, value, arg_name);
      }
    } else {
      AssertBinding(arg, value, arg_name);
    }
  }

  void AssertBinding(const Expr& lhs, const Expr& rhs, const std::string& arg_name = "argument") {
    if (auto lhs_prim = lhs.as<PrimExpr>()) {
      PrimExpr rhs_prim = rhs.as_or_throw<PrimExpr>();
      TVM_FFI_ICHECK(analyzer_->CanProve(lhs_prim.value() == rhs_prim))
          << "The buffer match constraint for " << arg_name << " unmet: " << lhs << "==" << rhs
          << ".";
    } else {
      TVM_FFI_ICHECK(ffi::StructuralEqual()(lhs, rhs))
          << "The buffer match constraint for " << arg_name << " unmet: " << lhs << "==" << rhs
          << ".";
    }
  }

 private:
  /*! \brief Buffer region mapping. */
  ffi::Map<Buffer, BufferRegion> match_buffers_;
  /*! \brief Var mapping for buffer signature (data, strides, element_offset, etc.) */
  ffi::Map<Var, Expr> var_map_;
  /*! \brief The analyzer */
  arith::Analyzer analyzer_;
};

namespace transform {

Pass LowerMatchBuffer() {
  auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
    auto fptr = f.CopyOnWrite();
    fptr->body = MatchBufferLower(f)(std::move(fptr->body));
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "s_tir.LowerMatchBuffer", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("s_tir.transform.LowerMatchBuffer", LowerMatchBuffer);
}

}  // namespace transform

}  // namespace s_tir
}  // namespace tvm
