/*
 * 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 storage_rewrite.cc
 * \brief Memory access pattern analysis and optimization.
 *  Re-write data access to enable memory sharing when possible.
 */
#include <tvm/arith/analyzer.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/type.h>
#include <tvm/runtime/logging.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <list>
#include <map>
#include <unordered_map>
#include <unordered_set>

#include "../../arith/int_operator.h"
#include "../../runtime/thread_storage_scope.h"
#include "../ir/buffer_common.h"
#include "ir_utils.h"

namespace tvm {
namespace tirx {

using runtime::StorageRank;
using runtime::StorageScope;

namespace {

ffi::Optional<Var> GetBufferDataVar(const ffi::Any& data) {
  if (auto var = data.as<Var>()) {
    return var;
  }
  if (const auto* call = data.as<CallNode>();
      call && call->op.same_as(builtin::buffer_data()) && call->args.size() == 1) {
    return call->args[0].as<Var>();
  }
  return std::nullopt;
}

struct PrimTypeHash {
  size_t operator()(const PrimType& ty) const {
    DLDataType dtype = ty->dtype;
    return (static_cast<size_t>(dtype.code) << 24) ^ (static_cast<size_t>(dtype.bits) << 16) ^
           static_cast<size_t>(dtype.lanes);
  }
};

struct PrimTypeEqual {
  bool operator()(const PrimType& lhs, const PrimType& rhs) const { return lhs == rhs; }
};

}  // namespace

// Find a linear pattern of storage access
// Used for liveness analysis.
// Composite scopes(loop/thread_launch/IfThen) is represented by two points:
// before_scope -> scope_body -> after_scope
//
// The linear_seq_ stores before_scope and after_scope.
// The access to the arrays are stored at the after_scope point.
//
// Define "scope" as the body of For/thread_launch/IfThenElse
// This pass tries to detect last point that we need to keep memory
// alive under the same scope as allocate.
// The storage need to be kept alive between allocate and last access.
// The free point is only inserted at the same scope of allocate.
//
class LinearAccessPatternFinder final : public StmtExprVisitor {
 public:
  explicit LinearAccessPatternFinder(const ffi::Array<Var>& params) {
    for (const Var& param : params) {
      if (auto buffer = param.as<BufferVar>()) {
        buffer_aliases_.Set(buffer.value().var(), buffer.value().var());
      }
    }
  }

  /*! \brief record the touch hist of statment. */
  struct StmtEntry {
    // The statment
    const ffi::Object* stmt;
    // The index in the linear_seq_ to point to end of the nested scope.
    // This is only set to non-zero if stmt is a nested scope.
    // if offset > 0, means this is the begin, the end entry is current_index + offset
    // if offset < 0, means this is the end, the begin entry is current_index + offset
    int64_t scope_pair_offset{0};
    // The buffer variables this statment touched.
    std::vector<const VarNode*> touched;
  };
  // The scope of each allocation
  struct AllocEntry {
    // The physical dimension of the allocation.
    size_t num_physical_dimensions{0};
    // scope level
    size_t level{0};
    // The AllocBuffer node that created this allocation.
    const AllocBufferNode* alloc{nullptr};
  };

  void VisitStmt_(const AllocBufferNode* op) final {
    size_t level = scope_.size();
    const VarNode* buf = op->buffer.get();
    buffer_aliases_.Set(op->buffer.var(), op->buffer.var());

    AllocEntry entry;
    entry.alloc = op;
    entry.level = level;
    entry.num_physical_dimensions = op->buffer->shape.size();
    alloc_info_[buf] = entry;

    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitStmt_(const DeclBufferNode* op) final { RegisterBufferAlias(op->buffer, op->data); }

  void VisitStmt_(const BufferStoreNode* op) final {
    scope_.push_back(StmtEntry());
    // visit subexpr
    StmtExprVisitor::VisitStmt_(op);
    // Add write access.
    const VarNode* buffer_var =
        buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get();
    auto it = alloc_info_.find(buffer_var);
    if (it != alloc_info_.end() && it->second.alloc) {
      TVM_FFI_ICHECK_LT(it->second.level, scope_.size());
      scope_[it->second.level].touched.push_back(buffer_var);

      TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions)
          << "BufferVar " << op->buffer.name() << " is allocated with "
          << it->second.num_physical_dimensions
          << " physical dimensions, but is accessed as having "
          << "1 physical dimension" << std::endl;
    }
    StmtEntry e = scope_.back();
    scope_.pop_back();
    if (e.touched.size() != 0) {
      e.stmt = op;
      linear_seq_.push_back(e);
    }
  }

  void VisitExpr_(const BufferLoadNode* op) final {
    // Add write access.
    StmtExprVisitor::VisitExpr_(op);

    const VarNode* buffer_var =
        buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get();
    auto it = alloc_info_.find(buffer_var);
    if (it != alloc_info_.end() && it->second.alloc) {
      TVM_FFI_ICHECK_LT(it->second.level, scope_.size())
          << "Load memory in places other than store.";
      scope_[it->second.level].touched.push_back(buffer_var);

      TVM_FFI_ICHECK_EQ(1, it->second.num_physical_dimensions)
          << "BufferVar " << op->buffer.name() << " is allocated with "
          << it->second.num_physical_dimensions
          << " physical dimensions, but is accessed as having "
          << "1 physical dimension" << std::endl;
    }
  }

  void VisitStmt_(const EvaluateNode* op) final {
    scope_.push_back(StmtEntry());
    // visit subexpr
    StmtExprVisitor::VisitStmt_(op);
    StmtEntry e = scope_.back();
    scope_.pop_back();
    if (e.touched.size() != 0) {
      e.stmt = op;
      linear_seq_.push_back(e);
    }
  }

  void VisitStmt_(const ReturnNode* op) final {
    scope_.push_back(StmtEntry());
    StmtExprVisitor::VisitStmt_(op);
    StmtEntry e = scope_.back();
    scope_.pop_back();
    if (e.touched.size() != 0) {
      e.stmt = op;
      linear_seq_.push_back(e);
    }
  }

  void VisitExpr_(const VarNode* buf) final {
    // Directly reference to the variable count as a read.
    if (buf->ty.as<BufferTypeNode>()) {
      Var var = ffi::GetRef<Var>(buf);
      buf = buffer_aliases_.Get(var).value_or(var).get();
    }
    auto it = alloc_info_.find(buf);
    if (it != alloc_info_.end() && it->second.alloc) {
      TVM_FFI_ICHECK_LT(it->second.level, scope_.size()) << " buf=" << buf->name;
      scope_[it->second.level].touched.push_back(buf);
    }
  }

  template <typename T>
  void VisitNewScope(const T* op) {
    scope_.push_back(StmtEntry());
    StmtEntry e;
    e.stmt = op;
    int64_t begin_index = static_cast<int64_t>(linear_seq_.size());
    // before scope.
    linear_seq_.push_back(e);
    StmtExprVisitor::VisitStmt_(op);
    // after scope.
    e.touched = std::move(scope_.back().touched);
    scope_.pop_back();
    int64_t end_index = static_cast<int64_t>(linear_seq_.size());
    TVM_FFI_ICHECK_GT(end_index, begin_index);
    e.scope_pair_offset = begin_index - end_index;
    linear_seq_.push_back(e);
    // record the pointer to end index.
    TVM_FFI_ICHECK_NE(end_index, 0U);
    linear_seq_[begin_index].scope_pair_offset = end_index - begin_index;
  }

  void VisitStmt_(const AttrStmtNode* op) final {
    // Only record the outer most thread extent.
    if (op->attr_key == attr::thread_extent && !in_thread_env_) {
      in_thread_env_ = true;
      VisitNewScope(op);
      in_thread_env_ = false;
    } else if (op->attr_key == attr::extern_scope) {
      VisitNewScope(op);
    } else if (op->attr_key == s_tir::attr::virtual_thread) {
      VisitNewScope(op);
    } else {
      StmtExprVisitor::VisitStmt_(op);
    }
  }

  void VisitStmt_(const IfThenElseNode* op) final { VisitNewScope(op); }

  void VisitStmt_(const ForNode* op) final { VisitNewScope(op); }

  void VisitStmt_(const WhileNode* op) final { VisitNewScope(op); }

  void VisitStmt_(const AssertStmtNode* op) final { VisitNewScope(op); }

  void VisitStmt_(const BindNode* op) final {
    scope_.push_back(StmtEntry());
    // visit subexpr (the value may contain BufferLoad)
    StmtExprVisitor::VisitStmt_(op);
    StmtEntry e = scope_.back();
    scope_.pop_back();
    if (e.touched.size() != 0) {
      e.stmt = op;
      linear_seq_.push_back(e);
    }
  }

  // linearized access sequence.
  std::vector<StmtEntry> linear_seq_;
  // The storage scope of each buffer
  std::unordered_map<const VarNode*, AllocEntry> alloc_info_;
  // Physical roots of buffer aliases, flattened when each DeclBuffer is visited.
  ffi::Map<Var, Var> buffer_aliases_;

 private:
  void RegisterBufferAlias(BufferVar buffer, const Expr& data) {
    Var root = buffer.var();
    if (auto source = GetBufferDataVar(data);
        source.has_value() && source.value()->ty.as<BufferTypeNode>()) {
      auto source_root = buffer_aliases_.Get(source.value());
      TVM_FFI_ICHECK(source_root.has_value()) << "Buffer alias source " << source.value()->name
                                              << " must be registered before its DeclBuffer alias";
      root = source_root.value();
    }
    buffer_aliases_.Set(buffer.var(), root);
  }

