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

/*!
 * \brief Replace copy from global to shared with async copy
 * \file inject_ptx_async_copy.cc
 */
#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/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt_functor.h>

#include "../../tirx/ir/buffer_common.h"
#include "storage_access.h"
#include "tvm/s_tir/stmt.h"
#include "tvm/tirx/stmt.h"

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

class PTXAsyncCopyInjector : public StmtMutator {
 public:
  Stmt VisitStmt_(const AttrStmtNode* attr) {
    if (attr->attr_key == s_tir::attr::async_scope) {
      TVM_FFI_ICHECK(in_async == false) << "Nested async scopes not supported";
      in_async = true;
      auto body = this->VisitStmt(attr->body);
      in_async = false;
      return body;
    }
    return StmtMutator::VisitStmt_(attr);
  }

  Stmt InjectPTX(const BufferLoadNode* load, const BufferStoreNode* store, bool predicated = false,
                 PrimExpr predicate_value = PrimExpr()) {
    if (load->buffer.scope() == "global") {
      TVM_FFI_ICHECK(load->indices.size() == 1 && store->indices.size() == 1);
      TVM_FFI_ICHECK(load->indices[0].ty().lanes() == store->indices[0].ty().lanes());

      const int indices_lanes = load->indices[0].ty().lanes();
      const int bytes = indices_lanes * static_cast<int>(load->buffer->dtype.StorageBytes());

      if (bytes == 4 || bytes == 8 || bytes == 16) {
        auto dst_elem_type = GetPointerType(store->buffer->data->ty);
        auto src_elem_type = GetPointerType(load->buffer->data->ty);
        TVM_FFI_ICHECK(dst_elem_type.has_value() && src_elem_type.has_value())
            << "Both store and load buffer should have a pointer type annotation.";

        int index_factor = 1;
        if (dst_elem_type.value() != src_elem_type.value()) {
          // The only case where src and dst have different dtypes is when the dst shared memory
          // is a byte buffer generated by merging dynamic shared memory.
          TVM_FFI_ICHECK(store->buffer.scope() == "shared.dyn");
          TVM_FFI_ICHECK((dst_elem_type.value() == PrimType::UInt(8)));
          // BufferStore/Load have the "pointer reinterpret" semantics according to their
          // "value" dtype. Their "indices" are supposed to be applied after such pointer cast,
          // for example: ((*float16)(byte_buffer))[buffer->indices] = fp16_value;
          // To replace BufferStore/Load with cp.async, we need to multiply the store index by
          // the byte size of the "value" dtype, to get the correct offset into the byte buffer.
          index_factor = static_cast<int>(src_elem_type.value().StorageBytes());
        }

        if (indices_lanes == 1) {
          auto src_offset = load->indices[0];
          auto dst_offset = store->indices[0];
          ffi::Array<Expr> args = {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
                                   load->buffer->data, src_offset, PrimExpr(bytes)};
          // use arguments size to indicate whether or not to use predicated cp.async
          if (predicated) {
            args.push_back(predicate_value);
          }
          static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw");
          return Evaluate(
              Call(store->buffer->dtype, ptx_cp_async_op, args).as_or_throw<PrimExpr>());
        }

        // Predicated load don't support vectorized indexing.
        if (!predicated) {
          // Only some vectorized indexing patterns are supported for now.
          auto src_offset = [=]() -> PrimExpr {
            if (load->indices[0]->IsInstance<RampNode>()) {
              return load->indices[0].as<RampNode>()->base;
            }
            return PrimExpr();
          }();

          auto dst_offset = [=]() -> PrimExpr {
            if (store->indices[0].as<RampNode>()) {
              return store->indices[0].as<RampNode>()->base;
            } else if (store->indices[0].as<AddNode>()) {
              // The case where the dst buffer is a byte buffer generated by merging dynamic
              // shared memory.
              // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)]
              auto* add = store->indices[0].as<AddNode>();
              if (!add->a->IsInstance<RampNode>()) return PrimExpr();
              if (!add->b->IsInstance<BroadcastNode>()) return PrimExpr();
              return tirx::Add(add->a.as<RampNode>()->base, add->b.as<BroadcastNode>()->value);
            }
            return PrimExpr();
          }();
          if (src_offset.defined() && dst_offset.defined()) {
            static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw");
            return Evaluate(Call(store->buffer->dtype, ptx_cp_async_op,
                                 {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
                                  load->buffer->data, src_offset, PrimExpr(bytes)})
                                .as_or_throw<PrimExpr>());
          }
        } else {
          // Only some vectorized indexing patterns are supported for now.
          auto src_offset = [=]() -> PrimExpr {
            if (load->indices[0]->IsInstance<RampNode>()) {
              return load->indices[0].as<RampNode>()->base;
            }
            return PrimExpr();
          }();

          auto dst_offset = [=]() -> PrimExpr {
            if (store->indices[0].as<RampNode>()) {
              return store->indices[0].as<RampNode>()->base;
            } else if (store->indices[0].as<AddNode>()) {
              // The case where the dst buffer is a byte buffer generated by merging dynamic
              // shared memory.
              // A_shared.dyn[(ramp(...), 1, 8) + x8(17408))] = A_global[ramp(...),1, 8)]
              auto* add = store->indices[0].as<AddNode>();
              if (!add->a->IsInstance<RampNode>()) return PrimExpr();
              if (!add->b->IsInstance<BroadcastNode>()) return PrimExpr();
              return tirx::Add(add->a.as<RampNode>()->base, add->b.as<BroadcastNode>()->value);
            }
            return PrimExpr();
          }();

          if (src_offset.defined() && dst_offset.defined()) {
            static const Op& ptx_cp_async_op = Op::Get("tirx.ptx.cp_async_raw");
            return Evaluate(Call(store->buffer->dtype, ptx_cp_async_op,
                                 {store->buffer->data, mul(dst_offset, PrimExpr(index_factor)),
                                  load->buffer->data, src_offset, PrimExpr(bytes), predicate_value})
                                .as_or_throw<PrimExpr>());
          }
        }
      }
    }
    return StmtMutator::VisitStmt_(store);
  }

