/*
 * 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.
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/transform.h>

#include "../utils.h"

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

class ThreadExtentChecker : private StmtVisitor {
 public:
  static bool Check(const Stmt& stmt, int thread_warp_size) {
    try {
      TVM_FFI_ICHECK(thread_warp_size > 0);
      ThreadExtentChecker checker(thread_warp_size);
      checker.VisitStmt(stmt);
      return true;
    } catch (const std::exception&) {
      return false;
    }
  }

 private:
  explicit ThreadExtentChecker(int thread_warp_size) : thread_warp_size_(thread_warp_size) {}

  void VisitStmt_(const ForNode* loop) {
    runtime::ThreadScope thread_scope = GetThreadScope(loop);
    if (IsThreadIdx(thread_scope)) {
      if (const int64_t* p_ext = GetLoopIntExtent(loop)) {
        int64_t ext = *p_ext;
        if (thread_scope.dim_index == 0) {
          std::swap(thread_idx_x, ext);
          StmtVisitor::VisitStmt_(loop);
          std::swap(thread_idx_x, ext);
        } else if (thread_scope.dim_index == 1) {
          std::swap(thread_idx_y, ext);
          StmtVisitor::VisitStmt_(loop);
          std::swap(thread_idx_y, ext);
        } else if (thread_scope.dim_index == 2) {
          std::swap(thread_idx_z, ext);
          StmtVisitor::VisitStmt_(loop);
          std::swap(thread_idx_z, ext);
        } else {
          StmtVisitor::VisitStmt_(loop);
        }
        return;
      } else {
        throw std::runtime_error("Dynamic thread extent");
      }
    }
    StmtVisitor::VisitStmt_(loop);
  }

  void VisitStmt_(const SBlockNode* block) {
    int old_thread_idx_x = thread_idx_x;
    if (block->annotations.count(s_tir::attr::warp_execution)) {
      thread_idx_x = thread_warp_size_;
    }
    if (ffi::Optional<int64_t> low_inclusive =
            GetAnn<int64_t>(block, s_tir::attr::meta_schedule_thread_extent_low_inclusive)) {
      if (ffi::Optional<int64_t> high_inclusive =
              GetAnn<int64_t>(block, s_tir::attr::meta_schedule_thread_extent_high_inclusive)) {
        int64_t low = low_inclusive.value();
        int64_t high = high_inclusive.value();
        int64_t thread_extent_product = thread_idx_x * thread_idx_y * thread_idx_z;
        if (!(low <= thread_extent_product && thread_extent_product <= high)) {
          throw std::runtime_error("Thread extent");
        }
      }
    }
    StmtVisitor::VisitStmt_(block);
    thread_idx_x = old_thread_idx_x;
  }

  int64_t thread_idx_x = 1;
  int64_t thread_idx_y = 1;
  int64_t thread_idx_z = 1;
  int thread_warp_size_ = -1;
};

}  // namespace s_tir
}  // namespace tvm

namespace tvm {
namespace s_tir {
namespace meta_schedule {

/*! \brief Extract attribute from a target. */
IntImm Extract(const Target& target, const char* name) {
  TVM_FFI_ICHECK(target.defined());
  if (ffi::Optional<int64_t> v = target->GetAttr<int64_t>(name)) {
    return IntImm::Int64(v.value());
  }
  TVM_FFI_THROW(AttributedError) << "\"" << name << "\" is not defined in the target";
  throw;
}

/*! \brief Verify the correctness of the generated GPU code. */
class VerifyGPUCodeNode : public PostprocNode {
 public:
  Target target_{ffi::UnsafeInit()};
  ffi::Map<ffi::String, PrimExpr> target_constraints_{ffi::UnsafeInit()};
  int thread_warp_size_ = -1;

  void InitializeWithTuneContext(const TuneContext& context) final {
    TVM_FFI_ICHECK(context->target.has_value());
    this->target_ = context->target.value();
    this->target_constraints_ = ffi::Map<ffi::String, PrimExpr>{
        {"max_shared_memory_per_block", Extract(this->target_, "max_shared_memory_per_block")},
        {"max_threads_per_block", Extract(this->target_, "max_threads_per_block")},
        {"max_vthread", IntImm::Int32(8)},
        {"max_vector_bytes", IntImm::Int32(16)},
    };
    thread_warp_size_ = static_cast<int>(Extract(this->target_, "thread_warp_size")->value);
  }