  // Whether already in thread env.
  bool in_thread_env_{false};
  // The scope stack.
  std::vector<StmtEntry> scope_;
};

// Verify if the statement can be run safely via inplace fashion
//
// Detect pattern: dst[index] = f(src[index])
//
// WARNING: the current detection algorithm cannot handle the case
// when a location in an array is written multiple times
//
// For example, the following program will pass the check,
// but we cannot make A and B to be the same array.
//
//  A[0] = B[0] + 1
//  A[0] = B[0] + 1
//
// The high level code generator needs to ensure that the generated
// code only write each location of the target array once.
//
// This is the case with IR generated by the current compute schedule.
// We explicitly return false if we find there is an extern block
// which can be arbitrary IR.
//
// Neve-the-less, inplace detector should be used with care in mind.
// We may also consider introduce a condition checker that checks
// if every index only visited once for an absolute sufficient condition.
//
// The code after inplace transformation is no longer idempotent.
//
class InplaceOpVerifier : public StmtExprVisitor {
 public:
  bool Check(const ffi::Object* stmt, const VarNode* dst, const VarNode* src) {
    dst_ = dst;
    src_ = src;
    result_ = true;
    if (stmt->IsInstance<AttrStmtNode>()) {
      VisitStmt_(static_cast<const AttrStmtNode*>(stmt));
    } else if (stmt->IsInstance<ForNode>()) {
      VisitStmt_(static_cast<const ForNode*>(stmt));
    } else if (stmt->IsInstance<IfThenElseNode>()) {
      VisitStmt_(static_cast<const IfThenElseNode*>(stmt));
    } else if (stmt->IsInstance<WhileNode>()) {
      VisitStmt_(static_cast<const WhileNode*>(stmt));
    } else if (stmt->IsInstance<BufferStoreNode>()) {
      VisitStmt_(static_cast<const BufferStoreNode*>(stmt));
    } else {
      return false;
    }
    return result_;
  }

  using StmtExprVisitor::VisitStmt_;

  void VisitStmt(const Stmt& n) final {
    if (!result_) return;
    StmtExprVisitor::VisitStmt(n);
  }
  void VisitExpr(const Expr& n) final {
    if (!result_) return;
    StmtExprVisitor::VisitExpr(n);
  }

  void VisitExpr_(const VarNode* op) final {
    // assume all opaque access is unsafe
    if (op == dst_ || op == src_) {
      result_ = false;
      return;
    }
  }

  void VisitStmt_(const BufferStoreNode* op) final {
    ++mem_nest_;
    for (const auto& index : op->indices) {
      this->VisitExpr(index);
    }
    --mem_nest_;
    if (op->buffer.get() == dst_) {
      store_ = op;
      this->VisitExpr(op->value);
      store_ = nullptr;
    } else {
      this->VisitExpr(op->value);
    }
  }

