src/tirx/transform/ir_utils.cc

/*
 * 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 ir_utils.cc
 * \brief Helper functions to construct and compose IR nodes.
 */
#include "ir_utils.h"

#include <tvm/arith/analyzer.h>
#include <tvm/arith/int_solver.h>
#include <tvm/ffi/cast.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/scope_stack.h>
#include <tvm/s_tir/stmt.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>

#include <unordered_map>
#include <unordered_set>
#include <utility>

namespace tvm {
namespace tirx {

Stmt MergeNest(const std::vector<Stmt>& nest, Stmt body) {
  // use reverse iteration
  for (auto ri = nest.rbegin(); ri != nest.rend(); ++ri) {
    Stmt s = *ri;
    if (const auto* for_ = s.as<ForNode>()) {
      auto n = ffi::make_object<ForNode>(*for_);
      TVM_FFI_ICHECK(is_no_op(n->body));
      n->body = body;
      body = Stmt(n);
    } else if (const auto* bind = s.as<BindNode>()) {
      // Bind has no body -- prepend it before the accumulated body in a SeqStmt.
      body = SeqStmt::Flatten(ffi::GetRef<Stmt>(bind), body);
    } else if (const auto* attr = s.as<AttrStmtNode>()) {
      auto n = ffi::make_object<AttrStmtNode>(*attr);
      TVM_FFI_ICHECK(is_no_op(n->body));
      n->body = body;
      body = Stmt(n);
    } else if (const auto* ite = s.as<IfThenElseNode>()) {
      auto n = ffi::make_object<IfThenElseNode>(*ite);
      TVM_FFI_ICHECK(is_no_op(n->then_case));
      TVM_FFI_ICHECK(!n->else_case);
      n->then_case = body;
      body = Stmt(n);
    } else if (const auto* seq = s.as<SeqStmtNode>()) {
      auto n = ffi::make_object<SeqStmtNode>(*seq);
      TVM_FFI_ICHECK(n->size() != 0 && is_no_op(n->seq[n->size() - 1]));
      n->seq.Set(n->size() - 1, body);
      body = Stmt(n);
    } else if (s.as<AssertStmtNode>()) {
      body = SeqStmt({s, body});
    } else if (s.as<AllocBufferNode>() || s.as<DeclBufferNode>()) {
      body = SeqStmt::Flatten(s, body);
    } else {
      TVM_FFI_THROW(InternalError) << "not supported nest type";
    }
  }
  return body;
}

Stmt MergeNest(const std::vector<std::vector<Stmt>>& nest, Stmt body) {
  for (auto ri = nest.rbegin(); ri != nest.rend(); ++ri) {
    body = MergeNest(*ri, body);
  }
  return body;
}

class IRConvertSSA final : public StmtExprMutator {
 public:
  PrimFunc VisitPrimFunc(PrimFunc func) {
    // Remap parameters, if they were used in another function.
    // Function-scope remaps use function_scope_var_remap_ (not the scope stack),
    // because they persist across the entire function body.
    auto params = func->params.Map([&](const tirx::Var& var) -> tirx::Var {
      if (defined_.count(var.get())) {
        Var new_var = MakeNewVar(var);
        PushVarRemap(var, new_var);
        return new_var;
      } else {
        defined_.insert(var.get());
        return var;
      }
    });

    // Remap implicitly defined buffer parameters
    {
      std::unordered_set<const VarNode*> defined_params;
      for (const auto& var : func->params) {
        defined_params.insert(var.get());
      }
      std::unordered_set<const VarNode*> defined_match_vars;
      for (const Var& param : func->params) {
        auto buffer = param.as<BufferVar>();
        if (!buffer) continue;
        auto check_var = [&](const Var& var) {
          const VarNode* var_ptr = var.get();
          if (defined_params.count(var_ptr)) return;
          if (!defined_match_vars.insert(var_ptr).second) return;

          // Buffer-parameter shape vars use "match" semantics: first occurrence
          // defines the var, subsequent occurrences (in other buffers) are
          // just consistent uses of the same var -- not redefinitions.
          if (defined_.count(var_ptr)) {
            Var new_var = MakeNewVar(var);
            PushVarRemap(var, new_var);
          } else {
            defined_.insert(var_ptr);
          }
        };
        for (const auto& dim : buffer.value()->shape) {
          PostOrderVisit(dim, [&](const ffi::ObjectRef& obj) {
            if (auto var = obj.as<Var>()) check_var(var.value());
          });
        }
        for (const auto& stride : buffer.value()->strides) {
          if (auto var = stride.as<Var>()) check_var(var.value());
        }
        if (auto var = buffer.value()->elem_offset.as<Var>()) check_var(var.value());
      }
    }

    // Update the buffer parameters, based on the redefined parameters
    bool buffer_params_changed = false;
    for (size_t i = 0; i < func->params.size(); ++i) {
      if (auto buffer = func->params[i].as<BufferVar>()) {
        BufferVar new_buffer = GetRemappedBuffer(buffer.value());
        if (!new_buffer.same_as(buffer.value()) || !params[i].same_as(new_buffer)) {
          buffer_params_changed = true;
          params.Set(i, new_buffer.var());
        }
      }
    }

    auto attrs = [&]() -> DictAttrs {
      ffi::Map<ffi::String, ffi::Any> dict;
      bool made_change = false;

      for (const auto& [key, old_value] : func->attrs->dict) {
        auto value = old_value;
        if (auto expr = value.as<PrimExpr>()) {
          value = VisitPrimExpr(expr.value());
        } else if (auto* stmt = value.as<StmtNode>()) {
          value = VisitStmt(ffi::GetRef<Stmt>(stmt));
        }

        made_change = made_change || !value.same_as(old_value);
        dict.Set(key, value);
      }

      if (made_change) {
        return DictAttrs(dict);
      } else {
        return func->attrs;
      }
    }();

    auto body = VisitStmt(func->body);

    // If anything changed, update the returned function
    if (!params.same_as(func->params) || buffer_params_changed || !attrs.same_as(func->attrs) ||
        !body.same_as(func->body)) {
      func = PrimFunc(params, body, func->ret_type, attrs);
    }

    // Pop function-scope remaps in reverse order
    PopAllRemapsInCurrentScope();
    function_scope_var_remap_.clear();
    return func;
  }

