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

/*!
 *  Lower TVM related builtin intrinsics such as packed call.
 * \file tirx/transform/lower_tvm_builtin.cc
 */
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/op.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <unordered_set>

#include "ir_utils.h"

namespace tvm {
namespace tirx {

namespace {

TVM_FFI_INLINE int GetVectorBytes(const PrimType& dtype) {
  return static_cast<int>(dtype.StorageBytes());
}

}  // namespace

// Calculate the statistics of packed function.
// These information are needed during codegen.
class BuiltinLower : public StmtExprMutator {
 public:
  static PrimFunc Build(PrimFunc func) {
    ffi::Optional<PrimExpr> device_type = std::nullopt;
    if (auto target = func->GetAttr<Target>(tvm::attr::kTarget)) {
      device_type = IntImm::Int32(target.value()->kind->default_device_type);
    }

    BuiltinLower mutator(device_type);
    func.CopyOnWrite()->body = mutator.VisitBodyAndRealizeAlloca(func->body);
    return func;
  }

  explicit BuiltinLower(ffi::Optional<PrimExpr> device_type = std::nullopt)
      : device_type_(device_type) {}

  // NOTE: Right now, we make the following scoping requirement
  // for memory allocated by the following primitives
  // - tvm_stack_make_array
  // - tvm_stack_make_shape
  // - arg stack
  //
  // Scoping and liveness rules:
  // - Every call_packed introduce a new scope.
  // - The memory allocated by tvm_stack_make_array/make_shape will
  //   no longer become valid outside the scope (and may be reused by
  //   subsequent call_packed.
  // - TODO(tvm-team): we might consider a root scope so stack_make_shape
  //   can be called out-side call_packed.
  //
  //  Example:
  //  {
  //    call_packed(make_shape1(...),
  //                call_packed(make_shape2(...))
  //    call_packed(make_shape3(...))
  //  }
  //
  //  In this case, make_shape1 and make_shape2 should not share memory,
  //  but they can share memory with make_shape3.
  //
  //  Rationale: most of the packed calls needs their own internal
  //  argument stack, and those stack can be shared across calls.
  //  Scoping is a quick way to enable sharing without having
  //  to do full-scale liveness analysis and it does its job.
  //  Alternative approaches can also be used.
  struct StackSizes {
    // If a tvm_stack_make_shape call has no arguments, it is still
    // valid and represents a scalar shape ().  Therefore, -1 is used
    // to represent "no shape arguments exist", while 0 represents
    // "shape arguments exist, all of which are size 0".
    int64_t shape_stack{-1};
    uint64_t array_stack{0};
    uint64_t arg_stack{0};
  };

  // Record stack frame for existing scope.
  struct AllocaScope {
    Buffer stack_shape;
    Var stack_array = Var("stack_array", PointerType::VoidPointerTy());
    Var stack_ffi_any = Var("stack_ffi_any", PointerType::VoidPointerTy());

    StackSizes max_sizes;
    StackSizes run_sizes;

    void UpdateMax() {
      max_sizes.shape_stack = std::max(max_sizes.shape_stack, run_sizes.shape_stack);
      max_sizes.array_stack = std::max(max_sizes.array_stack, run_sizes.array_stack);
      max_sizes.arg_stack = std::max(max_sizes.arg_stack, run_sizes.arg_stack);
    }

    void AssertMaxIsValid() const {
      TVM_FFI_ICHECK((max_sizes.shape_stack >= run_sizes.shape_stack) ||
                     (max_sizes.array_stack >= run_sizes.array_stack) ||
                     (max_sizes.arg_stack >= run_sizes.arg_stack));
    }
  };

  Stmt Build(Stmt stmt) { return this->VisitBodyAndRealizeAlloca(stmt); }

  StackSizes GetMaxStack(Stmt stmt) {
    BuiltinLower precheck;
    precheck.is_precheck_ = true;
    precheck.device_id_ = this->device_id_;
    precheck.device_type_ = this->device_type_;

    precheck.alloca_scope_.emplace_back();
    {
      // NOTE: this scope reference is invalid after any mutation is applied to alloca_scope_.
      auto& scope = precheck.alloca_scope_.back();
      scope.stack_shape = decl_buffer({IntImm::Int64(0)}, PrimType::Int(64), "stack_shape");
    }

    precheck.VisitStmt(stmt);

    TVM_FFI_ICHECK_EQ(precheck.alloca_scope_.size(), 1);
    return precheck.alloca_scope_[0].max_sizes;
  }