  void VisitStmt_(const AttrStmtNode* op) final {
    // always reject extern code
    if (op->attr_key == attr::extern_scope) {
      result_ = false;
      return;
    }
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitStmt_(const AllocBufferNode* op) final {
    // reject inplace for volatile buffers
    if (op->annotations.count(attr::kVolatile)) {
      result_ = false;
      return;
    }
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitExpr_(const BufferLoadNode* op) final {
    const VarNode* buf = op->buffer.get();
    // cannot read from dst_ (no reduction)
    if (buf == dst_) {
      result_ = false;
      return;
    }
    // do not allow indirect memory load
    if (mem_nest_ != 0) {
      result_ = false;
      return;
    }
    if (src_ == buf) {
      if (store_ == nullptr || store_->value.ty() != op->ty.as_or_throw<PrimType>()) {
        result_ = false;
        return;
      }
      TVM_FFI_ICHECK_EQ(store_->indices.size(), op->indices.size())
          << "Store/Load occur to the same buffer " << buf->name
          << " with differing number of indices";
      for (size_t i = 0; i < store_->indices.size(); i++) {
        if (!tirx::ExprDeepEqual()(store_->indices[i], op->indices[i])) {
          result_ = false;
          return;
        }
      }
    }
    ++mem_nest_;
    StmtExprVisitor::VisitExpr_(op);
    --mem_nest_;
  }

 private:
  // result of the check
  bool result_{true};
  // destination memory
  const VarNode* dst_;
  // source variable
  const VarNode* src_;
  // counter of load,
  // it is not safe to inplace when there is nested load like A[B[i]]
  int mem_nest_{0};
  // The current store to be inspected
  const BufferStoreNode* store_{nullptr};
};

/* \brief Rewrite and merge memory allocation.
 *
 * Using LinearAccessPatternFinder, determines which buffers could share an
 * allocation.  This includes both sequential usage of the same buffer and
 * merging small allocations at the same scope into a single larger allocation.
 * The merging of small allocations requires the codegen to cast the resulting
 * value from the storage type to the output type after access.
 */
class StoragePlanRewriter : public StmtExprMutator {
 public:
  using StmtEntry = LinearAccessPatternFinder::StmtEntry;
  using AllocEntry = LinearAccessPatternFinder::AllocEntry;

  Stmt Rewrite(Stmt stmt, const ffi::Array<Var>& params, bool detect_inplace, bool enable_reuse,
               bool reuse_require_exact_matched_dtype) {
    detect_inplace_ = detect_inplace;
    // plan the rewrite
    LinearAccessPatternFinder finder(params);
    finder(stmt);
    this->LivenessAnalysis(finder.linear_seq_);
    this->PlanMemory(finder.linear_seq_, finder.alloc_info_, enable_reuse,
                     reuse_require_exact_matched_dtype);
    buffer_aliases_ = std::move(finder.buffer_aliases_);
    this->PrepareNewAlloc();
    // start rewrite
    stmt = operator()(std::move(stmt));
    if (attach_map_.count(nullptr)) {
      return MakeAttach(attach_map_.at(nullptr), stmt);
    }
    return stmt;
  }

  template <typename Node>
  Node VisitBufferAccess(Node node) {
    const VarNode* root =
        buffer_aliases_.Get(node->buffer.var()).value_or(node->buffer.var()).get();
    auto it = alloc_map_.find(root);
    if (it != alloc_map_.end()) {
      BufferVar buf = RemapBuffer(node->buffer, it->second->alloc_var);

      ffi::Array<PrimExpr> indices = node->indices;
      indices.Set(indices.size() - 1,
                  RemapIndex(node->buffer->dtype, indices[indices.size() - 1], it->second));

      auto writer = node.CopyOnWrite();
      writer->buffer = buf;
      writer->indices = indices;
    }
    return node;
  }

  BufferVar RemapBuffer(BufferVar buf, Var new_backing_array) {
    auto key = buf.get();
    auto it = buffer_remap_.find(key);
    if (it != buffer_remap_.end()) {
      TVM_FFI_ICHECK_EQ(remapped_backing_.at(it->second.get()).get(), new_backing_array.get())
          << "Cannot remap buffer " << buf.name() << " to use backing array "
          << new_backing_array->name << ", previously used backing array "
          << remapped_backing_.at(it->second.get()).name();
      return it->second;
    }

    BufferVar backing(new_backing_array);
    BufferVar remapped = buf.same_as(backing)
                             ? buf
                             : RebuildBufferVar(buf, CopyBufferType(buf), new_backing_array->name);
    buffer_remap_[key] = remapped;
    remapped_backing_[remapped.get()] = backing;
    return remapped;
  }

  Stmt VisitStmt_(const BufferStoreNode* op) final {
    auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<BufferStore>();
    return VisitBufferAccess(std::move(node));
  }

  Expr VisitExpr_(const BufferLoadNode* op) final {
    auto node = StmtExprMutator::VisitExpr_(op).as_or_throw<BufferLoad>();
    return VisitBufferAccess(std::move(node));
  }

  Expr VisitExpr_(const VarNode* op) final {
    const VarNode* root = op;
    if (op->ty.as<BufferTypeNode>()) {
      Var var = ffi::GetRef<Var>(op);
      root = buffer_aliases_.Get(var).value_or(var).get();
    }
    auto it = alloc_map_.find(root);
    if (it != alloc_map_.end()) {
      if (it->second->bits_offset != 0) {
        LOG(WARNING) << "Use a merged buffer variable address, could cause error";
      }
      return it->second->alloc_var;
    } else {
      return ffi::GetRef<Var>(op);
    }
  }
  Expr VisitExpr_(const CallNode* op) final {
    if (op->op.same_as(builtin::tvm_access_ptr())) {
      TVM_FFI_ICHECK_EQ(op->args.size(), 5U);
      PrimExpr dtype_marker = op->args[0].as_or_throw<PrimExpr>();
      PrimType dtype = dtype_marker.ty();
      auto buffer_var = GetBufferDataVar(op->args[1]);
      if (!buffer_var.has_value()) {
        return StmtExprMutator::VisitExpr_(op);
      }
      const VarNode* buffer = buffer_var.value().get();
      if (buffer->ty.as<BufferTypeNode>()) {
        Var var = buffer_var.value();
        buffer = buffer_aliases_.Get(var).value_or(var).get();
      }
      auto it = alloc_map_.find(buffer);
      if (it == alloc_map_.end()) {
        return StmtExprMutator::VisitExpr_(op);
      }
      const StorageEntry* se = it->second;
      PrimExpr offset = this->VisitPrimExpr(op->args[2].as_or_throw<PrimExpr>());
      PrimExpr extent = this->VisitPrimExpr(op->args[3].as_or_throw<PrimExpr>());
      uint64_t elem_bits = dtype.bits() * dtype.lanes();
      TVM_FFI_ICHECK_EQ(se->bits_offset % elem_bits, 0U);
      if (se->bits_offset != 0) {
        offset = MakeConst(offset.ty(), se->bits_offset / elem_bits) + offset;
      }
      return Call(
          op->ty, op->op,
          {dtype_marker, se->alloc_var, offset, extent, op->args[4].as_or_throw<PrimExpr>()},
          op->attrs, {}, op->span);
    } else {
      return StmtExprMutator::VisitExpr_(op);
    }
  }

  Stmt VisitStmt_(const AttrStmtNode* op) final {
    if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread ||
        attr::IsPragmaKey(op->attr_key)) {
      // remake all the allocation at the attach scope.
      if (attach_map_.count(op)) {
        auto& svec = attach_map_[op];
        Stmt stmt = StmtExprMutator::VisitStmt_(op);
        op = stmt.as<AttrStmtNode>();
        return AttrStmt(op->node, op->attr_key, op->value, MakeAttach(svec, op->body));
      } else {
        return StmtExprMutator::VisitStmt_(op);
      }
    } else {
      return StmtExprMutator::VisitStmt_(op);
    }
  }

  Stmt VisitStmt_(const ForNode* op) final {
    TVM_FFI_ICHECK(op->kind != ForKind::kVectorized) << "VectorizeLoop before LiftStorageAlloc";
    // remake all the allocation at the attach scope.
    if (attach_map_.count(op)) {
      auto& svec = attach_map_[op];
      Stmt stmt = StmtExprMutator::VisitStmt_(op);
      op = stmt.as<ForNode>();
      return For(op->loop_var, op->min, op->extent, op->kind, MakeAttach(svec, op->body),
                 op->thread_binding, op->annotations, op->step);
    } else {
      return StmtExprMutator::VisitStmt_(op);
    }
  }

  Stmt VisitStmt_(const AllocBufferNode* op) final {
    // AllocBuffer combines allocation and buffer declaration.
    // Storage rewrite may merge this allocation with others.
    if (auto it = alloc_map_.find(op->buffer.get()); it != alloc_map_.end()) {
      if (it->second->alloc_var.get() == op->buffer.get()) {
        // This is the "winner" allocation -- its AllocBuffer was already
        // hoisted by PrepareNewAlloc.  Strip this one.
        return Evaluate(0);
      }
      // This allocation was merged into another.  Emit a DeclBuffer
      // aliasing the winner's data variable.
      BufferVar buf = RemapBuffer(op->buffer, it->second->alloc_var);
      return DeclBuffer(buf, BufferVar(it->second->alloc_var).data());
    }
    // If not in alloc_map (e.g. unused), strip entirely.
    return Evaluate(0);
  }

  Stmt VisitStmt_(const DeclBufferNode* op) final {
    const VarNode* root = buffer_aliases_.Get(op->buffer.var()).value_or(op->buffer.var()).get();
    auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<DeclBuffer>();
    auto it = alloc_map_.find(root);
    if (it != alloc_map_.end()) {
      auto writer = node.CopyOnWrite();
      writer->buffer = RemapBuffer(node->buffer, it->second->alloc_var);
      writer->data = BufferVar(it->second->alloc_var).data();
    }
    return node;
  }

 private:
  struct StorageEntry {
    // The scope that this alloc attaches after
    // For shared/local memory it is beginning of the thread extent.
    // for global memory it is nullptr, means beginning of everything.
    const ffi::Object* attach_scope_{nullptr};
    // The constant size of the buffer in bits, only used if it is constant
    uint64_t const_nbits{0};
    // The storage scope.
    StorageScope scope;
    // The physical dimensionality of the allocations.  Since
    // StorageRewrite is applied after FlattenBuffer,
    // this is the number of extent dimensions.
    size_t ndim;
    // Allocs that shares this entry.
    std::vector<const AllocBufferNode*> allocs;
    // The children of this entry, not including itself.
    std::vector<StorageEntry*> merged_children;
    // The replacement AllocBuffer, if any.
    std::vector<Stmt> alloc_nest;
    // The var expr of new allocation.
    Var alloc_var;
    // The allocation element type.
    PrimType elem_type = PrimType::Void();
    // Whether any constituent allocation was marked volatile.
    bool is_volatile{false};
    // This is non-zero if this alloc_buffer is folded into another one
    // the address(in bits) becomes alloc_var + bits_offset;
    // can be effectively converted to the element type.
    // We need to convert bit_offset to offset of specific element type later.
    //
    // We use bits(instead of bytes) to support non-conventional indexing in hardware.
    // When we are merging buffer together, the bits_offset are set to be aligned
    // to certain value given by the max_simd_bits property of the special memory.
    //
    // This allows effective sharing among different types as long as their alignment
    // requirement fits into the max_simd_bits.
    uint64_t bits_offset{0};
  };

  // Checks whether the storage_scope is especially tagged for a specific memory.
  // Special memory is all combined into a single allocation.
  bool IsSpecialTaggedMemory(const StorageScope& scope) {
    return scope.tag.length() != 0 && scope.tag != ".dyn" && scope.tag != ".workspace" &&
           scope.tag != ".vtcm";
  }

  // Alllocate entry of node.
  // Event entry in liveness analysis
  struct EventEntry {
    // variables we generate
    std::vector<const VarNode*> gen;
    // variables we kill
    std::vector<const VarNode*> kill;
  };

  Stmt MakeAttach(const std::vector<StorageEntry*>& svec, Stmt body) {
    for (auto it = svec.rbegin(); it != svec.rend(); it++) {
      body = MergeNest((*it)->alloc_nest, body);
    }
    return body;
  }
  // Remap the index
  PrimExpr RemapIndex(PrimType dtype, PrimExpr index, StorageEntry* e) {
    if (e->bits_offset == 0) return index;
    uint64_t elem_bits = dtype.bits();
    TVM_FFI_ICHECK_EQ(e->bits_offset % elem_bits, 0U);
    return MakeConst(index.ty(), e->bits_offset / elem_bits) + index;
  }
  // Prepare the new allocations
  void PrepareNewAlloc() {
    for (size_t i = 0; i < alloc_vec_.size(); ++i) {
      StorageEntry* e = alloc_vec_[i].get();
      attach_map_[e->attach_scope_].push_back(e);
    }
    // find allocation via attach map.
    for (auto& kv : attach_map_) {
      // find the element with the most amount of bytes.
      std::vector<StorageEntry*>& vec = kv.second;
      // try to find merge, for tagged memory
      for (size_t i = 0; i < vec.size(); ++i) {
        StorageEntry* e = vec[i];
        if (IsSpecialTaggedMemory(e->scope)) {
          TVM_FFI_ICHECK_NE(e->const_nbits, 0U) << "Special tagged memory must be const size";
          for (size_t j = 0; j < i; ++j) {
            if (e->scope == vec[j]->scope) {
              vec[j]->merged_children.push_back(e);
              break;
            }
          }
        }
      }
      // Start allocation
      for (size_t i = 0; i < vec.size(); ++i) {
        StorageEntry* e = vec[i];
        // already merged
        if (e->bits_offset != 0) continue;
        if (e->merged_children.size() != 0) {
          NewAllocTagMerged(e);
          continue;
        }
        if (e->allocs.size() == 1 && e->allocs[0]->buffer->dtype.IsScalableVector()) {
          // Scalable vector lanes are runtime-dependent.  Keep these allocations exact rather
          // than trying to compare or merge their compile-time bit size.
          e->alloc_var = e->allocs[0]->buffer.var();
          BufferVar buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var);
          ffi::Map<ffi::String, ffi::Any> annotations;
          if (e->is_volatile) {
            annotations.Set(attr::kVolatile, true);
          }
          e->alloc_nest.push_back(AllocBuffer(buf, annotations));
          continue;
        }
        // Get the allocation size;
        e->alloc_var = e->allocs[0]->buffer.var();
        PrimType alloc_type = e->allocs[0]->buffer->dtype;
        for (const AllocBufferNode* op : e->allocs) {
          if (op->buffer->dtype.lanes() > alloc_type.lanes()) {
            alloc_type = op->buffer->dtype;
          }
        }

        bool all_allocs_identical = std::all_of(
            e->allocs.begin() + 1, e->allocs.end(), [&](const AllocBufferNode* op) -> bool {
              const AllocBufferNode* first = *e->allocs.begin();
              if (op->buffer->dtype->dtype != first->buffer->dtype->dtype) {
                return false;
              }
              if (op->buffer->shape.size() != first->buffer->shape.size()) {
                return false;
              }
              ExprDeepEqual expr_equal;
              for (size_t i = 0; i < op->buffer->shape.size(); i++) {
                if (!expr_equal(op->buffer->shape[i], first->buffer->shape[i])) {
                  return false;
                }
              }
              return true;
            });

        if (all_allocs_identical) {
          // Emit AllocBuffer for the hoisted allocation.
          BufferVar buf = RemapBuffer(e->allocs[0]->buffer, e->alloc_var);
          ffi::Map<ffi::String, ffi::Any> annotations;
          if (e->is_volatile) {
            annotations.Set(attr::kVolatile, true);
          }
          e->alloc_nest.push_back(AllocBuffer(buf, annotations));
        } else {
          // Build a merged allocation
          PrimExpr combo_size;
          for (const AllocBufferNode* op : e->allocs) {
            TVM_FFI_ICHECK_EQ(op->buffer->shape.size(), 1)
                << "BufferVar var " << op->buffer.name()
                << " was identified as a re-usable allocation, but has " << op->buffer->shape.size()
                << " physical dimensions.  "
                << "Currently, only flat 1-d memory spaces should be identified as re-usable "
                   "allocations.";
            PrimExpr sz = op->buffer->shape[0];
            auto nbits = op->buffer->dtype.bits() * op->buffer->dtype.lanes();
            if (const auto* imm = sz.as<IntImmNode>()) {
              if (imm->value > std::numeric_limits<int>::max() / nbits) {
                LOG(WARNING) << "The allocation requires : " << imm->value << " * " << nbits
                             << " bits, which is greater than the maximum of"
                                " int32. The size is cast to int64."
                             << "\n";
                sz = IntImm::Int64(imm->value);
              }
            }
            // transform to bits
            auto sz_nbits = sz * nbits;
            if (combo_size.defined()) {
              combo_size = max(combo_size, sz_nbits);
            } else {
              combo_size = sz_nbits;
            }
          }
          // transform to alloc bytes
          auto type_bits = alloc_type.bits() * alloc_type.lanes();
          bool divided = analyzer_->CanProve(indexmod(combo_size, type_bits) == 0);
          combo_size = indexdiv(combo_size, type_bits);
          // round up for can not divided
          if (!divided) {
            combo_size = combo_size + IntImm::Int32(1);
          }
          combo_size = analyzer_->Simplify(combo_size);
          BufferVar buf(e->alloc_var->name, BufferType(e->scope.to_string(), alloc_type,
                                                       {combo_size}, {}, PrimExpr(), 0, 0));
          e->alloc_var = buf.var();
          ffi::Map<ffi::String, ffi::Any> annotations;
          if (e->is_volatile) {
            annotations.Set(attr::kVolatile, true);
          }
          e->alloc_nest.push_back(AllocBuffer(buf, annotations));
        }
      }
    }
  }
  // New allocation for merged data
  void NewAllocTagMerged(StorageEntry* e) {
    TVM_FFI_ICHECK_NE(e->scope.tag.length(), 0U);
    // allocate with element type.
    TVM_FFI_ICHECK_NE(e->const_nbits, 0U);
    uint64_t total_bits = e->const_nbits;
    // By default, align to 32 bits.
    size_t align = 32;
    // Always align to max_simd_bits
    // so we can remap types by keeping this property
    if (total_bits % align != 0) {
      total_bits += align - (total_bits % align);
    }
    e->alloc_var = e->allocs[0]->buffer.var();
    for (StorageEntry* child : e->merged_children) {
      TVM_FFI_ICHECK_NE(child->const_nbits, 0U);
      TVM_FFI_ICHECK_NE(total_bits, 0U);
      child->bits_offset = total_bits;
      child->alloc_var = e->alloc_var;
      total_bits += child->const_nbits;
      if (total_bits % align != 0) {
        total_bits += align - (total_bits % align);
      }
    }
    uint64_t type_bits = e->elem_type.bits() * e->elem_type.lanes();
    PrimExpr alloc_size =
        MakeConst(e->allocs[0]->buffer->shape[0].ty(), (total_bits + type_bits - 1) / type_bits);
    BufferVar buf(e->alloc_var->name, BufferType(e->scope.to_string(), e->elem_type, {alloc_size},
                                                 {}, PrimExpr(), 0, 0));
    e->alloc_var = buf.var();
    for (StorageEntry* child : e->merged_children) {
      child->alloc_var = e->alloc_var;
    }
    bool any_volatile = e->is_volatile;
    for (StorageEntry* child : e->merged_children) {
      if (child->is_volatile) any_volatile = true;
    }
    ffi::Map<ffi::String, ffi::Any> annotations;
    if (any_volatile) {
      annotations.Set(attr::kVolatile, true);
    }
    e->alloc_nest.push_back(AllocBuffer(buf, annotations));
  }
  // Liveness analysis to find gen and kill point of each variable.
  void LivenessAnalysis(const std::vector<StmtEntry>& seq) {
    // find kill point, do a reverse linear scan.
    std::unordered_set<const VarNode*> touched;
    for (size_t i = seq.size(); i != 0; --i) {
      const StmtEntry& s = seq[i - 1];
      for (const VarNode* buffer : s.touched) {
        if (!touched.count(buffer)) {
          touched.insert(buffer);
          event_map_[s.stmt].kill.push_back(buffer);
        }
      }
    }
    // find gen point, do forward scan
    touched.clear();
    for (size_t i = 0; i < seq.size(); ++i) {
      int64_t offset = seq[i].scope_pair_offset;
      if (offset < 0) continue;
      const StmtEntry& s = seq[i + offset];
      for (const VarNode* buffer : s.touched) {
        if (!touched.count(buffer)) {
          touched.insert(buffer);
          event_map_[s.stmt].gen.push_back(buffer);
        }
      }
    }
  }
  void PlanNewScope(const ffi::Object* op) {
    if (thread_scope_ != nullptr) {
      TVM_FFI_ICHECK(thread_scope_ == op);
      // erase all memory atatched to this scope.
      for (auto it = const_free_map_.begin(); it != const_free_map_.end();) {
        if (it->second->attach_scope_ == op) {
          it = const_free_map_.erase(it);
        } else {
          ++it;
        }
      }
      for (auto it = sym_free_list_.begin(); it != sym_free_list_.end();) {
        if ((*it)->attach_scope_ == op) {
          it = sym_free_list_.erase(it);
        } else {
          ++it;
        }
      }
      thread_scope_ = nullptr;
    } else {
      thread_scope_ = op;
    }
  }