  // Do not use the base VisitBufferDef for buffer remapping.
  //
  // IRConvertSSA has its own scoped buffer remapping via GetRemappedBuffer and
  // buf_remap_, which handles SSA conversion of buffer data vars, shape, strides,
  // and elem_offset with proper scope tracking. The base StmtMutator::VisitBufferDef
  // would create a conflicting second remap (into base buffer_remap_) when called
  // from the default DeclBuffer/AllocBuffer handlers, producing buffers with
  // undefined SSA-renamed variables.
  BufferVar VisitBufferDef(const BufferVar& buffer, bool alloc_data) override { return buffer; }

  Expr VisitExpr_(const VarNode* op) final { return GetRemappedVar(ffi::GetRef<Var>(op)); }
  Expr VisitExpr_(const LetNode* op) final {
    const Var& v = op->var;
    if (defined_.count(v.get())) {
      PrimExpr value = this->VisitPrimExpr(op->value);
      Var new_var = MakeNewVar(v);
      PushVarRemap(v, new_var);
      PrimExpr body = this->VisitPrimExpr(op->body);
      PopVarRemap(v, new_var);
      return Let(new_var, value, body);
    } else {
      defined_.insert(v.get());
      return StmtExprMutator::VisitExpr_(op);
    }
  }

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

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

  Stmt VisitStmt_(const DeclBufferNode* op) final {
    Var v = op->buffer.var();
    if (defined_.count(v.get())) {
      Var new_var = MakeNewVar(v);
      PushVarRemap(v, new_var);
    } else {
      defined_.insert(v.get());
    }
    DeclBuffer decl = StmtExprMutator::VisitStmt_(op).as_or_throw<DeclBuffer>();
    BufferVar new_buffer = GetRemappedBuffer(decl->buffer);
    if (!new_buffer.same_as(decl->buffer)) {
      decl.CopyOnWrite()->buffer = std::move(new_buffer);
    }
    return decl;
  }

  Stmt VisitStmt_(const SBlockNode* op) final {
    SBlock block = ffi::GetRef<SBlock>(op);

    // The SBlockNode is the point of definition for the IterVar
    // instances.  These re-defines must be present before visiting
    // the body of the SBlockNode.
    return scope_.WithNewScope([&]() -> Stmt {
      ffi::Array<IterVar> iter_vars = op->iter_vars.Map([&](IterVar iter_var) {
        if (defined_.count(iter_var->var.get())) {
          Var new_var = MakeNewVar(iter_var->var);
          PushVarRemap(iter_var->var, new_var);
          iter_var.CopyOnWrite()->var = new_var.as_or_throw<PrimVar>();
        } else {
          defined_.insert(iter_var->var.get());
        }
        return iter_var;
      });
      ffi::Array<BufferRegion> reads =
          block->reads.Map([&](const auto& region) { return VisitBufferAccess(region); });
      ffi::Array<BufferRegion> writes =
          block->writes.Map([&](const auto& region) { return VisitBufferAccess(region); });

      if (!reads.same_as(block->reads) || !writes.same_as(block->writes) ||
          !iter_vars.same_as(op->iter_vars)) {
        auto write_ptr = block.CopyOnWrite();
        write_ptr->reads = reads;
        write_ptr->writes = writes;
        write_ptr->iter_vars = iter_vars;
      }

      return StmtExprMutator::VisitStmt_(block.get()).as_or_throw<SBlock>();
    });
  }

  template <typename Node>
  Node VisitBufferAccess(Node node) {
    BufferVar new_buf = GetRemappedBuffer(node->buffer);
    if (!new_buf.same_as(node->buffer)) {
      auto writer = node.CopyOnWrite();
      writer->buffer = new_buf;
    }

    return node;
  }

