/*
 * 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/arith/analyzer.h>
#include <tvm/arith/iter_affine_map.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/s_tir/transform.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>

#include "../../arith/const_fold.h"
#include "../../arith/pattern_match.h"

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

class PTXRewriter : public StmtMutator {
 public:
  Stmt AddAllocationsIfNeeded(Stmt body) {
    if (!needs_buffer || has_buffer_2) {
      return body;
    }
    EnsureBuffers();
    body = SeqStmt::Flatten(AllocBuffer(predicate_buffer), AllocBuffer(addr_buffer), body);
    has_buffer_2 = true;
    return body;
  }

  Stmt VisitStmt_(const AllocBufferNode* op) final {
    Stmt result = StmtMutator::VisitStmt_(op);
    if (needs_buffer && !has_buffer_2) {
      EnsureBuffers();
      has_buffer_2 = true;
      result = SeqStmt::Flatten(AllocBuffer(predicate_buffer), AllocBuffer(addr_buffer), result);
    }
    return result;
  }

  Stmt VisitStmt_(const BufferStoreNode* store) final {
    Stmt result = StmtMutator::VisitStmt_(store);
    BufferVar load_buffer = store->buffer;
    PrimExpr load_value = store->value;
    // const BufferLoadNode* gload = load_value.as<BufferLoadNode>(); // take
    // the place of instance of
    const CallNode* call = load_value.as<CallNode>();
    if (call != nullptr) {
      const OpNode* op = call->op.as<OpNode>();
      if (op != nullptr && op->name == "tirx.if_then_else") {
        PrimExpr predicate = call->args[0].as_or_throw<PrimExpr>();
        PrimExpr lhs = call->args[1].as_or_throw<PrimExpr>();
        PrimExpr rhs = call->args[2].as_or_throw<PrimExpr>();
        PrimExpr global_addr, local_addr;
        const BufferLoadNode* load = lhs.as<BufferLoadNode>();
        PrimExpr imm_value = rhs;
        if (load == nullptr) {
          load = rhs.as<BufferLoadNode>();
          imm_value = lhs;
          if (load == nullptr) {
            return result;
          }
        }
        global_addr = load->indices[0];
        const RampNode* ramp = global_addr.as<RampNode>();
        if (ramp != nullptr) {
          return result;
        }
        EnsureBuffers();
        needs_buffer = true;
        local_addr = store->indices[0];
        BufferStore addr_store(addr_buffer, global_addr, {IntImm::Int32(0)});
        BufferStore local_addr_store(addr_buffer, local_addr, {IntImm::Int32(1)});
        BufferStore predicate_store(predicate_buffer, predicate, {IntImm::Int32(0)});
        PrimExpr new_lhs, new_rhs, new_predicate, new_indice;
        new_lhs = BufferLoad(load->buffer, {BufferLoad(addr_buffer, {IntImm::Int32(0)})});
        new_rhs = IntImm::Int32(0);
        new_predicate = BufferLoad(predicate_buffer, {IntImm::Int32(0)});
        new_indice = BufferLoad(addr_buffer, {IntImm::Int32(1)});
        BufferStore value_store(store->buffer, imm_value, {new_indice});
        static const Op& ptx_ldg32_op = Op::Get("tirx.s_tir.ldg32");
        Evaluate ptx_load(Call(store->buffer->dtype, ptx_ldg32_op,
                               {store->buffer.data(), new_predicate, new_lhs, new_indice})
                              .as_or_throw<PrimExpr>());
        ffi::Array<Stmt> tmp_seq = {addr_store, local_addr_store, predicate_store, value_store,
                                    ptx_load};
        SeqStmt seq_stmt = SeqStmt(tmp_seq);
        return seq_stmt;
      }
    }
    return result;
  }

  void EnsureBuffers() {
    if (has_buffer_1) {
      return;
    }
    has_buffer_1 = true;
    // addr[0] -> global_addr /  addr[1] -> local_addr
    addr_buffer = decl_buffer({IntImm::Int32(2)}, PrimType::Int(32), "addr", "local");
    predicate_buffer = decl_buffer({IntImm::Int32(1)}, PrimType::Bool(), "predicate", "local");
  }

  bool has_buffer_1 = false, has_buffer_2 = false;
  bool needs_buffer = false;
  BufferVar addr_buffer, predicate_buffer;
};

namespace transform {

Pass InjectPTXLDG32(bool enable_inject_ptx_intrin) {
  auto pass_func = [enable_inject_ptx_intrin](PrimFunc f, IRModule m, PassContext ctx) {
    if (enable_inject_ptx_intrin) {
      auto target = f->GetAttr<Target>("target");
      if (!target.has_value() || target.value()->kind->name != "cuda") {
        return f;
      }
      auto* n = f.CopyOnWrite();
      PTXRewriter rewriter;
      Stmt body = rewriter(n->body);
      n->body = rewriter.AddAllocationsIfNeeded(body);
      // inject ptx
    }
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectPTXLDG32", {});
}

// The pass can now be invoked via the pass infrastructure, but we also add a
// Python binding for it
TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("s_tir.transform.InjectPTXLDG32", InjectPTXLDG32);
}

}  // namespace transform
}  // namespace s_tir
}  // namespace tvm