  // Allcoate stack frames, only at parallel-for or root.
  Stmt VisitBodyAndRealizeAlloca(Stmt stmt) {
    // Only perform the precheck up to the point where we would add a
    // new scope.
    if (is_precheck_) {
      return stmt;
    }

    alloca_scope_.emplace_back();
    {
      // NOTE: this scope reference is invalid after any mutation is applied to alloca_scope_.
      auto& scope = alloca_scope_.back();

      // Initial check to identify maximum stack sizes.  These are used
      // to construct Buffer objects to hold the stack, which are then
      // used when mutating.
      scope.max_sizes = GetMaxStack(stmt);

      // Build a flat list of Bind stmts followed by the body
      ffi::Array<Stmt> alloca_stmts;
      if (scope.max_sizes.arg_stack != 0) {
        alloca_stmts.push_back(
            Bind(scope.stack_ffi_any,
                 StackAlloca(scope.stack_ffi_any->ty, "tvm_ffi_any", scope.max_sizes.arg_stack)));
      }

      if (scope.max_sizes.array_stack != 0) {
        alloca_stmts.push_back(Bind(scope.stack_array, StackAlloca(scope.stack_array->ty, "array",
                                                                   scope.max_sizes.array_stack)));
      }

      if (scope.max_sizes.shape_stack != -1) {
        scope.stack_shape = decl_buffer({IntImm::Int64(scope.max_sizes.shape_stack)},
                                        PrimType::Int(64), "stack_shape");
        alloca_stmts.push_back(
            Bind(scope.stack_shape->data,
                 StackAlloca(scope.stack_shape->data->ty, "shape", scope.max_sizes.shape_stack)));
        stmt = SeqStmt::Flatten(DeclBuffer(scope.stack_shape), stmt);
      }

      if (!alloca_stmts.empty()) {
        alloca_stmts.push_back(stmt);
        stmt = SeqStmt(alloca_stmts);
      }
    }

    stmt = scope_.WithNewScope([&]() -> Stmt {
      Stmt visited = this->VisitStmt(stmt);
      return AppendPendingFrees(visited);
    });

    TVM_FFI_ICHECK(!alloca_scope_.empty());
    alloca_scope_.pop_back();

    return stmt;
  }

  Stmt VisitStmt(const Stmt& s) final {
    // allocate space to hold prepare stmts before s
    prep_seq_stack_.emplace_back(std::vector<Stmt>());

    auto scope_size = alloca_scope_.size();
    auto stmt = StmtExprMutator::VisitStmt(s);
    {
      // NOTE: this scope reference is invalid after any mutation is applied to alloca_scope_.
      auto& scope = alloca_scope_.back();
      // This invariant asserts the assumption that
      // make_stack_shape only happens within a call_packed.
      // We could relax this in the future if we want to
      // introduce root scope as a separate scope
      TVM_FFI_ICHECK_EQ(alloca_scope_.size(), scope_size)
          << "alloca_scope_ length is different before and after recursion";
      TVM_FFI_ICHECK_EQ(scope.run_sizes.shape_stack, -1)
          << "Expect no tvm_stack_make_shape outside of CallNodes";
      TVM_FFI_ICHECK_EQ(scope.run_sizes.array_stack, 0)
          << "Expect no tvm_stack_make_array outside of CallNodes";
    }

    auto prep_seq = std::move(prep_seq_stack_.back());
    prep_seq_stack_.pop_back();

    if (prep_seq.size() != 0) {
      Stmt ret = SeqStmt::Flatten(prep_seq, stmt);
      return ret;
    } else {
      return stmt;
    }
  }

  Stmt VisitStmt_(const BindNode* op) final {
    if (const CallNode* call = op->value.as<CallNode>()) {
      if (call->op.same_as(builtin::nd_mem_alloc_with_scope())) {
        return MakeNdMemAllocWithScope(op, call);
      }
    }
    return StmtExprMutator::VisitStmt_(op);
  }