  Var GetRemappedVar(Var var) {
    if (auto it = var_remap_.find(var.get()); it != var_remap_.end() && it->second.size()) {
      return it->second.back();
    } else if (auto it = function_scope_var_remap_.find(var.get());
               it != function_scope_var_remap_.end()) {
      return it->second;
    } else {
      return var;
    }
  }

  BufferVar GetRemappedBuffer(BufferVar buf) {
    // Determine the buffer var that should be in the updated buffer,
    // given the current scope.  If no redefines are present, then the
    // buffer var is unchanged.
    Var new_buffer_var = GetRemappedVar(buf.var());
    PrimExpr elem_offset = VisitPrimExpr(buf->elem_offset);
    auto visit_expr = [this](const PrimExpr& expr) { return VisitPrimExpr(expr); };
    ffi::Array<PrimExpr> shape = buf->shape.Map(visit_expr);
    ffi::Array<PrimExpr> strides = buf->strides.Map(visit_expr);

    // Rewrite the layout's per-iter extent/stride expressions in lockstep
    // with the shape. If we don't, SSA-renamed shape vars end up as fresh
    // Vars while the layout still references the original, producing
    // structurally-unequal buffers whose shape and layout disagree (e.g.,
    // test_dynamic_launch_thread).
    ffi::Optional<Layout> new_layout = buf->layout;
    bool layout_changed = false;
    if (buf->layout.has_value()) {
      if (auto opt_tile = buf->layout.value().as<TileLayoutNode>()) {
        auto remap_iter = [&](const Iter& it) -> Iter {
          PrimExpr new_extent = VisitPrimExpr(it->extent);
          PrimExpr new_stride = VisitPrimExpr(it->stride);
          if (new_extent.same_as(it->extent) && new_stride.same_as(it->stride)) {
            return it;
          }
          return Iter(new_extent, new_stride, it->axis);
        };
        auto new_shard = opt_tile->shard.Map(remap_iter);
        auto new_replica = opt_tile->replica.Map(remap_iter);
        if (!new_shard.same_as(opt_tile->shard) || !new_replica.same_as(opt_tile->replica)) {
          new_layout = TileLayout(new_shard, new_replica, opt_tile->offset);
          layout_changed = true;
        }
      }
    }

    // If no mapping is required, return the original buffer.
    if (new_buffer_var.same_as(buf.var()) && elem_offset.same_as(buf->elem_offset) &&
        shape.same_as(buf->shape) && strides.same_as(buf->strides) && !layout_changed) {
      return buf;
    }

    // If the current scope already has a mapping of this buffer, use
    // the mapped buffer.
    auto key = buf.get();
    std::vector<BufferVar>& buffers = buf_remap_[key];
    if (buffers.size() && buffers.back().same_as(new_buffer_var)) {
      return buffers.back();
    }

    // When only the buffer's identity changed, the remapped Var already has
    // the desired BufferType.  Reuse that exact Var so the definition and all
    // subsequent uses remain in SSA.
    if (const auto* type = new_buffer_var->ty.as<BufferTypeNode>()) {
      BufferVar candidate(new_buffer_var);
      if (shape.same_as(type->shape) && strides.same_as(type->strides) &&
          elem_offset.same_as(type->elem_offset) && !layout_changed) {
        buffers.push_back(candidate);
        return candidate;
      }
    }

    // Otherwise, make and return a new buffer object that uses the
    // new buffer, pushing it onto the scoped stack of existing
    // buffers.  This will be popped when the new_buffer_var
    // redefinition is popped.
    auto type = CopyBufferType(buf);
    type->shape = shape;
    type->strides = strides;
    type->elem_offset = elem_offset;
    if (layout_changed) {
      type->layout = std::move(new_layout);
    }
    BufferVar new_buf = RebuildBufferVar(buf, std::move(type), new_buffer_var->name);

    // A BufferVar's metadata lives in its Var type.  If rewriting the
    // metadata required a fresh Var, make it the active remap as well.  This
    // keeps BufferLoad/BufferStore and ordinary Var uses (such as
    // buffer_data) on the same identity.
    auto it = var_remap_.find(buf.get());
    if (it != var_remap_.end() && it->second.size() && it->second.back().same_as(new_buffer_var)) {
      it->second.back() = new_buf.var();
    } else if (auto function_it = function_scope_var_remap_.find(buf.get());
               function_it != function_scope_var_remap_.end() &&
               function_it->second.same_as(new_buffer_var)) {
      function_it->second = new_buf.var();
    } else {
      PushVarRemap(buf.var(), new_buf.var());
    }
    buffers.push_back(new_buf);
    return new_buf;
  }

  Stmt VisitStmt_(const BindNode* op) final {
    // Bind var remaps are tracked in the current scope so they persist
    // across SeqStmt siblings and are cleaned up when the enclosing
    // body-carrying statement's scope exits.
    const Var& v = op->var;
    if (defined_.count(v.get())) {
      Expr value = this->VisitExpr(op->value);
      Var new_var = MakeNewVar(v);
      PushVarRemap(v, new_var);
      return Bind(new_var, value);
    } else {
      defined_.insert(v.get());
      return StmtExprMutator::VisitStmt_(op);
    }
  }

