/*
 * 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_async_dma.cc
 */

#include <tvm/arith/analyzer.h>
#include <tvm/arith/bound.h>
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/analysis.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/buffer.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>

#include <optional>
#include <set>

#include "../../arith/ir_mutator_with_analyzer.h"
#include "../../tirx/transform/ir_utils.h"

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

class AsyncDMALowerer : public arith::IRMutatorWithAnalyzer {
 public:
  explicit AsyncDMALowerer(bool dma_bypass_cache, const arith::Analyzer& analyzer)
      : IRMutatorWithAnalyzer(analyzer), dma_bypass_cache_(dma_bypass_cache) {}

  // TODO(leiwang1999): split lower async DMA support for CUDA and Hexagon Backend
  Stmt VisitStmt_(const ForNode* loop) final {
    // if for loop is not within async_commit_queue_scope
    if (!async_queue_id_.has_value()) {
      return arith::IRMutatorWithAnalyzer::VisitStmt_(loop);
    }

    // if for loop is not a memcpy of a contiguous region, it might be a cuda cp.async behavior
    std::optional<s_tir::MemCpyDetails> mem_copy =
        s_tir::IdentifyMemCpy(ffi::GetRef<For>(loop), ffi::GetRef<arith::Analyzer>(analyzer_));
    if (!mem_copy.has_value() || mem_copy->dest->region.size() != 1 ||
        mem_copy->source->region.size() != 1) {
      return arith::IRMutatorWithAnalyzer::VisitStmt_(loop);
    }

    // now that we are about to perform the `copy` transform
    // save queue ID for inspection in `wait` transform
    // and, increment the number of DMA copies in the group
    queue_ids_.insert(async_queue_id_.value());
    dmas_in_group_++;

    tvm::PrimExpr src_min = mem_copy->source->region[0]->min;
    tvm::PrimExpr dst_min = mem_copy->dest->region[0]->min;
    tvm::PrimExpr dst_extent = mem_copy->dest->region[0]->extent;

    auto src = BufferLoad(mem_copy->source->buffer, {src_min});
    auto dst = BufferLoad(mem_copy->dest->buffer, {dst_min});
    PrimExpr dst_nbytes = dst_extent * static_cast<int>(src.ty().StorageBytes());
    return Evaluate(
        Call(PrimType::Int(32), builtin::dma_copy(),
             ffi::Array<Expr>{PrimExpr(async_queue_id_.value()),
                              Call(mem_copy->dest->buffer->data->ty, builtin::address_of(),
                                   ffi::Array<Expr>{dst}, Attrs(), {}, Span()),
                              Call(mem_copy->source->buffer->data->ty, builtin::address_of(),
                                   ffi::Array<Expr>{src}, Attrs(), {}, Span()),
                              dst_nbytes, PrimExpr(dma_bypass_cache_)},
             Attrs(), {}, Span())
            .as_or_throw<PrimExpr>());
  }

  Stmt VisitStmt_(const AttrStmtNode* op) final {
    // populate analyzer knowledge of loop iterators
    auto previsit = arith::IRMutatorWithAnalyzer::VisitStmt_(op);

    // Convert this, for example:
    // attr [0] "async_wait_queue_scope" = 0;
    // attr [0] "async_wait_inflight_count" = 0;
    //
    // To this:
    // @tirx.dma_wait(
    //   0, /* queue id */
    //   0, /* in flight count */
    //   dtype=int32
    // )
    if (op->attr_key == s_tir::attr::async_wait_queue_scope) {
      // get queue ID
      auto queue_id_node = op->value.as<IntImmNode>();
      TVM_FFI_ICHECK(queue_id_node);
      int queue_id = queue_id_node->value;

      // abort if we have not seen this queue ID in `copy` transform
      if (queue_ids_.find(queue_id) == queue_ids_.end()) {
        DLOG(INFO) << "AsyncDMALowerer exiting because the queue ID observed in the "
                      "`async_wait_queue_scope` transform has not been previously observed in the "
                      "`async_commit_queue_scope` transform";
        return previsit;
      }

      auto async_wait = op->body.as<AttrStmtNode>();
      if (!async_wait || async_wait->attr_key != s_tir::attr::async_wait_inflight_count) {
        DLOG(INFO) << "AsyncDMALowerer exiting because the body of the `AttrStmtNode` with key "
                      "`async_wait_queue_scope` does not contain an `AttrStmtNode` with key "
                      "`async_wait_inflight_count`";
        return previsit;
      }
      auto call_dma_wait = Evaluate(
          Call(PrimType::Int(32), builtin::dma_wait(), {PrimExpr(queue_id), async_wait->value})
              .as_or_throw<PrimExpr>());

      // concatenate the call with the body and return
      return SeqStmt({call_dma_wait, arith::IRMutatorWithAnalyzer::VisitStmt(async_wait->body)});

      // Convert this, for example:
      // attr [0] "async_commit_queue_scope" = 0;
      // attr [0] "async_scope" = 1;
      // for (ax0: int32, 0, 128) {
      //   A_global[ax0] = A[ax0]
      // }
      //
      // To this:
      // @tirx.dma_copy(
      //   0, /* queue id */
      //   @tirx.address_of(A_global[0], dtype=handle),
      //   @tirx.address_of(A[0], dtype=handle),
      //   128, /* size */
      //   dtype=int32
      // )
    } else if (op->attr_key == s_tir::attr::async_commit_queue_scope) {
      // get queue ID
      auto queue_id_node = op->value.as<IntImmNode>();
      TVM_FFI_ICHECK(queue_id_node);
      async_queue_id_ = queue_id_node->value;
      auto result = arith::IRMutatorWithAnalyzer::VisitStmt_(op);
      if (dmas_in_group_ > 1) {
        auto call_dma_start_group = Evaluate(
            Call(PrimType::Int(32), builtin::dma_start_group(), {PrimExpr(async_queue_id_.value())})
                .as_or_throw<PrimExpr>());
        auto call_dma_end_group = Evaluate(
            Call(PrimType::Int(32), builtin::dma_end_group(), {PrimExpr(async_queue_id_.value())})
                .as_or_throw<PrimExpr>());
        result = SeqStmt({call_dma_start_group, result, call_dma_end_group});
      }

      async_queue_id_ = std::nullopt;
      dmas_in_group_ = 0;
      return result;
    }
    return arith::IRMutatorWithAnalyzer::VisitStmt_(op);
  }

 private:
  int dmas_in_group_ = 0;
  std::set<int> queue_ids_;
  std::optional<int> async_queue_id_ = std::nullopt;
  bool dma_bypass_cache_;
  ffi::Map<Var, Range> input_iters = ffi::Map<Var, Range>();
};

namespace transform {

Pass LowerAsyncDMA() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    auto fptr = f.CopyOnWrite();
    arith::Analyzer analyzer;
    bool dma_bypass_cache =
        ctx->GetConfig<bool>("tirx.experimental_dma_bypass_cache", false).value();
    fptr->body = AsyncDMALowerer(dma_bypass_cache, analyzer)(std::move(fptr->body));
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "s_tir.LowerAsyncDMA", {});
}

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

}  // namespace s_tir
}  // namespace tvm