  Stmt VisitStmt_(const AllocBufferNode* op) final {
    // Lower AllocBuffer to device allocate when needed.
    // AllocBuffer is flat (no body). Visit buffer fields via base class.
    Stmt stmt = StmtExprMutator::VisitStmt_(op);
    op = stmt.as<AllocBufferNode>();
    if (op->annotations.count(transform::kDisableLowerTVMBuiltin)) {
      if (op->annotations[transform::kDisableLowerTVMBuiltin].as_or_throw<IntImm>()->value) {
        return stmt;
      }
    }
    if (op->buffer->dtype.IsScalableVector()) {
      return stmt;
    }
    int64_t nbytes = GetVectorBytes(op->buffer->dtype);
    if (const auto* dev_type = device_type_.as<IntImmNode>();
        dev_type && dev_type->value == kDLCPU) {
      auto storage_scope = op->buffer->data->ty.as_or_throw<PointerType>()->storage_scope;
      if (storage_scope == "global") {
        auto constant_size = stmt.as_or_throw<AllocBuffer>().ConstantAllocationSize();
        if (constant_size.has_value() && constant_size.value() > 0 &&
            static_cast<size_t>(constant_size.value()) * nbytes < runtime::kMaxStackAlloca) {
          return stmt;
        }
      }
    }
    PrimExpr total_bytes = IntImm(PrimType::UInt(64), nbytes);
    for (size_t i = 0; i < op->buffer->shape.size(); ++i) {
      total_bytes = total_bytes * op->buffer->shape[i];
    }
    TVM_FFI_ICHECK(device_type_) << "Unknown device type in current IR";
    TVM_FFI_ICHECK(device_id_) << "Unknown device id in current IR";
    Stmt throw_last_error = Evaluate(
        Call(PrimType::Int(32), builtin::tvm_throw_last_error(), {}).as_or_throw<PrimExpr>());

    Stmt alloc_nullptr_check = IfThenElse(
        Call(PrimType::Bool(), builtin::isnullptr(), {op->buffer->data}).as_or_throw<PrimExpr>(),
        throw_last_error);

    static const Op& free_workspace_op = Op::Get("tirx.TVMBackendFreeWorkspace");
    static const Op& alloc_workspace_op = Op::Get("tirx.TVMBackendAllocWorkspace");
    PrimExpr free_op = Call(PrimType::Int(32), free_workspace_op,
                            {cast(PrimType::Int(32), device_type_.value()),
                             cast(PrimType::Int(32), device_id_.value()), op->buffer->data})
                           .as_or_throw<PrimExpr>();
    Stmt free_stmt = IfThenElse(free_op != IntImm::Int32(0), throw_last_error);

    // Push free to enclosing scope's pending_frees (LIFO ordering preserved).
    scope_.Current().pending_frees.push_back(free_stmt);

    Stmt alloc_bind = Bind(
        op->buffer->data,
        Call(op->buffer->data->ty, alloc_workspace_op,
             {cast(PrimType::Int(32), device_type_.value()),
              cast(PrimType::Int(32), device_id_.value()), total_bytes,
              IntImm::Int32(op->buffer->dtype.code()), IntImm::Int32(op->buffer->dtype.bits())}));

    return SeqStmt({alloc_bind, alloc_nullptr_check});
  }

  Stmt VisitStmt_(const AttrStmtNode* op) final {
    if (op->attr_key == attr::device_id) {
      auto cache = device_id_;
      device_id_ = op->value;
      Stmt out = scope_.WithNewScope([&]() -> Stmt {
        Stmt body = this->VisitStmt(op->body);
        return AppendPendingFrees(body);
      });
      device_id_ = cache;
      return out;
    } else if (op->attr_key == attr::device_type) {
      auto cache = device_type_;
      device_type_ = op->value;
      Stmt out = scope_.WithNewScope([&]() -> Stmt {
        Stmt body = this->VisitStmt(op->body);
        return AppendPendingFrees(body);
      });
      device_type_ = cache;
      return out;
    } else {
      return scope_.WithNewScope([&]() -> Stmt {
        Stmt visited = StmtExprMutator::VisitStmt_(op);
        if (!scope_.Current().pending_frees.empty()) {
          const auto* attr = visited.as<AttrStmtNode>();
          if (attr) {
            return AttrStmt(attr->node, attr->attr_key, attr->value, AppendPendingFrees(attr->body),
                            attr->span);
          }
        }
        return visited;
      });
    }
  }
  Stmt VisitStmt_(const ForNode* op) final {
    PrimExpr min = this->VisitPrimExpr(op->min);
    PrimExpr extent = this->VisitPrimExpr(op->extent);
    Stmt body;

    if (op->kind == ForKind::kParallel) {
      body = this->VisitBodyAndRealizeAlloca(op->body);
    } else {
      body = scope_.WithNewScope([&]() -> Stmt {
        Stmt visited = this->VisitStmt(op->body);
        return AppendPendingFrees(visited);
      });
    }

    if (min.same_as(op->min) && extent.same_as(op->extent) && body.same_as(op->body)) {
      return ffi::GetRef<Stmt>(op);
    } else {
      auto n = CopyOnWrite(op);
      n->min = std::move(min);
      n->extent = std::move(extent);
      n->body = std::move(body);
      return Stmt(n);
    }
  }