  Stmt VisitStmt_(const IfThenElseNode* op) final {
    // Each branch gets its own scope so Bind remaps in one branch
    // do not leak into the other.
    PrimExpr condition = VisitPrimExpr(op->condition);
    Stmt then_case = scope_.WithNewScope([&]() -> Stmt { return VisitStmt(op->then_case); });
    ffi::Optional<Stmt> else_case;
    if (op->else_case) {
      else_case = scope_.WithNewScope([&]() -> Stmt { return VisitStmt(op->else_case.value()); });
    }
    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);
  }

  Stmt VisitStmt_(const ForNode* op) final {
    const Var& v = op->loop_var;
    if (defined_.count(v.get())) {
      return scope_.WithNewScope([&]() -> Stmt {
        Var new_var = MakeNewVar(v);
        PushVarRemap(v, new_var);
        Stmt stmt = StmtExprMutator::VisitStmt_(op);
        auto n = ffi::make_object<ForNode>(*stmt.as<ForNode>());
        n->loop_var = new_var.as_or_throw<PrimVar>();
        return For(n);
      });
    } else {
      defined_.insert(v.get());
      return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::VisitStmt_(op); });
    }
  }
  Stmt VisitStmt_(const WhileNode* op) final {
    return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::VisitStmt_(op); });
  }
  Stmt VisitStmt_(const AllocBufferNode* op) final {
    Var v = op->buffer.var();
    if (defined_.count(v.get())) {
      Var new_var = MakeNewVar(v);
      PushVarRemap(v, new_var);
    } else {
      defined_.insert(v.get());
    }
    Stmt stmt = StmtExprMutator::VisitStmt_(op);
    op = stmt.as<AllocBufferNode>();
    // Use GetRemappedBuffer so that the AllocBuffer's buffer is the same
    // object as the one used by BufferStore/BufferLoad in subsequent siblings.
    BufferVar new_buf = GetRemappedBuffer(op->buffer);
    if (!new_buf.same_as(op->buffer)) {
      auto node = stmt.as_or_throw<AllocBuffer>();
      node.CopyOnWrite()->buffer = std::move(new_buf);
      return node;
    }
    return stmt;
  }
  Stmt VisitStmt_(const AttrStmtNode* op) final {
    if (const IterVarNode* iter_var = op->node.as<IterVarNode>()) {
      Range dom = iter_var->dom;
      if (dom.defined()) {
        auto min = VisitPrimExpr(dom->min);
        auto extent = VisitPrimExpr(dom->extent);
        if (!min.same_as(iter_var->dom->min) || !extent.same_as(iter_var->dom->extent)) {
          dom = Range::FromMinExtent(min, extent);
        }
      }

      Var var = iter_var->var;
      bool delayed_define = false;
      if (auto it = function_scope_var_remap_.find(var.get());
          it != function_scope_var_remap_.end()) {
        var = it->second;
      } else if (defined_.count(var.get())) {
        Var new_var(var->name, var->ty);

        function_scope_var_remap_.insert({var.get(), new_var});
        var = new_var;
      } else {
        // The AttrStmt refers to an undefined variable.  This is
        // allowed for some attributes, such as
        // "pragma_parallel_launch_point", which annotates a variable
        // that is about to occur in a ForNode.  In these cases, the
        // ForNode and the AttrStmt must continue using the same
        // variable defintion.
        //
        // However, other AttrStmt, such as "thread_extent", act as
        // points of definition for the variable they annotate.  If
        // the variable has not been defined after visiting the body,
        // we should mark it as defined before exiting.  This ensures
        // correct de-duplication between multiple functions.
        //
        // This implementation may be simplified in the future by
        // moving "pragma_parallel_launch_point" to be an annotation
        // on the `ForNode`, rather than an `AttrStmt`.
        delayed_define = true;
      }

      IterVar new_iter_var;
      if (dom.same_as(iter_var->dom) && var.same_as(iter_var->var)) {
        new_iter_var = ffi::GetRef<IterVar>(iter_var);
      } else {
        new_iter_var = IterVar(dom, var.as_or_throw<PrimVar>(), iter_var->iter_type,
                               iter_var->thread_tag, iter_var->span);
      }

      auto value = VisitPrimExpr(op->value);
      auto body = scope_.WithNewScope([&]() -> Stmt { return VisitStmt(op->body); });

      Stmt output;
      if (new_iter_var.get() == iter_var && body.same_as(op->body) && value.same_as(op->value)) {
        output = ffi::GetRef<Stmt>(op);
      } else {
        output = AttrStmt(new_iter_var, op->attr_key, value, body, iter_var->span);
      }

      if (delayed_define) {
        if (!defined_.count(var.get())) {
          function_scope_var_remap_.insert({var.get(), var});
          defined_.insert(var.get());
        }
      }

      return output;

    } else if (const VarNode* v = op->node.as<VarNode>()) {
      Stmt stmt = scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::VisitStmt_(op); });
      op = stmt.as<AttrStmtNode>();
      if (var_remap_.count(v) && var_remap_[v].size() != 0) {
        return AttrStmt(var_remap_[v].back(), op->attr_key, op->value, op->body);
      } else {
        return stmt;
      }
    } else {
      return scope_.WithNewScope([&]() -> Stmt { return StmtExprMutator::VisitStmt_(op); });
    }
  }

 private:
  /*! \brief Record of a variable remap pushed to the current scope. */
  struct VarRemap {
    Var old_var;
    Var new_var;
  };

  /*! \brief Check whether a buffer uses a variable in any remapped field. */
  static bool BufferDependsOnVar(const BufferVar& buffer, const VarNode* var) {
    if (buffer.get() == var) return true;

    auto uses_var = [var](const PrimExpr& expr) {
      return expr.defined() && UsesVar(expr, [var](const VarNode* node) { return node == var; });
    };
    if (uses_var(buffer->elem_offset)) return true;
    for (const PrimExpr& dim : buffer->shape) {
      if (uses_var(dim)) return true;
    }
    for (const PrimExpr& stride : buffer->strides) {
      if (uses_var(stride)) return true;
    }
    if (buffer->layout.has_value()) {
      if (const auto* tile_layout = buffer->layout.value().as<TileLayoutNode>()) {
        for (const Iter& iter : tile_layout->shard) {
          if (uses_var(iter->extent) || uses_var(iter->stride)) return true;
        }
        for (const Iter& iter : tile_layout->replica) {
          if (uses_var(iter->extent) || uses_var(iter->stride)) return true;
        }
      }
    }
    return false;
  }

  /*! \brief Create a new variable with the same name and type as the original. */
  static Var MakeNewVar(const Var& old_var) { return Var(old_var->name, old_var->ty); }

  /*! \brief Push a variable remap to the current scope and the var_remap_ stack. */
  void PushVarRemap(const Var& old_var, const Var& new_var) {
    var_remap_[old_var.get()].push_back(new_var);
    auto& level = scope_.Current();
    level.parent = this;
    level.push_back({old_var, new_var});
  }

  /*! \brief Pop a single variable remap (used for expression-level Let scoping). */
  void PopVarRemap(const Var& old_var, const Var& new_var) {
    var_remap_[old_var.get()].pop_back();
    for (auto& kv : buf_remap_) {
      std::vector<BufferVar>& buffers = kv.second;
      if (buffers.size() && BufferDependsOnVar(buffers.back(), new_var.get())) {
        buffers.pop_back();
      }
    }
    // Also remove from the current scope's tracking vector
    auto& current = scope_.Current();
    if (current.size() && current.back().new_var.same_as(new_var)) {
      current.pop_back();
    }
  }

  /*! \brief Pop all remaps in the current scope level (used for function-scope cleanup). */
  void PopAllRemapsInCurrentScope() {
    auto& current = scope_.Current();
    while (current.size()) {
      auto& remap = current.back();
      var_remap_[remap.old_var.get()].pop_back();
      for (auto& kv : buf_remap_) {
        std::vector<BufferVar>& buffers = kv.second;
        if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) {
          buffers.pop_back();
        }
      }
      current.pop_back();
    }
  }

  /*! \brief Scope stack: each scope level holds the remaps introduced in that scope.
   *
   * When a body-carrying statement (For, SBlock, Allocate) calls
   * scope_.WithNewScope([&]{...}), a new scope level is pushed.
   * Bind statements push their remaps to the current scope.
   * On scope exit, the destructor of std::vector<VarRemap> triggers,
   * and we undo all remaps in that level.
   *
   * Note: ScopeStack<T>::WithNewScope calls T's destructor on exit.
   * std::vector's destructor destroys elements but does NOT call custom
   * cleanup.  So we wrap the vector in ScopeLevel which handles cleanup.
   */
  struct ScopeLevel {
    std::vector<VarRemap> remaps;
    IRConvertSSA* parent{nullptr};

    void push_back(VarRemap remap) { remaps.push_back(std::move(remap)); }
    size_t size() const { return remaps.size(); }
    VarRemap& back() { return remaps.back(); }
    void pop_back() { remaps.pop_back(); }

    ~ScopeLevel() {
      if (!parent) return;
      // Pop remaps in reverse order
      while (remaps.size()) {
        auto& remap = remaps.back();
        parent->var_remap_[remap.old_var.get()].pop_back();
        for (auto& kv : parent->buf_remap_) {
          std::vector<BufferVar>& buffers = kv.second;
          if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) {
            buffers.pop_back();
          }
        }
        remaps.pop_back();
      }
    }

    ScopeLevel() = default;
    ScopeLevel(const ScopeLevel&) = delete;
    ScopeLevel& operator=(const ScopeLevel&) = delete;
    ScopeLevel(ScopeLevel&& other) noexcept
        : remaps(std::move(other.remaps)), parent(other.parent) {
      other.parent = nullptr;  // prevent other's destructor from popping
    }
    ScopeLevel& operator=(ScopeLevel&& other) noexcept {
      if (this != &other) {
        // Run our destructor logic first
        if (parent) {
          while (remaps.size()) {
            auto& remap = remaps.back();
            parent->var_remap_[remap.old_var.get()].pop_back();
            for (auto& kv : parent->buf_remap_) {
              std::vector<BufferVar>& buffers = kv.second;
              if (buffers.size() && BufferDependsOnVar(buffers.back(), remap.new_var.get())) {
                buffers.pop_back();
              }
            }
            remaps.pop_back();
          }
        }
        remaps = std::move(other.remaps);
        parent = other.parent;
        other.parent = nullptr;
      }
      return *this;
    }
  };

  std::unordered_map<const VarNode*, std::vector<Var>> var_remap_;
  std::unordered_set<const VarNode*> defined_;
  std::unordered_map<const VarNode*, std::vector<BufferVar>> buf_remap_;
  std::unordered_map<const VarNode*, Var> function_scope_var_remap_;
  ScopeStack<ScopeLevel> scope_;
};