  // Memory plan algorithm
  void PlanMemory(const std::vector<StmtEntry>& seq,
                  const std::unordered_map<const VarNode*, AllocEntry>& alloc_info,
                  bool enable_reuse, bool reuse_require_exact_matched_dtype) {
    std::unordered_set<const VarNode*> inplace_flag;

    for (size_t i = 0; i < seq.size(); ++i) {
      const StmtEntry& s = seq[i];
      auto it = event_map_.find(seq[i].stmt);

      // scope_pair_offset >= 0 means it is either
      // - leaf stmt(offset = 0)
      // - beginning of scope(offset < 0)
      // In both cases, we need to handle the gen event correctly
      if (it != event_map_.end() && seq[i].scope_pair_offset >= 0) {
        // Inplace operation detection
        // specially handle this
        bool detect_inplace = detect_inplace_ && (it->second.gen.size() <= 2);

        for (const VarNode* var : it->second.gen) {
          TVM_FFI_ICHECK(alloc_info.count(var));
          const AllocEntry& entry = alloc_info.at(var);
          const AllocBufferNode* alloc = entry.alloc;
          auto storage_scope = StorageScope::Create(GetPtrStorageScope(ffi::GetRef<Var>(var)));
          StorageEntry* dst_entry = nullptr;
          // inplace detection
          if (detect_inplace) {
            // only one inplace var for s.stmt
            bool inplace_found = false;
            for (const VarNode* src : it->second.kill) {
              if (!inplace_flag.count(src) && alloc_map_.count(src)) {
                InplaceOpVerifier visitor;
                StorageEntry* src_entry = alloc_map_.at(src);
                if (src_entry->scope == storage_scope &&
                    src_entry->attach_scope_ == thread_scope_ &&
                    !alloc->buffer->dtype.IsScalableVector() &&
                    src_entry->elem_type == alloc->buffer->dtype.WithLanes(1) &&
                    visitor.Check(s.stmt, var, src)) {
                  int64_t const_size = AllocBuffer(ffi::GetRef<AllocBuffer>(alloc))
                                           .ConstantAllocationSize()
                                           .value_or(0);
                  uint64_t const_nbits = static_cast<uint64_t>(const_size) *
                                         alloc->buffer->dtype.bits() * alloc->buffer->dtype.lanes();
                  if (src_entry->const_nbits == const_nbits && !inplace_found) {
                    // successfully inplace
                    dst_entry = src_entry;
                    inplace_flag.insert(src);
                    inplace_found = true;
                  }
                }
              }
            }
          }
          if (dst_entry == nullptr) {
            dst_entry =
                FindAlloc(alloc, thread_scope_, storage_scope, entry.num_physical_dimensions,
                          enable_reuse, reuse_require_exact_matched_dtype);
          }
          dst_entry->allocs.emplace_back(alloc);
          if (alloc->annotations.count(attr::kVolatile)) {
            dst_entry->is_volatile = true;
          }
          alloc_map_[var] = dst_entry;
        }
      }
      // enter/exit new scope
      if (s.stmt->IsInstance<AttrStmtNode>()) {
        const auto* op = static_cast<const AttrStmtNode*>(s.stmt);
        if (op->attr_key == attr::thread_extent || op->attr_key == s_tir::attr::virtual_thread ||
            attr::IsPragmaKey(op->attr_key)) {
          PlanNewScope(op);
        } else {
          TVM_FFI_ICHECK(op->attr_key == attr::extern_scope);
        }
      } else if (s.stmt->IsInstance<ForNode>()) {
        const auto* op = static_cast<const ForNode*>(s.stmt);
        if (op->kind == ForKind::kParallel) {
          if (thread_scope_ == nullptr || thread_scope_ == op) {
            PlanNewScope(op);
          }
        }
      }
      // scope_pair_offset <= 0 means it is either
      // - leaf stmt(offset = 0)
      // - end of scope(offset < 0)
      // In both cases, we need to handle the kill event correctly
      if (it != event_map_.end() && seq[i].scope_pair_offset <= 0) {
        for (const VarNode* var : it->second.kill) {
          // skip space which are already replaced by inplace
          if (!inplace_flag.count(var)) {
            this->Free(var);
          }
        }
      }
    }
  }
  // Allocate new storage entry.
  StorageEntry* NewAlloc(const AllocBufferNode* op, const ffi::Object* attach_scope,
                         const StorageScope& scope, size_t const_nbits) {
    TVM_FFI_ICHECK(op != nullptr);
    // Re-use not successful, allocate a new buffer.
    auto entry = std::make_unique<StorageEntry>();
    entry->attach_scope_ = attach_scope;
    entry->scope = scope;
    entry->elem_type = op->buffer->dtype.WithLanes(1);
    entry->const_nbits = const_nbits;
    StorageEntry* e = entry.get();
    alloc_vec_.emplace_back(std::move(entry));
    return e;
  }