  Stmt VisitStmt_(const IfThenElseNode* op) final {
    PrimExpr condition = this->VisitPrimExpr(op->condition);
    // Each branch gets its own scope to prevent frees from leaking across branches.
    Stmt then_case = scope_.WithNewScope([&]() -> Stmt {
      Stmt visited = this->VisitStmt(op->then_case);
      return AppendPendingFrees(visited);
    });
    ffi::Optional<Stmt> else_case;
    if (op->else_case) {
      else_case = scope_.WithNewScope([&]() -> Stmt {
        Stmt visited = this->VisitStmt(op->else_case.value());
        return AppendPendingFrees(visited);
      });
    }
    if (condition.same_as(op->condition) && then_case.same_as(op->then_case) &&
        else_case.same_as(op->else_case)) {
      return ffi::GetRef<Stmt>(op);
    }
    return IfThenElse(condition, then_case, else_case, op->span);
  }

  Expr VisitExpr_(const CallNode* op) final {
    if (op->op.same_as(builtin::tvm_call_packed())) {
      return MakeCallPackedGeneric(op, 0, builtin::tvm_call_packed_lowered(),
                                   /* use_last_value_as_traced_value*/ false);
    } else if (op->op.same_as(builtin::tvm_call_cpacked())) {
      return MakeCallPackedGeneric(op, 0, builtin::tvm_call_cpacked_lowered(),
                                   /* use_last_value_as_traced_value*/ false);
    } else if (op->op.same_as(builtin::tvm_call_trace_packed())) {
      return MakeCallPackedGeneric(op, 0, builtin::tvm_call_trace_packed_lowered(),
                                   /* use_last_value_as_traced_value*/ true);
    } else if (op->op.same_as(builtin::anylist_setitem_call_packed())) {
      return MakeAnyListSetItemCallPacked(op, builtin::tvm_call_packed_lowered());
    } else if (op->op.same_as(builtin::anylist_setitem_call_cpacked())) {
      return MakeAnyListSetItemCallPacked(op, builtin::tvm_call_cpacked_lowered());
    } else if (op->op.same_as(builtin::tvm_stack_make_shape())) {
      return MakeShape(op);
    } else if (op->op.same_as(builtin::tvm_stack_make_array())) {
      return MakeArray(op);
    } else if (op->op.same_as(builtin::tvm_context_id())) {
      return IntImm(op->ty.as_or_throw<PrimType>(), 0);
    } else if (op->op.same_as(builtin::dma_copy())) {
      return MakeDMACopy(op);
    } else if (op->op.same_as(builtin::dma_wait())) {
      return MakeDMAWait(op);
    } else if (op->op.same_as(builtin::dma_start_group())) {
      return MakeDMAStartGroup(op);
    } else if (op->op.same_as(builtin::dma_end_group())) {
      return MakeDMAEndGroup(op);
    } else {
      return StmtExprMutator::VisitExpr_(op);
    }
  }

  StringImm GetDeviceMethodName(const char* method_name) const {
    TVM_FFI_ICHECK(device_type_) << "Method " << method_name << " requires the device type, "
                                 << "but occurred outside of a \"device_type\" annotation";

    auto as_int = device_type_.as<IntImmNode>();
    TVM_FFI_ICHECK(as_int) << "Method " << method_name
                           << " requires the device type to be a DLDeviceType enum value, "
                           << "but was instead the expression " << device_type_ << " with type "
                           << device_type_.value()->GetTypeKey();

    ffi::String device_name = runtime::DLDeviceType2Str(as_int->value);
    return StringImm("device_api." + device_name + "." + method_name);
  }