Stmt ConvertSSA(Stmt stmt) { return IRConvertSSA()(std::move(stmt)); }

ffi::String GetPtrStorageScope(Var buffer_var) {
  if (const auto* buffer_type = buffer_var->ty.as<BufferTypeNode>()) {
    return buffer_type->storage_scope;
  }
  const auto* ptr_type = buffer_var->ty.as<PointerTypeNode>();
  TVM_FFI_ICHECK(ptr_type)
      << "The provided variable is neither a pointer nor a buffer-typed variable";
  return ptr_type->storage_scope;
}

ffi::Array<PrimExpr> GetBufferAllocationShape(const BufferVar& buffer) {
  ffi::Array<PrimExpr> alloc_shape = buffer->shape;
  if (buffer->strides.size()) {
    TVM_FFI_ICHECK_EQ(buffer->shape.size(), buffer->strides.size());
    for (size_t i = buffer->strides.size() - 1; i > 0; --i) {
      TVM_FFI_ICHECK(arith::Analyzer()->CanProveEqual(
          floormod(buffer->strides[i - 1], buffer->strides[i]), 0));
      alloc_shape.Set(i, buffer->strides[i - 1] / buffer->strides[i]);
    }
  }
  return alloc_shape;
}

ffi::Array<PrimExpr> ConvertIndices(const MatchBufferRegion& match_buffer,
                                    const ffi::Array<PrimExpr>& indices) {
  const BufferVar& target = match_buffer->buffer;
  const BufferRegion& source = match_buffer->source;
  TVM_FFI_ICHECK_EQ(indices.size(), target->shape.size());

  arith::Analyzer analyzer;
  ffi::Array<PrimExpr> result;
  result.reserve(source->region.size());
  size_t offset = source->region.size() - indices.size();
  for (size_t i = 0; i < offset; ++i) {
    const Range& range = source->region[i];
    TVM_FFI_ICHECK(analyzer->CanProve(range->extent == 1));
    result.push_back(range->min);
  }
  for (size_t i = 0; i < indices.size(); ++i) {
    const Range& range = source->region[i + offset];
    const PrimExpr& index = indices[i];
    result.push_back(range->min + index);
  }
  return result;
}