  Stmt VisitStmt_(const BufferStoreNode* store) {
    if (in_async && (store->buffer.scope() == "shared" || store->buffer.scope() == "shared.dyn")) {
      if (auto* load = store->value.as<BufferLoadNode>()) {
        return InjectPTX(load, store);
      } else if (auto* call = store->value.as<CallNode>()) {
        // tirx.if_then_else is a call to tirx::builtin::if_then_else()
        if (call->op.same_as(builtin::if_then_else()) && call->args.size() == 3) {
          if (auto* load = call->args[1].as<BufferLoadNode>()) {
            // Only default value of 0 is supported since 0 is the default value used by cp.async
            // ptx. @see section 9.7.8.22.3. of
            // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#asynchronous-memory-operations
            bool else_value_is_zero = false;
            if (auto* b = call->args[2].as<BroadcastNode>()) {
              if (auto* f = b->value.as<FloatImmNode>()) {
                else_value_is_zero = f->value == 0.0f;
              }
            }
            if (auto* f = call->args[2].as<FloatImmNode>()) {
              else_value_is_zero = f->value == 0.0f;
            }
            if (else_value_is_zero) {
              return InjectPTX(load, store, true, call->args[0].as_or_throw<PrimExpr>());
            }
          }
        }
      }
    }
    return StmtMutator::VisitStmt_(store);
  }

 private:
  bool in_async{false};
};

namespace transform {

Pass InjectPTXAsyncCopy() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    auto* n = f.CopyOnWrite();
    n->body = PTXAsyncCopyInjector()(n->body);
    return f;
  };
  return CreatePrimFuncPass(pass_func, 0, "s_tir.InjectPTXAsyncCopy", {});
}

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

}  // namespace transform

}  // namespace s_tir
}  // namespace tvm