  PrimExpr MakeDMACopy(const CallNode* op) {
    PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
    Expr dst = op->args[1];
    Expr src = op->args[2];
    PrimExpr size = op->args[3].as_or_throw<PrimExpr>();
    PrimExpr bypass_cache = op->args[4].as_or_throw<PrimExpr>();

    auto method_name = GetDeviceMethodName("dma_copy");
    Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(),
                            {method_name, queue_id, dst, src, size, bypass_cache});
    return VisitPrimExpr(call_packed.as_or_throw<PrimExpr>());
  }

  PrimExpr MakeDMAWait(const CallNode* op) {
    PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();
    PrimExpr inflight = op->args[1].as_or_throw<PrimExpr>();

    auto method_name = GetDeviceMethodName("dma_wait");
    Call call_packed =
        Call(PrimType::Int(32), builtin::tvm_call_packed(), {method_name, queue_id, inflight});
    return VisitPrimExpr(call_packed.as_or_throw<PrimExpr>());
  }

  PrimExpr MakeDMAStartGroup(const CallNode* op) {
    PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();

    auto method_name = GetDeviceMethodName("dma_start_group");
    Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(), {method_name, queue_id});
    return VisitPrimExpr(call_packed.as_or_throw<PrimExpr>());
  }

  PrimExpr MakeDMAEndGroup(const CallNode* op) {
    PrimExpr queue_id = op->args[0].as_or_throw<PrimExpr>();

    auto method_name = GetDeviceMethodName("dma_end_group");
    Call call_packed = Call(PrimType::Int(32), builtin::tvm_call_packed(), {method_name, queue_id});
    return VisitPrimExpr(call_packed.as_or_throw<PrimExpr>());
  }