Region ConvertRegion(const MatchBufferRegion& match_buffer, const Region& region) {
  const BufferVar& target = match_buffer->buffer;
  const BufferRegion& source = match_buffer->source;
  TVM_FFI_ICHECK_EQ(region.size(), target->shape.size());

  arith::Analyzer analyzer;
  Region result;
  result.reserve(source->region.size());
  size_t offset = source->region.size() - region.size();
  for (size_t i = 0; i < offset; ++i) {
    const Range& source_range = source->region[i];
    TVM_FFI_ICHECK(analyzer->CanProve(source_range->extent == 1));
    result.push_back(Range::FromMinExtent(source_range->min, 1));
  }
  for (size_t i = 0; i < region.size(); ++i) {
    const Range& source_range = source->region[i + offset];
    const Range& target_range = region[i];
    result.push_back(
        Range::FromMinExtent(source_range->min + target_range->min, target_range->extent));
  }
  return result;
}

ffi::Optional<arith::IntConstraints> ConditionalBoundsContext::TrySolveCondition() {
  // extract equations and related vars from condition expression.
  // currently only extract simple integral equations which could be solvable.
  arith::Analyzer analyzer;
  PrimExpr condition = analyzer->Simplify(condition_);
  if (is_const_int(condition)) {
    return std::nullopt;
  }
  ffi::Array<PrimExpr> equations;
  ffi::Array<PrimVar> vars;
  std::function<void(const PrimExpr&)> fvisit = [&equations, &vars, &fvisit](const PrimExpr& e) {
    if (e->IsInstance<GENode>() || e->IsInstance<GTNode>() || e->IsInstance<LENode>() ||
        e->IsInstance<LTNode>() || e->IsInstance<EQNode>() || e->IsInstance<NENode>()) {
      bool is_simple = true;
      std::vector<PrimVar> cand_vars;
      PostOrderVisit(e, [&cand_vars, &is_simple, &e](const ffi::ObjectRef& obj) {
        if (obj.same_as(e)) {
          return;
        } else if (const VarNode* var = obj.as<VarNode>()) {
          PrimType var_ty = var->ty.as_or_throw<PrimType>();
          if (var_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
            cand_vars.push_back(ffi::GetRef<Var>(var).as_or_throw<PrimVar>());
          }
        } else {
          is_simple &= obj->IsInstance<AddNode>() || obj->IsInstance<SubNode>() ||
                       obj->IsInstance<MulNode>() || obj->IsInstance<FloorDivNode>() ||
                       obj->IsInstance<FloorModNode>() || obj->IsInstance<IntImmNode>();
        }
      });
      if (is_simple && !cand_vars.empty()) {
        for (const PrimVar& new_var : cand_vars) {
          if (!std::any_of(vars.begin(), vars.end(),
                           [&new_var](const PrimVar& v) { return v.same_as(new_var); })) {
            vars.push_back(new_var);
          }
        }
        equations.push_back(e.as_or_throw<PrimExpr>());
      }
    } else if (e->IsInstance<AndNode>()) {
      And op = e.as_or_throw<And>();
      fvisit(op->a);
      fvisit(op->b);
    } else if (e->IsInstance<CallNode>()) {
      Call op = e.as_or_throw<Call>();
      if (op->op.same_as(builtin::likely())) {
        fvisit(op->args[0].as_or_throw<PrimExpr>());
      }
    }
  };
  fvisit(condition);
  if (equations.empty() || vars.empty()) {
    return std::nullopt;
  }
  // build dom ranges for related vars
  ffi::Map<Var, Range> ranges;
  for (const Var& v : vars) {
    arith::IntSet dom;
    auto relax_it = relax_map_->find(v.get());
    if (relax_it != relax_map_->end()) {
      dom = relax_it->second;
    } else {
      auto hint_it = hint_map_->find(v.get());
      if (hint_it != hint_map_->end()) {
        dom = hint_it->second;
      }
    }
    if (dom.defined()) {
      ranges.Set(v, Range::FromMinExtent(dom.min(), analyzer->Simplify(dom.max() - dom.min() + 1)));
    }
  }
  // solve constraints
  arith::IntConstraints constraint(vars, ranges, equations);
  arith::IntConstraints result = arith::SolveInequalitiesToRange(constraint);
  if (!result->relations.empty()) {
    return std::nullopt;
  }
  return result;
}