  StorageEntry* FindAlloc(const AllocBufferNode* op, const ffi::Object* attach_scope,
                          const StorageScope& scope, size_t num_physical_dimensions,
                          bool enable_reuse, bool reuse_require_exact_matched_dtype) {
    TVM_FFI_ICHECK(op != nullptr);
    // skip plan for local variable,
    // compiler can do a better job with register allocation.
    const uint64_t match_range = 16;
    bool is_scalable_vector = op->buffer->dtype.IsScalableVector();
    uint64_t op_elem_bits =
        is_scalable_vector ? 0 : op->buffer->dtype.bits() * op->buffer->dtype.lanes();
    int64_t const_size =
        AllocBuffer(ffi::GetRef<AllocBuffer>(op)).ConstantAllocationSize().value_or(0);
    uint64_t const_nbits =
        is_scalable_vector ? 0 : static_cast<uint64_t>(const_size * op_elem_bits);

    // If the size of the array isn't known at compile-time, it must
    // have its own allocation with size determined at runtime.
    bool is_known_size = (const_nbits != 0);

    // Currently, only flat memory spaces can be re-used.  Packing
    // into N-d space (e.g. 2-d texture memory on GPUs) will require
    // more in-depth algorithms.
    bool is_flat_memory_space = (num_physical_dimensions == 1);

    // disable reuse of small arrays, they will be lowered to registers in LLVM
    // This rules only apply if we are using non special memory
    bool is_small_array = (scope.tag.length() == 0) && (scope.rank >= StorageRank::kWarp ||
                                                        (is_known_size && const_nbits <= 32));

    if (is_scalable_vector || !enable_reuse || is_small_array || !is_flat_memory_space) {
      return NewAlloc(op, attach_scope, scope, const_nbits);
    }

    if (is_known_size) {
      // constant allocation.
      auto begin = const_free_map_.lower_bound(const_nbits / match_range);
      auto mid = const_free_map_.lower_bound(const_nbits);
      auto end = const_free_map_.upper_bound(const_nbits * match_range);
      // start looking at the buffer that is bigger than the required size first
      for (auto it = mid; it != end; ++it) {
        StorageEntry* e = it->second;
        if (e->attach_scope_ != attach_scope) continue;
        if (e->scope != scope) continue;
        // when not divided, no reuse, eg, float4 vs float3
        if (e->bits_offset % op_elem_bits != 0) continue;
        if (reuse_require_exact_matched_dtype && e->elem_type != op->buffer->dtype) {
          continue;
        }
        e->const_nbits = std::max(const_nbits, e->const_nbits);
        const_free_map_.erase(it);
        return e;
      }
      // then start looking at smaller buffers.
      for (auto it = mid; it != begin;) {
        --it;
        StorageEntry* e = it->second;
        if (e->attach_scope_ != attach_scope) continue;
        if (e->scope != scope) continue;
        if (e->elem_type != op->buffer->dtype.WithLanes(1)) continue;
        if (reuse_require_exact_matched_dtype && e->elem_type != op->buffer->dtype) {
          continue;
        }
        e->const_nbits = std::max(const_nbits, e->const_nbits);
        const_free_map_.erase(it);
        return e;
      }
    } else {
      // Simple strategy: round roubin.
      for (auto it = sym_free_list_.begin(); it != sym_free_list_.end(); ++it) {
        StorageEntry* e = *it;
        if (e->attach_scope_ != attach_scope) continue;
        if (e->scope != scope) continue;
        if (e->elem_type != op->buffer->dtype.WithLanes(1)) continue;
        sym_free_list_.erase(it);
        return e;
      }
    }
    return NewAlloc(op, attach_scope, scope, const_nbits);
  }
  // simulated free.
  void Free(const VarNode* var) {
    auto it = alloc_map_.find(var);
    TVM_FFI_ICHECK(it != alloc_map_.end());
    StorageEntry* e = it->second;
    TVM_FFI_ICHECK_NE(e->allocs.size(), 0U);

    // disable reuse of small arrays, they will be lowered to registers in LLVM
    // This rules only apply if we are using non special memory
    if (e->scope.tag.length() == 0) {
      // Disable sharing of local memory.
      if (e->scope.rank >= StorageRank::kWarp || e->allocs[0]->buffer->dtype.IsScalableVector()) {
        return;
      }
      // disable reuse of small arrays
      if (e->const_nbits > 0 && e->const_nbits <= 32) return;
    }
    // normal free.
    if (e->const_nbits != 0) {
      const_free_map_.insert({e->const_nbits, e});
    } else {
      sym_free_list_.push_back(e);
    }
  }
  // thread scope.
  const ffi::Object* thread_scope_{nullptr};
  // whether enable inplace detection.
  bool detect_inplace_{false};
  // Locations of free ops.
  std::unordered_map<const ffi::Object*, EventEntry> event_map_;
  // constant size free map.
  std::multimap<uint64_t, StorageEntry*> const_free_map_;
  // symbolic free list, for non constant items.
  std::list<StorageEntry*> sym_free_list_;
  // The allocation attach map
  std::unordered_map<const ffi::Object*, std::vector<StorageEntry*>> attach_map_;
  // The allocation assign map
  std::unordered_map<const VarNode*, StorageEntry*> alloc_map_;
  // The allocations
  std::vector<std::unique_ptr<StorageEntry>> alloc_vec_;
  // The buffer objects being remapped
  std::unordered_map<const VarNode*, BufferVar> buffer_remap_;
  // Physical backing chosen for each remapped buffer view.
  std::unordered_map<const VarNode*, BufferVar> remapped_backing_;
  // Physical roots of buffer aliases, flattened by LinearAccessPatternFinder.
  ffi::Map<Var, Var> buffer_aliases_;
  // analyzer
  arith::Analyzer analyzer_;
};

/* Helper struct containing information on how a buffer is declared and used
 *
 */
struct BufferVarInfo {
  enum DeclarationLocation {
    kPrimFuncBufferParam = (1 << 0),
    kPrimFuncPointerParam = (1 << 1),
    kAllocBufferNode = (1 << 2),
    kLetNode = (1 << 3),
    kDeclBufferNode = (1 << 4),
  };

  // The tirx::Var that represents this buffer.
  Var var;

  // The data type of an element of the buffer.
  PrimType element_dtype;

  /* The extent of the buffer.
   *
   * If multidimensional, the extent of the last dimension of the buffer.  If the
   * size is unknown (e.g. pointer arguments to PrimFunc without a BufferType
   * annotation), then extent is zero.
   */
  PrimExpr extent;

  // Where the buffer was declared
  DeclarationLocation declaration_location;

  // When accessed, which element type is it accessed as.  This may
  // differ both in base type (e.g. int32* cast to float32* after
  // packing in StorageRewrite) or in number of lanes (e.g. float16*
  // cast to float16x4*).
  std::unordered_set<PrimType, PrimTypeHash, PrimTypeEqual> access_dtype;
  // Data types used for scalar reads. This is used to record vectorized read dtypes that can be
  // shuffled for scalar reads when rewrite_scalar_read_to_vector_shuffle is enabled.
  std::unordered_set<PrimType, PrimTypeHash, PrimTypeEqual> scalar_read_dtype;

  PrimType get_preferred_dtype() const {
    std::unordered_set<PrimType, PrimTypeHash, PrimTypeEqual> base_access_dtype;
    for (auto dtype : access_dtype) {
      base_access_dtype.insert(dtype.WithLanes(1));
    }
    for (auto dtype : scalar_read_dtype) {
      base_access_dtype.insert(dtype.WithLanes(1));
    }
    // If the array is accessed as multiple base types within a
    // function, no point in changing the declared type.  CodeGenC can
    // handle this with a type-cast prior to indexing.  Vulkan will
    // raise an error at code-gen time, if a later pass doesn't split
    // it out.
    if (base_access_dtype.size() != 1) {
      return element_dtype;
    }

    PrimType preferred_base_type = *base_access_dtype.begin();

    // If there is only one vectorizable size used to access the
    // buffer, and if that access size is compatible with the array
    // size, then the buffer is vectorizable.  In the future, this
    // could be improved to allow vectorized buffer access of size
    // GCD(*lanes_used), if necessary.
    // When there are scalar reads and no writes, access_dtype can be empty and we should avoid
    // rewriting.
    int preferred_lanes = element_dtype.lanes();
    if (element_dtype.lanes() == 1 && (access_dtype.size() == 1)) {
      int lanes = access_dtype.begin()->lanes();
      // Check the scalar read dtypes are compatible with the vectorized access dtype.
      for (auto dtype : scalar_read_dtype) {
        if (dtype.lanes() % lanes != 0) {
          return element_dtype;
        }
      }
      arith::Analyzer analyzer_;
      arith::ModularSet me = analyzer_->modular_set(extent);
      if ((me->coeff % lanes == 0) && (me->base % lanes == 0)) {
        preferred_lanes = lanes;
      }
    }

    return preferred_base_type.WithLanes(preferred_lanes);
  }
};

/* Checks whether buffers are accessed as scalar or vector parameters in a
 * function.
 *
 */
class VectorTypeAccessChecker : public StmtExprVisitor {
 public:
  /* Constructor
   *
   * @param params The parameters passed to a PrimFunc
   *
   * @param allow_untyped_handles If a buffer or pointer variable is
   * missing a type annotation, assume that it has the same underlying
   * type as it is later accessed, with scalar element types.
   */
  VectorTypeAccessChecker(const ffi::Array<tirx::Var>& params, bool allow_untyped_pointers = false,
                          bool detect_scalar_read_patterns = true)
      : allow_untyped_pointers_(allow_untyped_pointers),
        detect_scalar_read_patterns_(detect_scalar_read_patterns) {
    for (Var buffer_var : params) {
      if (auto buffer = buffer_var.as<BufferVar>()) {
        buffer_aliases_.Set(buffer_var, buffer_var);
        PrimExpr extent = buffer.value()->shape.size()
                              ? buffer.value()->shape[buffer.value()->shape.size() - 1]
                              : PrimExpr(0);
        OnArrayDeclaration(buffer_var, buffer.value()->dtype, extent,
                           BufferVarInfo::kPrimFuncBufferParam);
        continue;
      }
      auto pointer_type = GetPointerType(buffer_var->ty);
      if (pointer_type.has_value() && !pointer_type.value().IsVoid()) {
        PrimType dtype = pointer_type.value();
        PrimExpr extent = 0;
        OnArrayDeclaration(buffer_var, dtype, extent, BufferVarInfo::kPrimFuncPointerParam);
      } else if (pointer_type.has_value() && allow_untyped_pointers_) {
        OnArrayDeclaration(buffer_var, PrimType::Void(), 0, BufferVarInfo::kPrimFuncPointerParam);
      }
    }
  }