  // call shape
  Expr MakeShape(const CallNode* op) {
    // if args.size() == 0, it represents a scalar shape ()
    TVM_FFI_ICHECK(!alloca_scope_.empty());
    auto& scope = alloca_scope_.back();
    auto& prep_seq = prep_seq_stack_.back();
    if (scope.run_sizes.shape_stack == -1) {
      scope.run_sizes.shape_stack = 0;
    }
    int64_t stack_begin = scope.run_sizes.shape_stack;
    scope.run_sizes.shape_stack += op->args.size();
    Expr expr = StmtExprMutator::VisitExpr_(op);
    op = expr.as<CallNode>();
    // no need to perform any store for a scalar shape
    for (size_t i = 0; i < op->args.size(); ++i) {
      prep_seq.emplace_back(BufferStore(
          scope.stack_shape, cast(PrimType::Int(64), op->args[i].as_or_throw<PrimExpr>()),
          {ConstInt32(stack_begin + i)}));
    }
    return AddressOffset(scope.stack_shape->data, PrimType::Int(64), stack_begin);
  }
  // make array
  Expr MakeArray(const CallNode* op) {
    TVM_FFI_ICHECK(!alloca_scope_.empty());
    auto& scope = alloca_scope_.back();
    auto& prep_seq = prep_seq_stack_.back();

    size_t idx = scope.run_sizes.array_stack;
    scope.run_sizes.array_stack += 1;
    Expr expr = StmtExprMutator::VisitExpr_(op);
    op = expr.as<CallNode>();

    prep_seq.emplace_back(
        TVMStructSet(scope.stack_array, idx, builtin::kDLTensorData, op->args[0]));
    prep_seq.emplace_back(
        TVMStructSet(scope.stack_array, idx, builtin::kDLTensorShape, op->args[1]));
    Expr strides = op->args[2];
    if (auto prim_strides = strides.as<PrimExpr>(); prim_strides && is_zero(prim_strides.value())) {
      strides = ConstHandle(0);
    }
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorStrides, strides));
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorNDim,
                                       op->args[3].as_or_throw<PrimExpr>()));
    PrimType dtype = op->args[4].as_or_throw<PrimExpr>().ty();
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeCode,
                                       IntImm(PrimType::UInt(8), static_cast<int>(dtype.code()))));
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeBits,
                                       IntImm(PrimType::UInt(8), dtype.bits())));
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorTypeLanes,
                                       IntImm(PrimType::UInt(16), dtype.lanes())));
    // set byte offset
    int data_bytes = GetVectorBytes(dtype);
    PrimExpr elem_offset = op->args[5].as_or_throw<PrimExpr>();
    PrimExpr byte_offset;
    if (!is_zero(elem_offset)) {
      byte_offset = elem_offset * IntImm(elem_offset.ty(), data_bytes);
    } else {
      byte_offset = elem_offset;
    }
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorByteOffset,
                                       cast(PrimType::UInt(64), byte_offset)));
    TVM_FFI_ICHECK(device_type_) << "Unknown device type in current IR";
    TVM_FFI_ICHECK(device_id_) << "Unknown device id in current IR";
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorDeviceId,
                                       cast(PrimType::Int(32), device_id_.value())));
    prep_seq.emplace_back(TVMStructSet(scope.stack_array, idx, builtin::kDLTensorDeviceType,
                                       cast(PrimType::Int(32), device_type_.value())));
    return TVMStructGet(PointerType::VoidPointerTy(), scope.stack_array, idx,
                        builtin::kDLTensorAddr);
  }

  void SetPackedArg(Expr arg, const Var& args_stack, size_t stack_offset,
                    std::vector<tirx::Stmt>* prep_seq) {
    auto* call_pattern = arg.as<CallNode>();
    if (call_pattern && call_pattern->op.same_as(builtin::anylist_getitem())) {
      // call runtime function to set anylist
      static const Op& anylist_set_packed_arg_op = Op::Get("tirx.TVMBackendAnyListSetPackedArg");
      prep_seq->emplace_back(Evaluate(Call(
          PrimType::Int(32), anylist_set_packed_arg_op,
          {call_pattern->args[0], call_pattern->args[1], args_stack, ConstInt32(stack_offset)})));
    } else {
      int arg_type_index;
      if (arg.as<StringImmNode>()) {
        arg_type_index = ffi::TypeIndex::kTVMFFIRawStr;
        arg = reinterpret(PointerType::VoidPointerTy(), std::move(arg));
      } else if (arg->ty.as<PointerTypeNode>()) {
        arg_type_index = IsArrayHandle(arg) ? ffi::TypeIndex::kTVMFFIDLTensorPtr
                                            : ffi::TypeIndex::kTVMFFIOpaquePtr;
      } else {
        PrimExpr prim_arg = arg.as_or_throw<PrimExpr>();
        PrimType arg_ty = prim_arg.ty();
        PrimType api_ty = APIType(arg_ty);
        if (arg_ty != api_ty) {
          arg = Cast(api_ty, prim_arg);
        }
        if (api_ty.MatchesCode(DLDataTypeCode::kDLBool)) {
          arg_type_index = ffi::TypeIndex::kTVMFFIBool;
        } else if (api_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
          arg_type_index = ffi::TypeIndex::kTVMFFIInt;
        } else if (api_ty.code() == DLDataTypeCode::kDLFloat) {
          arg_type_index = ffi::TypeIndex::kTVMFFIFloat;
        } else {
          TVM_FFI_THROW(InternalError) << "Unsupported type: " << api_ty;
        }
      }

      // opaque handle need to set the kind properly
      if (arg_type_index == ffi::TypeIndex::kTVMFFIOpaquePtr) {
        prep_seq->emplace_back(
            IfThenElse(Call(PrimType::Bool(), builtin::isnullptr(), {arg}).as_or_throw<PrimExpr>(),
                       TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyTypeIndex,
                                    ConstInt32(ffi::TypeIndex::kTVMFFINone)),
                       TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyTypeIndex,
                                    ConstInt32(ffi::TypeIndex::kTVMFFIOpaquePtr))));
      } else {
        prep_seq->emplace_back(TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyTypeIndex,
                                            ConstInt32(arg_type_index)));
      }
      // set zero padding to ensure compatibility with FFI convention
      prep_seq->emplace_back(
          TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyZeroPadding, ConstInt32(0)));
      // handle arg value
      // NOTE: the intrinsic codegen will handle padding value clear for 32bit
      // types or types that are smaller than 64 bits.
      prep_seq->emplace_back(
          TVMStructSet(args_stack, stack_offset, builtin::kTVMFFIAnyUnionValue, arg));
    }
  }

  PrimExpr MakeAnyListSetItemCallPacked(const CallNode* op, const Op& lowered_op) {
    Expr list_handle = op->args[0];
    PrimExpr list_index = op->args[1].as_or_throw<PrimExpr>();

    Call call = MakeCallPackedGeneric(op, 2, lowered_op, false);
    Expr args_stack = call->args[1];
    // The stack offset of return value stack_end
    PrimExpr ret_offset = call->args[3].as_or_throw<PrimExpr>();
    auto& prep_seq = prep_seq_stack_.back();
    prep_seq.emplace_back(Evaluate(call.as_or_throw<PrimExpr>()));
    static const Op& anylist_move_from_packed_return_op =
        Op::Get("tirx.TVMBackendAnyListMoveFromPackedReturn");
    return Call(PrimType::Int(32), anylist_move_from_packed_return_op,
                {list_handle, list_index, args_stack, ret_offset})
        .as_or_throw<PrimExpr>();
  }
  /*!
   * \brief Generic tool to make low-level
   *  packed_call(other_args..., func_name, packed_arg0, packed_arg1...)
   *
   * \param op The call
   * \param name_offset The beginning of function name and call packed section.
   * \param lowered_packed_op The target lowered op.
   * \param pass_last_arg_as_traced_value Whether to pass last argument as traced value
   */
  Call MakeCallPackedGeneric(const CallNode* op, size_t name_offset, const Op& lowered_packed_op,
                             bool pass_last_arg_as_traced_value) {
    auto& scope = alloca_scope_.back();
    auto& prep_seq = prep_seq_stack_.back();

    int64_t restore_shape_stack = scope.run_sizes.shape_stack;
    size_t restore_array_stack = scope.run_sizes.array_stack;
    size_t arg_stack_begin = scope.run_sizes.arg_stack;

    size_t args_begin = name_offset + 1;
    size_t args_end = op->args.size();

    size_t num_args = args_end - args_begin;

    // The extra one slot is for return value.
    scope.run_sizes.arg_stack += num_args + 1;
    // Specially handle the buffer packed intrinsic
    Expr expr = StmtExprMutator::VisitExpr_(op);
    op = expr.as<CallNode>();

    for (size_t i = 0; i < num_args; ++i) {
      const Expr& arg = op->args[args_begin + i];
      TVM_FFI_CHECK(arg.as<PrimExpr>() || arg->ty.as<PointerTypeNode>(), TypeError)
          << "Packed call argument must have a primitive or pointer type, but got " << arg->ty;
      this->SetPackedArg(arg, scope.stack_ffi_any, arg_stack_begin + i, &prep_seq);
    }
    // explicitly set return value to None to avoid bad state interpretation
    prep_seq.emplace_back(TVMStructSet(scope.stack_ffi_any, num_args, builtin::kTVMFFIAnyTypeIndex,
                                       ConstInt32(ffi::TypeIndex::kTVMFFINone)));
    prep_seq.emplace_back(
        TVMStructSet(scope.stack_ffi_any, num_args, builtin::kTVMFFIAnyZeroPadding, ConstInt32(0)));
    prep_seq.emplace_back(TVMStructSet(scope.stack_ffi_any, num_args, builtin::kTVMFFIAnyUnionValue,
                                       IntImm::Int64(0)));
    // Verify stack size matches earlier value.
    if (is_precheck_) {
      scope.UpdateMax();
    } else {
      scope.AssertMaxIsValid();
    }
    scope.run_sizes.shape_stack = restore_shape_stack;
    scope.run_sizes.array_stack = restore_array_stack;
    scope.run_sizes.arg_stack = arg_stack_begin;
    ffi::Array<Expr> packed_args = {op->args[name_offset], scope.stack_ffi_any,
                                    ConstInt32(arg_stack_begin),
                                    ConstInt32(arg_stack_begin + num_args)};
    if (pass_last_arg_as_traced_value) {
      // pass in last element as traced value
      // used by call_packed_traced
      packed_args.push_back(op->args[op->args.size() - 1]);
    }
    return Call(op->ty, lowered_packed_op, packed_args);
  }

  Stmt MakeNdMemAllocWithScope(const BindNode* let, const CallNode* call) {
    TVM_FFI_ICHECK(device_type_) << "Unknown device type in current IR";
    TVM_FFI_ICHECK(device_id_) << "Unknown device id in current IR";
    Stmt throw_last_error = Evaluate(
        Call(PrimType::Int(32), builtin::tvm_throw_last_error(), {}).as_or_throw<PrimExpr>());

    const auto* dtype_node = let->var->ty.as<PointerTypeNode>()->element_type.as<PrimTypeNode>();
    TVM_FFI_ICHECK(dtype_node);
    PrimType dtype = ffi::GetRef<PrimType>(dtype_node);

    ffi::Array<Expr> args = {
        GetDeviceMethodName("alloc_nd"), device_type_.value(),        device_id_.value(),
        IntImm::Int32(dtype.code()),     IntImm::Int32(dtype.bits()),
    };

    for (size_t i = 0; i < call->args.size(); ++i) {
      args.push_back(call->args[i]);
    }

    Call call_packed = Call(let->var->ty, builtin::tvm_call_packed(), args);
    Stmt null_check =
        IfThenElse(Call(PrimType::Bool(), builtin::isnullptr(), ffi::Array<Expr>{let->var})
                       .as_or_throw<PrimExpr>(),
                   throw_last_error);

    // Construct free_nd call and register in current scope.
    // The free will be emitted on scope exit, matching the old LetStmt body semantics.
    PrimExpr storage_scope = call->args[0].as_or_throw<PrimExpr>();
    Call free_op = Call(PrimType::Int(32), builtin::tvm_call_packed(),
                        {GetDeviceMethodName("free_nd"), device_type_.value(), device_id_.value(),
                         storage_scope, let->var});
    Stmt free_stmt =
        IfThenElse(free_op.as_or_throw<PrimExpr>() != IntImm::Int32(0), throw_last_error);
    // Visit the free_stmt so tvm_call_packed builtins inside it get lowered.
    free_stmt = StmtExprMutator::VisitStmt(free_stmt);
    scope_.Current().pending_frees.push_back(free_stmt);

    // Re-visit so tvm_call_packed in the Bind value and null_check get lowered.
    return StmtExprMutator::VisitStmt(SeqStmt({Bind(let->var, call_packed), null_check}));
  }

 private:
  bool IsArrayHandle(const Expr& arg) {
    // specially set array handle.
    if (const CallNode* buf = arg.as<CallNode>()) {
      if (buf->op.same_as(builtin::tvm_struct_get()) &&
          buf->args[2].as<IntImmNode>()->value == builtin::kDLTensorAddr) {
        return true;
      }
    }
    return false;
  }

  /*!
   * \brief Scope level for tracking nd_mem_alloc_with_scope deallocations.
   *
   * When a Bind allocates via nd_mem_alloc_with_scope, the corresponding
   * free_nd stmt is pushed to the current scope's pending_frees. Body-carrying
   * stmts (For, IfThenElse, AttrStmt) create new scopes via
   * WithNewScope. On scope exit, pending_frees are appended after the body.
   * AllocBuffer (flat, no body) pushes its free to the enclosing scope.
   */
  struct ScopeLevel {
    std::vector<Stmt> pending_frees;
  };

  /*! \brief Emit any pending frees from the current scope after the given body stmt. */
  Stmt AppendPendingFrees(Stmt body) {
    auto& frees = scope_.Current().pending_frees;
    if (frees.empty()) return body;
    ffi::Array<Stmt> stmts;
    stmts.push_back(body);
    for (auto it = frees.rbegin(); it != frees.rend(); ++it) {
      stmts.push_back(*it);
    }
    frees.clear();
    return SeqStmt::Flatten(stmts);
  }

  // The prepration sequence to be emitted before the current statement.
  std::vector<std::vector<Stmt>> prep_seq_stack_;
  ffi::Optional<PrimExpr> device_type_{std::nullopt};
  ffi::Optional<PrimExpr> device_id_{std::nullopt};

  bool is_precheck_{false};

  // Record all stack frames.
  std::vector<AllocaScope> alloca_scope_;
  // Scope stack for nd_mem_alloc_with_scope free tracking.
  ScopeStack<ScopeLevel> scope_;
};

namespace transform {

Pass LowerTVMBuiltin() {
  auto pass_func = [](PrimFunc func, IRModule m, PassContext ctx) {
    if (IsHostFunc(func).value_or(false)) {
      func = BuiltinLower::Build(func);
      VLOG(2) << "LowerTVMBuiltin: " << func;
    }
    return func;
  };
  return CreatePrimFuncPass(pass_func, 0, "tirx.LowerTVMBuiltin", {});
}

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

}  // namespace transform
}  // namespace tirx
}  // namespace tvm