ConditionalBoundsContext::ConditionalBoundsContext(
    const PrimExpr& condition, std::unordered_map<const VarNode*, arith::IntSet>* relax_map,
    std::unordered_map<const VarNode*, arith::IntSet>* hint_map,
    std::vector<PrimExpr>* pending_conditions)
    : condition_(condition),
      relax_map_(relax_map),
      hint_map_(hint_map),
      pending_conditions_(pending_conditions),
      origin_pending_conditions_num_(pending_conditions->size()) {}

void ConditionalBoundsContext::EnterWithScope() {
  ffi::Optional<arith::IntConstraints> constraints = TrySolveCondition();
  if (!constraints.has_value()) {
    // fail to process the condition, add to unresolved
    pending_conditions_->push_back(condition_);
    return;
  }
  // update solved var ranges
  for (const auto& kv : constraints.value()->ranges) {
    const VarNode* var = kv.first.get();
    arith::IntSet new_dom = arith::IntSet::FromRange(kv.second);
    auto relax_it = relax_map_->find(var);
    if (relax_it != relax_map_->end()) {
      // this is a bound for relaxed var
      origin_map_.emplace(var, relax_it->second);
      relax_it->second = arith::Intersect({relax_it->second, new_dom});
    } else {
      // this is a bound for free var
      auto hint_it = hint_map_->find(var);
      if (hint_it != hint_map_->end()) {
        origin_map_.emplace(var, hint_it->second);
        hint_it->second = arith::Intersect({hint_it->second, new_dom});
      } else {
        origin_map_.emplace(var, arith::IntSet::Nothing());
        hint_map_->insert(hint_it, {var, new_dom});
      }
    }
  }
}

void ConditionalBoundsContext::ExitWithScope() {
  pending_conditions_->resize(origin_pending_conditions_num_);
  for (const auto& p : origin_map_) {
    const auto* var = p.first;
    auto relax_it = relax_map_->find(var);
    if (relax_it != relax_map_->end()) {
      // recover bound for relaxed var
      relax_it->second = p.second;
    } else {
      // recover bound for free var
      auto hint_it = hint_map_->find(var);
      TVM_FFI_ICHECK(hint_it != hint_map_->end());
      if (p.second.IsNothing()) {
        hint_map_->erase(hint_it);
      } else {
        hint_it->second = p.second;
      }
    }
  }
}