  void VisitExpr_(const BufferLoadNode* op) final {
    OnArrayAccess(op->ty.as_or_throw<PrimType>(), op->buffer.get(), op->indices,
                  /*is_buffer_load=*/true);
    StmtExprVisitor::VisitExpr_(op);
  }

  void VisitStmt_(const BufferStoreNode* op) final {
    OnArrayAccess(op->value.ty(), op->buffer.get(), op->indices, /*is_buffer_load=*/false);
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitExpr_(const CallNode* op) final {
    if (op->op.same_as(builtin::tvm_access_ptr())) {
      PrimType dtype = op->args[0].as_or_throw<PrimExpr>().ty();
      auto buffer_var = GetBufferDataVar(op->args[1]);
      PrimExpr index = op->args[2].as_or_throw<PrimExpr>();
      // args[1] may be a nested Call (e.g. another tvm_access_ptr) rather
      // than a raw Var; OnArrayAccess derefs `buffer` so skip the record
      // here and let the recursive visit handle any inner buffer var.
      if (buffer_var.has_value()) {
        OnArrayAccess(dtype, buffer_var.value().get(), {index}, false);
      }
    } else if (op->op.same_as(builtin::address_of())) {
      if (const auto* load = op->args[0].as<BufferLoadNode>()) {
        OnArrayAccess(load->ty.as_or_throw<PrimType>(), load->buffer.get(), load->indices,
                      /*is_buffer_load=*/false);
      }
    }
    StmtExprVisitor::VisitExpr_(op);
  }

  void VisitStmt_(const AllocBufferNode* op) final {
    buffer_aliases_.Set(op->buffer.var(), op->buffer.var());
    const ffi::Array<PrimExpr>& shape = op->buffer->shape;
    PrimExpr extent = shape.size() ? shape[shape.size() - 1] : PrimExpr(0);
    OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent,
                       BufferVarInfo::kAllocBufferNode);

    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitStmt_(const DeclBufferNode* op) final {
    RegisterBufferAlias(op->buffer, op->data);
    const ffi::Array<PrimExpr>& shape = op->buffer->shape;
    PrimExpr extent = shape.size() ? shape.back() : PrimExpr(0);
    OnArrayDeclaration(op->buffer.var(), op->buffer->dtype, extent, BufferVarInfo::kDeclBufferNode);
    StmtExprVisitor::VisitStmt_(op);
  }

  void VisitExpr_(const LetNode* op) final {
    HandleLetNode(op->var);
    StmtExprVisitor::VisitExpr_(op);
  }

  void VisitStmt_(const BindNode* op) final {
    HandleLetNode(op->var);
    StmtExprVisitor::VisitStmt_(op);
  }

  void HandleLetNode(Var let_var) {
    auto pointer_type = GetPointerType(let_var->ty);
    if (pointer_type.has_value()) {
      if (!pointer_type.value().IsVoid()) {
        OnArrayDeclaration(let_var, pointer_type.value(), 0, BufferVarInfo::kLetNode);
      } else if (allow_untyped_pointers_) {
        OnArrayDeclaration(let_var, PrimType::Void(), 0, BufferVarInfo::kLetNode);
      }
      return;
    }
  }

  /* Update the type map for a buffer based on its declaration
   *
   * @param buffer The VarNode representing the buffer.
   *
   * @param element_dtype The dtype of a single element of the buffer.
   * If unknown, when used with the allow_untyped_pointers option,
   * should be the primitive void sentinel.
   *
   * @param extent The extent of the buffer.  Zero if size is unknown.
   *
   * @param declaration_location How the buffer was allocated, so that
   * some locations can be rewritten without others.
   */
  void OnArrayDeclaration(Var buffer, PrimType element_dtype, PrimExpr extent,
                          BufferVarInfo::DeclarationLocation declaration_location) {
    TVM_FFI_ICHECK(info_map_.find(buffer.get()) == info_map_.end())
        << "Array declaration of " << buffer->name << " occurred multiple times.";

    if (element_dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
      element_dtype = PrimType::Int(8, element_dtype.lanes());
    }
    info_map_.insert_or_assign(buffer.get(),
                               BufferVarInfo{buffer, element_dtype, extent, declaration_location});
  }

  /* Update the type map for a buffer based on its usage
   *
   * @param value_dtype The dtype of the value being stored to or
   * loaded from the buffer.
   *
   * @param buffer The VarNode representing the buffer.
   *
   * @param indices The index at which the value is being stored/loaded.
   *
   * @param is_buffer_load Whether the access is BufferLoad
   */
  void OnArrayAccess(PrimType value_dtype, const VarNode* buffer,
                     const ffi::Array<PrimExpr>& indices, bool is_buffer_load) {
    Var buffer_var = ffi::GetRef<Var>(buffer);
    const VarNode* root = buffer_aliases_.Get(buffer_var).value_or(buffer_var).get();
    if (info_map_.count(root)) {
      buffer = root;
    }
    auto it = info_map_.find(buffer);
    TVM_FFI_ICHECK(it != info_map_.end()) << "Load/Store of buffer " << buffer->name << " ("
                                          << buffer << ") occurred before its declaration.";

    if (value_dtype.IsScalableVector()) {
      // Scalable types are not currently supported in storage_rewrite. Scalable buffer
      // accesses are not currently checked and therefore are not rewritten.
      return;
    }

    BufferVarInfo& var_info = it->second;

    if (value_dtype.WithLanes(1).MatchesCode(DLDataTypeCode::kDLBool)) {
      value_dtype = PrimType::Int(8, value_dtype.lanes());
    }

    if (var_info.element_dtype.IsVoid()) {
      TVM_FFI_ICHECK(allow_untyped_pointers_)
          << "Variable " << buffer->name << " was missing a type annotation in its declaration";
      var_info.element_dtype = value_dtype.WithLanes(1);
    }

    for (int i = 0; i < static_cast<int>(indices.size()) - 1; i++) {
      TVM_FFI_ICHECK(indices[i].ty().IsScalar())
          << "Only the last index of a buffer access may be a vector type.";
    }
    int index_lanes = indices.size() ? indices.back().ty().lanes() : 1;

    PrimType access_dtype = value_dtype;

    int lanes_used = var_info.element_dtype.lanes();

    // This can happen due to a previous pass that had rewrite_store_load =
    // false.  This occurs from the StorageRewrite in tvm::lower, followed by the
    // PointerValueTypeRewrite in BuildSPIRV.  The rewrite_store_load = false is
    // necessary because the C-based codegens do not yet support vectorized
    // pointer types (e.g. float16x4*).  Once they do, this if statement should
    // instead be replaced by the below TVM_FFI_ICHECK_EQ.
    if (index_lanes * var_info.element_dtype.lanes() != value_dtype.lanes()) {
      TVM_FFI_ICHECK_EQ(index_lanes, value_dtype.lanes());
      lanes_used = 1;
      var_info.element_dtype = var_info.element_dtype.WithLanes(1);
    }

    // TODO(Lunderberg): Uncomment this check once it can be applied.
    // See https://discuss.tvm.apache.org/t/pre-rfc-vectorized-tirx-buffers/10615
    // for discussion.

    // TVM_FFI_ICHECK_EQ(index_lanes * var_info.element_dtype.lanes(), value_dtype.lanes())
    //     << "Attempting to retrieve " << value_dtype.lanes() << " lanes of data with "
    //     << index_lanes << " indices into an array whose elements have "
    //     << var_info.element_dtype.lanes() << " lanes.  "
    //     << "Expected output with " << index_lanes * var_info.element_dtype.lanes()
    //     << " lanes.";

    // If the index is a RampNode with stride of 1 and offset
    // divisible by the number of number of lanes, and the predicate
    // does not apply any masking, then this array access could be
    // vectorized.
    if (indices.size()) {
      const RampNode* ramp_index = indices[indices.size() - 1].as<RampNode>();
      if (ramp_index && is_one(ramp_index->stride)) {
        if (ramp_index->lanes->IsInstance<IntImmNode>()) {
          int lanes = static_cast<int>(ramp_index->lanes.as_or_throw<IntImm>()->value);
          arith::ModularSet me = analyzer_->modular_set(ramp_index->base);
          if ((me->coeff % lanes == 0) && (me->base % lanes == 0)) {
            lanes_used = lanes;
          }
        }
      }
    }

    if (detect_scalar_read_patterns_ && is_buffer_load && indices.size()) {
      const PrimExpr last_dim_index = indices[indices.size() - 1];
      if (last_dim_index.ty().lanes() == 1) {
        arith::ModularSet me = analyzer_->modular_set(last_dim_index);
        var_info.scalar_read_dtype.emplace(access_dtype.WithLanes(me->coeff));
        return;
      }
    }
    var_info.access_dtype.insert(access_dtype.WithLanes(lanes_used));
  }

 private:
  void RegisterBufferAlias(BufferVar buffer, const Expr& data) {
    Var root = buffer.var();
    if (auto source = GetBufferDataVar(data);
        source.has_value() && source.value()->ty.as<BufferTypeNode>()) {
      auto source_root = buffer_aliases_.Get(source.value());
      TVM_FFI_ICHECK(source_root.has_value()) << "Buffer alias source " << source.value()->name
                                              << " must be registered before its DeclBuffer alias";
      root = source_root.value();
    }
    buffer_aliases_.Set(buffer.var(), root);
  }

 public:
  // Map of buffer variable information determined
  std::unordered_map<const VarNode*, BufferVarInfo> info_map_;
  // Physical roots of buffer aliases, flattened at each declaration.
  ffi::Map<Var, Var> buffer_aliases_;

 public:
  //
  bool allow_untyped_pointers_{false};
  // Whether to detect scalar read patterns for rewriting to vector shuffle
  bool detect_scalar_read_patterns_{true};

  // internal analyzer
  arith::Analyzer analyzer_;
};

/* \brief Rewrites buffer/pointer variables from scalar types to vectorized
 * types.
 *
 * Some runtimes do not allow casting between composite types and the underlying
 * base type (e.g. Vulkan, casting from 1-lane float16* to 4-lane float16x4*).
 * In these cases, in order to have vectorized load/store on an array, the
 * element type of that array must be vectorized.  This is in contrast to C-style
 * runtimes, in which `float16x4* vec = *(float16x4*)(float_arr + offset)` is
 * valid.
 *
 * By default, VectorTypeRewriter will attempt to rewrite all buffer variables to
 * vectorized access, if the load/store occurring in the PrimFunc are all
 * vectorized.  This includes adjusting the indices being used to access the
 * array.  (e.g. If `float16* scalar_arr` is being converted to `float16x4*
 * vec_arr`, then `scalar_arr[Ramp(offset, 1, 4)]` will be converted to
 * `vec_arr[offset/4]`.)
 *
 * Currently, several of the C-style runtimes do not support buffers whose
 * elements are vectorized types, or rely on the presence of the Ramp nodes to
 * identify vectorized loads.  The boolean parameters in the constructor are to
 * mimic the previous behavior of VectorTypeRewriter, to avoid breaking these
 * runtimes.  Once all runtimes support vectorized buffer elements, these
 * parameters can be removed.
 */
class VectorTypeRewriter : public StmtExprMutator {
 public:
  /* Constructor
   *
   * @param checker The VectorTypeAccessChecker that has previously read out
   * information from the PrimFunc
   *
   * @param rewrite_buffer_params Whether BufferType-annotated parameters should
   * be rewritten from scalar element types to vectorized element types.
   *
   * @param rewrite_pointer_params Whether pointer-typed parameters should be
   * rewritten from scalar types to vectorized types.
   *
   * @param rewrite_alloc_buffer_node Whether the buffer variable associated with
   * AllocBufferNodes should be rewritten from scalar types to vectorized types.
   *
   * @param rewrite_indices Whether the indices to the Load and Store nodes
   * should be rewritten to correspond to the new buffer_var type.
   *
   * @param rewrite_let_node Whether pointer declarations in let nodes
   * should be re-written.
   */
  VectorTypeRewriter(const std::unordered_map<const VarNode*, BufferVarInfo>& info_map,
                     const ffi::Map<Var, Var>& buffer_aliases, bool rewrite_buffer_params = true,
                     bool rewrite_pointer_params = true, bool rewrite_alloc_buffer_node = true,
                     bool rewrite_indices = true, bool rewrite_let_node = true,
                     bool rewrite_scalar_read_to_vector_shuffle = true)
      : rewrite_indices_(rewrite_indices), buffer_aliases_(buffer_aliases) {
    int rewrite_mask = 0;
    if (rewrite_buffer_params) {
      rewrite_mask |= BufferVarInfo::kPrimFuncBufferParam;
    }
    if (rewrite_pointer_params) {
      rewrite_mask |= BufferVarInfo::kPrimFuncPointerParam;
    }
    if (rewrite_alloc_buffer_node) {
      rewrite_mask |= BufferVarInfo::kAllocBufferNode;
    }
    if (rewrite_let_node) {
      rewrite_mask |= BufferVarInfo::kLetNode;
    }
    // Rewrite any buffer variables whose preferred type isn't their current type.
    for (const auto& pair : info_map) {
      const auto& var_info = pair.second;
      PrimType preferred = var_info.get_preferred_dtype();
      if (preferred != var_info.element_dtype && (rewrite_mask & var_info.declaration_location)) {
        Var old_buffer_var = var_info.var;
        Var new_buffer_var = [&]() -> Var {
          if (old_buffer_var->ty.as<BufferTypeNode>()) {
            BufferVar old_buffer(old_buffer_var);
            auto type = CopyBufferType(old_buffer);
            type->dtype = preferred;
            if (!type->shape.empty()) {
              PrimExpr last_dim = type->shape.back();
              int factor = preferred.lanes() / var_info.element_dtype.lanes();
              type->shape.Set(type->shape.size() - 1, last_dim / MakeConst(last_dim.ty(), factor));
            }
            type->layout = std::nullopt;
            return RebuildBufferVar(old_buffer, std::move(type)).var();
          }
          return Var(old_buffer_var->name,
                     PointerType(preferred, GetPtrStorageScope(old_buffer_var)),
                     old_buffer_var->span);
        }();

        rewrite_map_.insert_or_assign(
            var_info.var.get(),
            RewriteInfo{var_info.var, new_buffer_var, var_info.element_dtype, preferred});
      }
    }
  }

