/*
 * 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 src/relax/transform/call_tir_rewrite.cc
 * \brief Perform explicit tensor allocation for call_tir,
 *        call_tir_inplace, and call_dps_packed.
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/relax/attrs/op.h>
#include <tvm/relax/expr_functor.h>
#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
#include <tvm/tirx/op.h>

#include "utils.h"

namespace tvm {
namespace relax {

// ==================
// CallTIRMutator
// Perform explicit tensor allocation for call_tir, call_tir_inplace, or call_dps_packed.
// Example:
// lv0: Tensor(n, m) = rx.call_tir(func, (x), (n, m), dtype="float32")
// -->
// gv0 = rx.call("relax.builtin.alloc_tensor", [n, m], dtype="float32")
// rx.call_packed(func, x, gv0)

class CallTIRMutator : public ExprMutator {
 public:
  explicit CallTIRMutator(const IRModule& mod) : ExprMutator(mod), mod_(std::move(mod)) {}

  IRModule Run() {
    for (const auto& [gv, func] : mod_->functions) {
      if (func->IsInstance<FunctionNode>()) {
        auto updated_func = this->VisitExpr(func).as_or_throw<Function>();
        builder_->UpdateFunction(gv, updated_func);
      }
    }
    return builder_->GetContextIRModule();
  }

 private:
  using ExprMutator::VisitExpr_;
  Expr VisitExpr_(const CallNode* call) override {
    // post-order mutation
    Expr expr = VisitExprPostOrder_(call);
    call = expr.as<CallNode>();

    static const Op& call_tir_op = Op::Get("relax.call_tir");
    static const Op& call_tir_inplace_op = Op::Get("relax.call_tir_inplace");
    static const Op& call_dps_packed_op = Op::Get("relax.call_dps_packed");
    static const Op& alloc_tensor_op = Op::Get("relax.builtin.alloc_tensor");
    static const Op& call_tir_dyn_op = Op::Get("relax.vm.call_tir_dyn");

    if (call->op.same_as(call_tir_op) || call->op.same_as(call_tir_inplace_op) ||
        call->op.same_as(call_dps_packed_op)) {
      bool is_inplace = call->op.same_as(call_tir_inplace_op);
      const auto* inplace_attrs = call->attrs.as<CallTIRInplaceAttrs>();
      ffi::Array<Expr> outs;
      if (const auto& tensor_ty = MatchType<TensorType>(expr)) {
        // single output case
        const TensorType& output_ty = tensor_ty.value();
        TVM_FFI_ICHECK(output_ty->shape.has_value())
            << "the TensorType shape of call_tir has not populated";
        int dev_index = 0;
        ffi::String scope = "global";
        if (output_ty->vdevice.has_value()) {
          dev_index = GetDeviceIndex(mod_, output_ty->vdevice.value());
          scope = output_ty->vdevice.value()->memory_scope;
        } else {
          dev_index = GetDeviceIndexByScope(mod_, scope);
        }

        if (!is_inplace) {
          outs.push_back(builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
                                             {output_ty->shape.value().as_or_throw<ShapeExpr>(),
                                              DataTypeImm(output_ty->dtype.value()->dtype),
                                              IntImm::Int64(dev_index), StringImm(scope)},
                                             Attrs(), {output_ty}),
                                        "alloc"));
        } else {
          // if there is only one output, it must be an in-place argument, but check anyway
          TVM_FFI_ICHECK(inplace_attrs->inplace_indices[0] != -1)
              << "If calling call_tir_inplace and there is one output, its in-place index must not"
                 " be -1.";
          outs.push_back(
              call->args[1].as_or_throw<Tuple>()->fields[inplace_attrs->inplace_indices[0]]);
        }
      } else if (const auto& tuple_ty = MatchType<TupleType>(expr)) {
        // multiple output case
        const TupleType& output_ty = tuple_ty.value();
        for (size_t i = 0; i < output_ty->fields.size(); ++i) {
          const auto& field = output_ty->fields[i];

          TVM_FFI_ICHECK(field->IsInstance<TensorTypeNode>())
              << "call_tir expects Tuple of TensorType, but got " << field
              << " as an element of TupleType";
          const auto& field_tensor = field.as_or_throw<TensorType>();
          TVM_FFI_ICHECK(field_tensor->shape.has_value())
              << "call_tir expects all TensorType has shape, but got " << field_tensor
              << " as an element of TupleType";

          int dev_index = 0;
          ffi::String scope = "global";
          if (field_tensor->vdevice.has_value()) {
            dev_index = GetDeviceIndex(mod_, field_tensor->vdevice.value());
            scope = field_tensor->vdevice.value()->memory_scope;
          }

          if (!is_inplace || inplace_attrs->inplace_indices[i] == -1) {
            outs.push_back(
                builder_->Emit(Call(Type::Missing(), alloc_tensor_op,
                                    {field_tensor->shape.value().as_or_throw<ShapeExpr>(),
                                     DataTypeImm(field_tensor->dtype.value()->dtype),
                                     IntImm::Int64(dev_index), StringImm(scope)},
                                    Attrs(), {field_tensor}),
                               "alloc"));
          } else {
            outs.push_back(
                call->args[1].as_or_throw<Tuple>()->fields[inplace_attrs->inplace_indices[i]]);
          }
        }
      } else {
        TVM_FFI_THROW(TypeError) << "The type of call_tir expects to be TensorType or "
                                    "TupleType, but got"
                                 << expr->ty;
      }

      ffi::Array<Expr> args;
      if (call->args[1].as<TupleNode>()) {
        args = call->args[1].as_or_throw<Tuple>()->fields;
        // for call_tir_inplace, don't reinsert in-place args, only the newly allocated ones
        if (!is_inplace) {
          args.insert(args.end(), outs.begin(), outs.end());
        } else {
          for (size_t i = 0; i < outs.size(); i++) {
            if (inplace_attrs->inplace_indices[i] == -1) {
              args.push_back(outs[i]);
            }
          }
        }

        if (call->args.size() == 2) {
          builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
        } else {
          // unpack semantics
          args.push_back(call->args[2]);
          builder_->Emit(Call(Type::Missing(), call_tir_dyn_op, {call->args[0], Tuple(args)}), "_");
        }
      } else {
        if (!is_inplace) {
          args = outs;
          args.insert(args.begin(), call->args[1]);
        } else {
          args.push_back(call->args[1]);
        }
        builder_->Emit(Call(Type::Missing(), call->args[0], args), "_");
      }

      if (outs.size() == 1) {
        return outs[0];
      }
      return std::move(Tuple(outs));
    }

    return ffi::GetRef<Expr>(call);
  }

  /*! \brief The context IRModule. */
  IRModule mod_;
};

namespace transform {

Pass CallTIRRewrite() {
  auto pass_func = [=](IRModule mod, PassContext pc) { return CallTIRMutator(mod).Run(); };
  return CreateModulePass(/*pass_function=*/pass_func,
                          /*opt_level=*/0,
                          /*pass_name=*/"CallTIRRewrite",
                          /*required=*/{});
}

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

}  // namespace transform

}  // namespace relax
}  // namespace tvm