std::pair<PrimExpr, PrimExpr> GetAsyncWaitAttributes(const AttrStmtNode* op) {
  TVM_FFI_ICHECK(op && op->attr_key == s_tir::attr::async_wait_queue_scope);
  auto inner = op->body.as<AttrStmtNode>();
  TVM_FFI_ICHECK(inner && inner->attr_key == s_tir::attr::async_wait_inflight_count);
  return std::make_pair(op->value, inner->value);
}

/*! \brief Collect storage alignment information from annotations. */
class StorageAlignCollector : public StmtVisitor {
 private:
  friend std::unordered_map<Var, StorageAlignAnnotation> CollectStorageAlignAnnotation(
      const Stmt& body);

  /*! \brief For s-stir, the alignment annotations reside in block annotations. */
  void VisitStmt_(const SBlockNode* op) final {
    auto it = op->annotations.find(s_tir::attr::buffer_dim_align);
    if (it != op->annotations.end()) {
      auto storage_align_annotation = (*it).second.as_or_throw<StorageAlignAnnotation>();
      for (const auto& storage_align_tuple : storage_align_annotation) {
        int buffer_index = storage_align_tuple.get<0>();
        const BufferVar& buffer = op->writes[buffer_index]->buffer;
        storage_align_[buffer.var()].push_back(storage_align_tuple);
      }
    }
    StmtVisitor::VisitStmt_(op);
  }

  /*! \brief AllocBuffer: check for buffer_dim_align annotations. */
  void VisitStmt_(const AllocBufferNode* op) final {
    auto it = op->annotations.find(s_tir::attr::buffer_dim_align);
    if (it != op->annotations.end()) {
      auto storage_align_annotation = (*it).second.as_or_throw<StorageAlignAnnotation>();
      for (const auto& storage_align_tuple : storage_align_annotation) {
        int buffer_index = storage_align_tuple.get<0>();
        // the first buffer idx info is meaningless for alloc
        // stmt and should set as negative intentionally.
        TVM_FFI_ICHECK_EQ(buffer_index, -1);
        storage_align_[op->buffer.var()].push_back(storage_align_tuple);
      }
    }
    StmtVisitor::VisitStmt_(op);
  }

  /*! \brief The map from buffer var to its storage alignment information. */
  std::unordered_map<Var, StorageAlignAnnotation> storage_align_;
};

std::unordered_map<Var, StorageAlignAnnotation> CollectStorageAlignAnnotation(const Stmt& body) {
  StorageAlignCollector collector;
  collector(body);
  return std::move(collector.storage_align_);
}

int Stoi(const std::string& str) {
  try {
    return std::stoi(str);
  } catch (std::invalid_argument& e) {
    TVM_FFI_THROW(InternalError) << "Cannot convert \"" << str << "\" to int";
    throw;
  }
}

std::pair<int32_t, int32_t> GetWmmaFragmentDimSize(const std::string& shape_str,
                                                   const std::string& scope) {
  size_t m, n, k;
  size_t last_pos = 0, pos = 0;
  pos = shape_str.find(", ", last_pos);
  m = Stoi(shape_str.substr(last_pos, pos - last_pos));
  last_pos = pos + 2;
  pos = shape_str.find(", ", last_pos);
  n = Stoi(shape_str.substr(last_pos, pos - last_pos));
  last_pos = pos + 2;
  k = Stoi(shape_str.substr(last_pos, shape_str.length() - last_pos));
  if (scope == "wmma.matrix_a") {
    return std::pair<int32_t, int32_t>(m, k);
  } else if (scope == "wmma.matrix_b") {
    return std::pair<int32_t, int32_t>(k, n);
  } else if (scope == "wmma.accumulator") {
    return std::pair<int32_t, int32_t>(m, n);
  }
  return std::pair<int32_t, int32_t>(0, 0);
}

std::optional<bool> IsHostFunc(const PrimFunc& func) {
  if (func->HasNonzeroAttr(tvm::tirx::attr::kIsHostFunc)) {
    return true;
  } else if (auto target = func->GetAttr<Target>(tvm::attr::kTarget)) {
    return target.value()->HasKey("cpu");
  } else {
    return std::nullopt;
  }
}

namespace transform {
Pass ConvertSSA() {
  auto pass_func = [](IRModule mod, PassContext ctx) {
    tirx::IRConvertSSA converter;
    ffi::Map<GlobalVar, BaseFunc> functions;
    bool made_change = false;
    for (auto [gvar, base_func] : mod->functions) {
      if (auto* ptr = base_func.as<tirx::PrimFuncNode>()) {
        auto updated = converter.VisitPrimFunc(ffi::GetRef<tirx::PrimFunc>(ptr));
        if (!updated.same_as(base_func)) {
          made_change = true;
          base_func = updated;
        }
      }
      functions.Set(gvar, base_func);
    }
    if (made_change) {
      mod.CopyOnWrite()->functions = std::move(functions);
    }
    return mod;
  };
  return tvm::transform::CreateModulePass(pass_func, 0, "tirx.ConvertSSA", {});
}

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

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