  /*!
   * \brief Mutator for BufferLoad or BufferStore.
   * \return The rewritten node and the shuffle index. (Only for BufferLoad) When the shuffle index
   * is non-negative, the caller should generate Shuffle to extract the element from the vector.
   */
  template <typename Node>
  std::pair<Node, int> VisitBufferAccess(Node node) {
    int shuffle_index = -1;
    if (!rewrite_indices_) {
      return {node, shuffle_index};
    }

    Var root = buffer_aliases_.Get(node->buffer.var()).value_or(node->buffer.var());
    auto it = rewrite_map_.find(root.get());
    if (it == rewrite_map_.end()) {
      return {node, shuffle_index};
    }
    const auto& info = it->second;

    ffi::Array<PrimExpr> indices = node->indices;
    const PrimExpr& last_dim_index = indices[indices.size() - 1];
    const RampNode* ramp_index = indices[indices.size() - 1].as<RampNode>();

    if (node->buffer->dtype.IsScalableVector() || last_dim_index.ty().IsScalableVector()) {
      // Scalable types are not currently supported in storage_rewrite. Scalable buffer
      // accesses are not currently checked and therefore are not rewritten.
      return {node, shuffle_index};
    }

    if (ramp_index && is_one(ramp_index->stride) && ramp_index->lanes->IsInstance<IntImmNode>()) {
      int lanes = static_cast<int>(ramp_index->lanes.as_or_throw<IntImm>()->value);
      PrimExpr new_index = ramp_index->base / MakeConst(ramp_index->base.ty(), lanes);
      if (lanes != info.factor()) {
        TVM_FFI_ICHECK(info.factor() && lanes % info.factor() == 0);
        int new_lanes = lanes / info.factor();
        new_index = Ramp(new_index * new_lanes, ramp_index->stride, new_lanes, ramp_index->span);
      }
      indices.Set(indices.size() - 1, new_index);
    } else if (last_dim_index.ty().lanes() == 1 && info.factor() > 1) {
      arith::ModularSet me = analyzer_->modular_set(last_dim_index);
      TVM_FFI_ICHECK(me->coeff == 0 || info.factor() % me->coeff == 0);
      PrimExpr new_index = last_dim_index / MakeConst(last_dim_index.ty(), info.factor());
      shuffle_index = me->base % info.factor();
      indices.Set(indices.size() - 1, new_index);
    }

    auto writer = node.CopyOnWrite();
    writer->buffer = RemapBuffer(node->buffer);
    writer->indices = indices;
    return {node, shuffle_index};
  }

  Expr VisitExpr_(const BufferLoadNode* op) final {
    auto node = StmtExprMutator::VisitExpr_(op).as_or_throw<BufferLoad>();
    auto [modified, shuffle_index] = VisitBufferAccess(node);

    // Not needed for BufferStoreNode, so we can't just call
    // LegalizeDtype() in VisitBufferAccess.
    if (node.same_as(modified)) {
      return node;

    } else {
      auto writer = modified.CopyOnWrite();
      writer->LegalizeDType();
      if (shuffle_index >= 0) {
        return Shuffle::ExtractElement(std::move(modified), shuffle_index);
      }
      return modified;
    }
  }

  Stmt VisitStmt_(const BufferStoreNode* op) final {
    auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<BufferStore>();
    auto [modified, shuffle_index] = VisitBufferAccess(std::move(node));
    TVM_FFI_ICHECK(shuffle_index < 0);
    return modified;
  }

  Stmt VisitStmt_(const BindNode* op) final {
    auto it = rewrite_map_.find(op->var.get());
    Expr value = this->VisitExpr(op->value);
    Var var = (it == rewrite_map_.end()) ? op->var : it->second.new_buffer_var;
    if (!ffi::StructuralEqual()(value->ty, var->ty)) {
      auto call = value.as_or_throw<Call>();
      value = Call(var->ty, call->op, call->args, call->attrs, call->ty_args, call->span);
    }
    if (var.same_as(op->var) && value.same_as(op->value)) {
      return ffi::GetRef<Stmt>(op);
    }
    return Bind(var, value);
  }