  bool Verify(const IRModule& mod) const {
    for (const auto& kv : mod->functions) {
      if (auto prim_func = kv.second.as<tirx::PrimFunc>()) {
        if (!s_tir::VerifyGPUCode(prim_func.value(), this->target_constraints_)) {
          return false;
        }
      }
    }
    return true;
  }

  bool Apply(const s_tir::Schedule& sch) final {
    IRModule mod = sch->mod();
    for (const auto& kv : mod->functions) {
      const GlobalVar& g_var = kv.first;
      const BaseFunc& base_func = kv.second;
      if (const auto* prim_func = base_func.as<tirx::PrimFuncNode>()) {
        if (!s_tir::ThreadExtentChecker::Check(prim_func->body, thread_warp_size_)) {
          return false;
        }
        IRModule lowered{ffi::UnsafeInit()};
        try {
          auto pass_list = ffi::Array<tvm::transform::Pass>();
          // Phase 1
          pass_list.push_back(s_tir::transform::LowerCrossThreadReduction());
          pass_list.push_back(s_tir::transform::LowerInitBlock());
          pass_list.push_back(s_tir::transform::PlanAndUpdateBufferAllocationLocation());
          pass_list.push_back(s_tir::transform::ConvertBlocksToOpaque());
          pass_list.push_back(s_tir::transform::LiftThreadBinding());
          pass_list.push_back(s_tir::transform::ManifestSharedMemoryLocalStage());
          pass_list.push_back(s_tir::transform::CompactBufferAllocation());
          pass_list.push_back(tirx::transform::StmtSimplify());
          pass_list.push_back(s_tir::transform::LowerAutoCopy());
          pass_list.push_back(s_tir::transform::UnifyThreadBinding());
          pass_list.push_back(s_tir::transform::LowerMatchBuffer());
          pass_list.push_back(s_tir::transform::InjectSoftwarePipeline());
          pass_list.push_back(s_tir::transform::LowerOpaqueBlock());
          pass_list.push_back(tirx::transform::FlattenBuffer());
          pass_list.push_back(tirx::transform::BF16ComputeLegalize());
          pass_list.push_back(tirx::transform::NarrowDataType(32));
          pass_list.push_back(tirx::transform::StmtSimplify());
          // Phase 2
          pass_list.push_back(tirx::transform::VectorizeLoop(true));
          pass_list.push_back(s_tir::transform::InjectVirtualThread());
          pass_list.push_back(s_tir::transform::InjectDoubleBuffer());
          pass_list.push_back(tirx::transform::StorageRewrite());
          pass_list.push_back(s_tir::transform::MergeSharedMemoryAllocations());
          pass_list.push_back(tirx::transform::LowerIntrin());
          // Convert Function to IRModule
          tvm::transform::PassContext pass_ctx = tvm::transform::PassContext::Current();
          tirx::PrimFunc f = WithAttr(ffi::GetRef<tirx::PrimFunc>(prim_func), "global_symbol",
                                      ffi::String(g_var->name_hint));
          f = WithAttr(f, tvm::attr::kTarget, this->target_);  // Required for LowerIntrin
          bool noalias = pass_ctx->GetConfig<bool>("tirx.noalias", true).value();
          if (noalias) {
            f = WithAttr(std::move(f), "tirx.noalias", true);
          }
          IRModule mod =
              IRModule(ffi::Map<GlobalVar, BaseFunc>({{GlobalVar(g_var->name_hint), f}}));
          lowered = tvm::transform::Sequential(pass_list)(std::move(mod));
        } catch (const std::exception&) {
          return false;
        }
        if (!Verify(lowered)) {
          return false;
        }
      }
    }
    return true;
  }

  Postproc Clone() const {
    ffi::ObjectPtr<VerifyGPUCodeNode> n = ffi::make_object<VerifyGPUCodeNode>(*this);
    n->target_constraints_ = this->target_constraints_;
    return Postproc(n);
  }

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<VerifyGPUCodeNode>();
  }

  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.VerifyGPUCode", VerifyGPUCodeNode,
                                    PostprocNode);
};

Postproc Postproc::VerifyGPUCode() {
  ffi::ObjectPtr<VerifyGPUCodeNode> n = ffi::make_object<VerifyGPUCodeNode>();
  return Postproc(n);
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  VerifyGPUCodeNode::RegisterReflection();
  refl::GlobalDef().def("s_tir.meta_schedule.PostprocVerifyGPUCode", Postproc::VerifyGPUCode);
}

}  // namespace meta_schedule
}  // namespace s_tir
}  // namespace tvm