  Stmt VisitStmt_(const AllocBufferNode* op) final {
    BufferVar new_buf = RemapBuffer(op->buffer);
    if (new_buf.same_as(op->buffer)) {
      return ffi::GetRef<Stmt>(op);
    }
    auto n = CopyOnWrite(op);
    n->buffer = std::move(new_buf);
    return Stmt(n);
  }

  Stmt VisitStmt_(const DeclBufferNode* op) final {
    auto node = StmtExprMutator::VisitStmt_(op).as_or_throw<DeclBuffer>();
    BufferVar new_buf = RemapBuffer(node->buffer);
    if (!new_buf.same_as(node->buffer)) {
      node.CopyOnWrite()->buffer = new_buf;
    }
    return node;
  }

  BufferVar RemapBuffer(BufferVar buf) {
    auto cache_key = buf.get();

    auto cache_it = buffer_map_.find(cache_key);
    if (cache_it != buffer_map_.end()) {
      return cache_it->second;
    }

    Var root = buffer_aliases_.Get(buf.var()).value_or(buf.var());
    auto info_it = rewrite_map_.find(root.get());
    if (info_it != rewrite_map_.end()) {
      auto& info = info_it->second;
      if (root.same_as(buf.var())) {
        buf = BufferVar(info.new_buffer_var);
      } else {
        auto type = CopyBufferType(buf);
        type->dtype = info.new_element_dtype;
        if (!type->shape.empty()) {
          PrimExpr last_dim = type->shape.back();
          type->shape.Set(type->shape.size() - 1,
                          last_dim / MakeConst(last_dim.ty(), info.factor()));
        }
        type->layout = std::nullopt;
        buf = RebuildBufferVar(buf, std::move(type));
      }
    }

    buffer_map_[cache_key] = buf;
    return buf;
  }

  Expr VisitExpr_(const CallNode* op) final {
    if (op->op.same_as(builtin::buffer_data()) && op->args.size() == 1) {
      if (auto var = op->args[0].as<Var>();
          var.has_value() && var.value()->ty.as<BufferTypeNode>()) {
        return RemapBuffer(BufferVar(var.value())).data();
      }
    }
    if (op->op.same_as(builtin::tvm_access_ptr())) {
      Expr expr = StmtExprMutator::VisitExpr_(op);
      op = expr.as<CallNode>();

      if (!rewrite_indices_) {
        return expr;
      }

      auto buffer = GetBufferDataVar(op->args[1]);
      if (!buffer.has_value()) {
        return expr;
      }
      Var var = buffer.value();
      Var root = buffer_aliases_.Get(var).value_or(var);
      auto it = rewrite_map_.find(root.get());
      if (it == rewrite_map_.end()) {
        return expr;
      }
      const auto& info = it->second;

      PrimExpr index = op->args[2].as_or_throw<PrimExpr>();
      PrimExpr extent = op->args[3].as_or_throw<PrimExpr>();
      PrimExpr flag = op->args[4].as_or_throw<PrimExpr>();

      PrimExpr e_dtype = tirx::TypeAnnotation(info.new_element_dtype);
      int factor = info.factor();
      extent = extent / MakeConst(extent.ty(), factor);
      index = index / MakeConst(index.ty(), factor);
      Expr data = info.new_buffer_var->ty.as<BufferTypeNode>()
                      ? BufferVar(info.new_buffer_var).data()
                      : Expr(info.new_buffer_var);
      ffi::Array<Expr> acc_args{e_dtype, data, index, extent, flag};
      auto old_pointer_type = op->ty.as_or_throw<PointerType>();
      Type new_pointer_type = PointerType(info.new_element_dtype, old_pointer_type->storage_scope);
      return Call(new_pointer_type, builtin::tvm_access_ptr(), acc_args);

    } else {
      return StmtExprMutator::VisitExpr_(op);
    }
  }

  /* Update the parameters and all remaining variable references
   *
   * Should be called after calling operator() on the body of the
   * function.
   *
   * @param func A pointer to the PrimFunc being modified.
   */
  void Finalize(PrimFunc* func_ptr) {
    TVM_FFI_ICHECK(func_ptr) << "Finalize expects a non-null pointer";
    auto& func = *func_ptr;
    auto* n = func.CopyOnWrite();

    std::unordered_map<const VarNode*, Var> var_remap;
    for (const auto& [_, info] : rewrite_map_) {
      var_remap.emplace(info.old_buffer_var.get(), info.new_buffer_var);
    }
    class PointerVarSubstituter : public StmtExprMutator {
     public:
      explicit PointerVarSubstituter(const std::unordered_map<const VarNode*, Var>& var_remap)
          : var_remap_(var_remap) {}

     private:
      using StmtExprMutator::VisitExpr_;

      Expr VisitExpr_(const VarNode* op) final {
        if (auto it = var_remap_.find(op); it != var_remap_.end()) {
          return it->second;
        }
        return ffi::GetRef<Var>(op);
      }

      BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) final {
        auto it = var_remap_.find(buffer.get());
        if (it != var_remap_.end()) {
          BufferVar new_buffer(it->second);
          buffer_remap_.Set(buffer, new_buffer);
          return new_buffer;
        }
        return StmtExprMutator::VisitBufferDef(buffer, alloc_data);
      }

      Stmt VisitStmt_(const AttrStmtNode* op) final {
        Stmt stmt = StmtExprMutator::VisitStmt_(op);
        op = stmt.as<AttrStmtNode>();
        TVM_FFI_ICHECK(op != nullptr);
        if (auto var = op->node.as<Var>()) {
          if (auto it = var_remap_.find(var.value().get()); it != var_remap_.end()) {
            return AttrStmt(it->second, op->attr_key, op->value, op->body, op->span);
          }
        }
        return stmt;
      }

      const std::unordered_map<const VarNode*, Var>& var_remap_;
    };

    n->body = PointerVarSubstituter(var_remap)(n->body);

    // Remap the argument list to use the new buffer variables.
    ffi::Array<Var> new_params;
    for (const auto& old_param : n->params) {
      auto it = rewrite_map_.find(old_param.get());
      if (it == rewrite_map_.end()) {
        new_params.push_back(old_param);
      } else {
        const auto& info = it->second;
        new_params.push_back(info.new_buffer_var);
      }
    }
    n->params = new_params;
  }

 private:
  struct RewriteInfo {
    Var old_buffer_var;
    Var new_buffer_var;
    PrimType old_element_dtype;
    PrimType new_element_dtype;

    int factor() const {
      int old_lanes = old_element_dtype.lanes();
      int new_lanes = new_element_dtype.lanes();
      TVM_FFI_ICHECK_EQ(new_lanes % old_lanes, 0);
      return new_lanes / old_lanes;
    }
  };

  bool rewrite_indices_{true};
  std::unordered_map<const VarNode*, RewriteInfo> rewrite_map_;
  std::unordered_map<const VarNode*, BufferVar> buffer_map_;
  const ffi::Map<Var, Var>& buffer_aliases_;
  arith::Analyzer analyzer_;
};

// Rewrite allocates, pointer parameters, and buffer parameters into vectorized versions
// if each access into a buffer is the same vector type.
PrimFunc PointerValueTypeRewrite(PrimFunc f, bool allow_untyped_pointers = false,
                                 bool rewrite_buffer_params = true,
                                 bool rewrite_pointer_params = true,
                                 bool rewrite_alloc_buffer_node = true, bool rewrite_indices = true,
                                 bool rewrite_let_node = true,
                                 bool rewrite_scalar_read_to_vector_shuffle = true) {
  VectorTypeAccessChecker checker(f->params, allow_untyped_pointers,
                                  rewrite_scalar_read_to_vector_shuffle);
  checker(f->body);

  VectorTypeRewriter rewriter(checker.info_map_, checker.buffer_aliases_, rewrite_buffer_params,
                              rewrite_pointer_params, rewrite_alloc_buffer_node, rewrite_indices,
                              rewrite_let_node, rewrite_scalar_read_to_vector_shuffle);
  PrimFuncNode* n = f.CopyOnWrite();
  n->body = rewriter(std::move(n->body));
  rewriter.Finalize(&f);

  return f;
}

namespace transform {

Pass StorageRewrite() {
  auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
    bool enable_reuse = true;
    bool reuse_require_exact_matched_dtype = false;
    bool merge_static_smem = ctx->GetConfig<bool>("tirx.merge_static_smem", false).value();
    if (merge_static_smem) {
      // When `merge_static_smem` is true, we will reuse and merge shared
      // memory in a dedicated pass `MergeSharedMemoryAllocations`.
      // And so we don't enable reuse in this pass.
      enable_reuse = false;
    }

    ffi::Optional<Target> target = f->GetAttr<Target>("target");
    if (target.has_value() &&
        (target.value()->kind->name == "vulkan" || target.value()->kind->name == "webgpu")) {
      // Require exactly same-dtype matching in smem reuse for Vulkan and WebGPU
      reuse_require_exact_matched_dtype = true;
    }
    auto* n = f.CopyOnWrite();
    n->body = StoragePlanRewriter().Rewrite(std::move(n->body), n->params, true, enable_reuse,
                                            reuse_require_exact_matched_dtype);
    // Parameters may not be rewritten, but internal allocations may.
    return PointerValueTypeRewrite(std::move(f), true, false, false, true, true, true, false);
  };
  return CreatePrimFuncPass(pass_func, 0, "tirx.StorageRewrite", {});
}

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

Pass PointerValueTypeRewrite() {
  auto pass_func = [](PrimFunc f, IRModule m, PassContext ctx) {
    return PointerValueTypeRewrite(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "tirx.PointerValueTypeRewrite", {});
}

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

}  // namespace transform

}  // namespace tirx
}  // namespace tvm
