/*!
 * \file op/parallel.cc
 * \brief Define Parallel for operator
 */

#include "parallel.h"
#include "support/check.h"
#include <tvm/ffi/extra/structural_equal.h>
#include <tvm/ir/cast.h>
#include <tvm/runtime/logging.h>

#include <algorithm>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/op.h>

#include "../layout/layout.h"
#include "arith/int_operator.h"

#include "../layout/utils.h"
#include "../transform/loop_partition.h"
#include "../transform/loop_vectorize.h"
#include "backend/common/target_utils.h"
#include "utils.h"

namespace tvm {
namespace tl {

using namespace tirx;
using namespace ffi;

namespace {

class IfBufferRemapLoopGenerator : public StmtExprMutator {
public:
  static For run(Stmt stmt, Map<Buffer, Buffer> buffer_remap,
                 Map<Buffer, Layout> layout_map) {
    IfBufferRemapLoopGenerator generator(buffer_remap, layout_map);
    return Downcast<For>(generator(std::move(stmt)));
  }

private:
  IfBufferRemapLoopGenerator(Map<Buffer, Buffer> buffer_remap,
                             Map<Buffer, Layout> layout_map)
      : buffer_remap_(buffer_remap), layout_map_(layout_map) {}

  PrimExpr VisitExpr_(const BufferLoadNode *op) final {
    auto load = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));

    if (buffer_remap_.count(load->buffer)) {
      auto new_indices = layout_map_[load->buffer]->Forward(load->indices);
      auto new_buffer = buffer_remap_[load->buffer];

      return BufferLoad(new_buffer, new_indices);
    }
    return load;
  }

  Stmt VisitStmt_(const BufferStoreNode *op) final {
    auto store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
    if (buffer_remap_.count(store->buffer)) {
      auto new_indices = layout_map_[store->buffer]->Forward(store->indices);
      auto new_buffer = buffer_remap_[store->buffer];
      return BufferStore(new_buffer, store->value, new_indices);
    }
    return store;
  }

  Map<Buffer, Buffer> buffer_remap_;
  Map<Buffer, Layout> layout_map_;
};

// Choose a vector width for non-fragment SIMT loops while keeping the full
// thread extent active.  The loop partitioner covers work in chunks of
// `thread_extent * vector_size`; when the logical loop size is ragged, the
// excess points are guarded as padding.  Prefer the candidate that introduces
// the fewest padded points, and keep the wider vector width on ties.
int SelectMinPaddingVectorSize(int max_vector_size, PrimExpr loop_total_size,
                               PrimExpr thread_extent,
                               arith::Analyzer *analyzer) {
  ICHECK_GT(max_vector_size, 0);
  const int64_t *total = as_const_int(analyzer->Simplify(loop_total_size));
  const int64_t *threads = as_const_int(analyzer->Simplify(thread_extent));
  // This is a compile-time planning heuristic.  If the size is symbolic, keep
  // the vectorizer's original choice and rely on the existing runtime guard.
  if (total == nullptr || threads == nullptr || *threads <= 0) {
    return max_vector_size;
  }

  int best_vector_size = max_vector_size;
  int64_t best_padding = 0;
  bool has_best = false;
  PrimExpr total_i64 = cast(DataType::Int(64), loop_total_size);
  // Candidate vector widths are generated by halving the maximum legal width,
  // matching the rest of the vectorization planner's power-of-two choices.
  for (int candidate = max_vector_size; candidate >= 1; candidate /= 2) {
    PrimExpr partition_width =
        IntImm(DataType::Int(64), (*threads) * static_cast<int64_t>(candidate));
    // Use tirx::ceildiv instead of hand-rolled integer arithmetic so the same
    // simplification and constant-folding rules are used as TIR expressions.
    PrimExpr covered_expr = analyzer->Simplify(
        ceildiv(total_i64, partition_width) * partition_width);
    const int64_t *covered = as_const_int(covered_expr);
    if (covered == nullptr) {
      return max_vector_size;
    }
    int64_t padding = *covered - *total;
    if (!has_best || padding < best_padding ||
        (padding == best_padding && candidate > best_vector_size)) {
      has_best = true;
      best_padding = padding;
      best_vector_size = candidate;
    }
    if (candidate == 1) {
      break;
    }
  }
  return best_vector_size;
}

} // anonymous namespace

/**
 * @brief Handle a parallel For node during traversal, collecting loop metadata.
 *
 * Visits a parallel loop, asserts the loop is parallel, records a data-parallel
 * IterVar for the loop, binds the loop variable range into the analyzer scope,
 * and extracts any reducer information from the loop's annotations into the
 * visitor's reducer_info_map_. Continues traversal into the loop body.
 */
void ParallelLoopNestVisitor::VisitStmt_(const ForNode *op) {
  if (op->kind == ForKind::kParallel)
    p->loop_vars_.push_back(IterVar(Range(op->min, op->extent), op->loop_var,
                                    IterVarType::kDataPar));
  else
    p->inner_vars_.Set(op->loop_var,
                       IterVar(Range(op->min, op->extent), op->loop_var,
                               IterVarType::kOrdered));
  p->analyzer_.Bind(op->loop_var, Range::FromMinExtent(op->min, op->extent));
  if (auto reducer_info_ref = op->annotations.Get(attr::kReducerInfo)) {
    if (auto reducer_info_map =
            reducer_info_ref.value().as<Map<Var, ReducerInfo>>()) {
      for (auto &&[buffer, info] : reducer_info_map.value())
        p->reducer_info_map_.Set(buffer, info);
    }
  }
  StmtExprVisitor::VisitStmt_(op);
}

void ParallelLoopNestVisitor::VisitStmt_(const BufferStoreNode *op) {
  if (IsFragmentBuffer(op->buffer)) {
    p->RecordBufferAccess(op->buffer, op->indices, /*is_write=*/true);
  }
  StmtExprVisitor::VisitStmt_(op);
}

void ParallelLoopNestVisitor::VisitExpr_(const BufferLoadNode *op) {
  if (IsFragmentBuffer(op->buffer)) {
    p->RecordBufferAccess(op->buffer, op->indices, /*is_write=*/false);
  }
  StmtExprVisitor::VisitExpr_(op);
}

ParallelOpNode::ParallelOpNode(For root) : root_(root), V(this) {
  V.VisitStmt(root);
  // Cache any annotated layout/predicate on the outermost loop.
  using namespace attr;
  if (root_->annotations.count(kParallelLoopLayout)) {
    annotated_layout_unbound_ =
        Downcast<Fragment>(root_->annotations.Get(kParallelLoopLayout).value());
  }
  if (root_->annotations.count(kParallelLoopPredicate)) {
    annotated_predicate_ = Downcast<PrimExpr>(
        root_->annotations.Get(kParallelLoopPredicate).value());
  }
  if (auto padding_guard_anno =
          root_->annotations.Get(kParallelLoopRequiresPaddingGuard)) {
    if (auto padding_guard_bool = padding_guard_anno.value().try_cast<Bool>()) {
      annotated_requires_padding_guard_ = padding_guard_bool.value()->value;
    } else if (auto padding_guard_int =
                   padding_guard_anno.value().as<IntImmNode>()) {
      annotated_requires_padding_guard_ = padding_guard_int->value != 0;
    } else {
      LOG(WARNING) << "Loop annotation `" << kParallelLoopRequiresPaddingGuard
                   << "` expects Bool value (True/False), but got "
                   << padding_guard_anno.value().GetTypeKey()
                   << ". Ignore override.";
    }
  }
  // Collect cross-thread access info and buffer store info.
  PostOrderVisit(root_, [&](const ObjectRef &obj) {
    if (const auto *store = obj.as<BufferStoreNode>()) {
      auto buffer = store->buffer;
      if (IsSharedBuffer(buffer) || IsGlobalBuffer(buffer)) {
        has_cross_thread_access_ = true;
        store_shared_global_buffers_.emplace_back(buffer);
      } else if (IsFragmentBuffer(buffer)) {
        store_fragment_buffers_.emplace_back(buffer);
      }
    } else if (const auto *load = obj.as<BufferLoadNode>()) {
      if (IsSharedBuffer(load->buffer) || IsGlobalBuffer(load->buffer)) {
        has_cross_thread_access_ = true;
      }
    }
  });
}

TileOperator ParallelOpNode::Clone() const {
  auto op = make_object<ParallelOpNode>(*this);
  return ParallelOp(op);
}

void ParallelOpNode::ExpandBindValues(
    const Map<Var, PrimExpr> &bind_var_to_expr) {
  if (bind_var_to_expr.empty())
    return;

  // Helper function to recursively find BufferLoads through Bind values
  std::function<void(const PrimExpr &)> expand = [&](const PrimExpr &expr) {
    PostOrderVisit(expr, [&](const ObjectRef &node) {
      if (auto bl = node.as<BufferLoadNode>()) {
        if (IsFragmentBuffer(bl->buffer)) {
          RecordBufferAccess(bl->buffer, bl->indices, /*is_write=*/false);
        }
      } else if (auto var_node = node.as<VarNode>()) {
        auto var = GetRef<Var>(var_node);
        if (bind_var_to_expr.count(var)) {
          expand(bind_var_to_expr[var]);
        }
      }
    });
  };

  // Only expand Bind values that are used in root_
  // First, collect all vars used in root_
  std::unordered_set<const VarNode *> used_vars;
  PostOrderVisit(root_, [&](const ObjectRef &node) {
    if (auto var_node = node.as<VarNode>()) {
      used_vars.insert(var_node);
    }
  });

  // Only expand Bind values for vars that are actually used in root_
  for (const auto &[var, expr] : bind_var_to_expr) {
    if (used_vars.count(var.get())) {
      expand(expr);
    }
  }
}

void ParallelOpNode::RecordBufferAccess(const Buffer &buffer,
                                        const Array<PrimExpr> &indices,
                                        bool is_write) {
  auto it = indice_map_.find(buffer);
  if (it != indice_map_.end()) {
    ICHECK(StructuralEqual()(it->second.indices, indices))
        << buffer << ": " << indices << " and " << it->second.indices;
  } else {
    BufferAccessInfo info;
    info.indices = indices;
    it = indice_map_.emplace(buffer, std::move(info)).first;
    access_order_.push_back(buffer);
  }
  if (is_write) {
    it->second.is_write = true;
  } else {
    it->second.is_read = true;
  }
}

const ParallelOpNode::BufferAccessInfo &
ParallelOpNode::GetAccessInfo(const Buffer &buffer) const {
  auto it = indice_map_.find(buffer);
  ICHECK(it != indice_map_.end())
      << "Missing access info for buffer " << buffer;
  return it->second;
}

bool ParallelOpNode::IsBufferCompletelyReplicated(
    const Buffer &buffer, const LayoutMap &layout_map) const {
  if (!IsFragmentBuffer(buffer))
    return false;
  auto frag = layout_map[buffer].as<Fragment>().value();
  // buffer indices should be IntImm
  for (const auto &index : GetAccessInfo(buffer).indices) {
    if (!index.as<IntImmNode>()) {
      return false;
    } else if (index.as<IntImmNode>()->value != 0) {
      LOG(FATAL) << "buffer " << buffer << " is not completed replicated";
    }
  }
  return frag->IsCompletedReplicated();
}

Stmt ParallelOpNode::Lower(const LowerArgs &lower_args,
                           arith::Analyzer *analyzer) const {
  return root_;
}

// (annotations parsed in ctor; adoption happens in InferLayout)

bool ParallelOpNode::IsCommonAccessIndice(const Buffer &buffer) const {
  auto common_indice = loop_vars_.Map([](const auto &iv) { return iv->var; });
  return StructuralEqual()(GetAccessInfo(buffer).indices, common_indice);
}

/*! \brief Infer the layout for parallel operations based on different inference
 * levels
 *
 * The inference level controls how aggressively we try to infer and optimize
 * layouts:
 * - kStrict (2): Most conservative level. Only allows explicitly defined
 * layouts. Returns empty layout map if loop_layout_ is not already defined.
 *                Used when exact layout control is required.
 *
 * - kCommon (1): Intermediate level between strict and free.
 *                Allows common layout patterns while maintaining some
 * constraints.
 *
 * - kFree (0):   Most permissive level. Allows maximum optimization freedom.
 *                Will attempt layout inference even without source buffers.
 *                Can generate new layouts based on vectorization and thread
 * bounds. Used when maximum performance optimization is desired.
 */
LayoutMap ParallelOpNode::InferLayout(const LayoutInferArgs &layout_args,
                                      InferLevel level) const {
  if (loop_layout_inferred_)
    return {};
  loop_layout_requires_padding_guard_ = false;

  // Expand Bind values to find fragment buffer accesses
  if (!layout_args.bind_var_to_expr.empty()) {
    const_cast<ParallelOpNode *>(this)->ExpandBindValues(
        layout_args.bind_var_to_expr);
  }

  if (level == InferLevel::kStrict) {
    LayoutMap results;
    // Deduce buffers that should be complicated replicated.
    // For example:
    // for i in T.Parallel(m):
    //   fragment[0] = x[i]
    // then fragment[0] must be replicated on all threads.
    for (const auto &buffer : access_order_) {
      const auto &access = GetAccessInfo(buffer);
      if (layout_args.layout_map.count(buffer)) {
        continue;
      }
      if (!IsFragmentBuffer(buffer))
        continue;

      // Check if all indices are zero
      bool all_indices_zero = true;
      for (const auto &index : access.indices) {
        if (const auto *imm = index.as<IntImmNode>()) {
          if (imm->value != 0) {
            all_indices_zero = false;
            LOG(FATAL)
                << "Fragment buffer access with non-zero index [" << imm->value
                << "] is not supported. "
                << "Only fragment[0] access is allowed within T.Parallel loop.";
          }
        } else {
          // Non-constant index, not all zero
          all_indices_zero = false;
        }
      }

      // Only set layout if all indices are zero
      if (all_indices_zero) {
        Array<IterVar> forward_vars;
        for (const auto &s : buffer->shape) {
          forward_vars.push_back(
              IterVar(Range(0, s), Var(), IterVarType::kDataPar));
        }
        Var rep;
        auto rep_iter = IterVar({0, layout_args.thread_bounds->extent}, rep,
                                IterVarType::kDataPar);

        // Use default fragment indexing (single output dim) to
        // stay consistent with other ops (e.g., ReduceOp), and
        // bind the thread range for comparability.
        const PrimExpr &forward_thread = rep;
        auto frag = Fragment(forward_vars, /*forward_index=*/{}, forward_thread,
                             rep_iter)
                        ->BindThreadRange(layout_args.thread_bounds);
        results.Set(buffer, frag);
      }
    }
    return results;
  }

  // Collect fragment buffers with const index and all fragment_buffers
  std::vector<Buffer> const_index_fragment_buffer, fragment_buffers;
  for (const auto &buffer : access_order_) {
    const auto &access = GetAccessInfo(buffer);
    if (!IsFragmentBuffer(buffer))
      continue;
    fragment_buffers.push_back(buffer);

    bool is_const_index = true;
    for (const auto &index : access.indices) {
      if (!index.as<IntImmNode>()) {
        is_const_index = false;
        break;
      }
    }
    if (is_const_index) {
      const_index_fragment_buffer.push_back(buffer);
    }
  }

  // Determine if common layout propagation should be applied.
  // If there are fragment buffers with non-constant indices, we need to
  // propagate the common layout pattern to ensure consistency across all
  // fragments. Example cases:
  //   - Need propagation: frag_a[0] = T.min(frag_a[0], frag_b[i])
  //     (const index frag_a interacts with non-const index frag_b)
  //   - No propagation needed: shared_a[i] = frag_a[0]
  //     (const index frag_a with non-fragment buffer)

  bool allow_layout_propgate =
      const_index_fragment_buffer.empty() ||
      (fragment_buffers.size() > const_index_fragment_buffer.size());

  // Step 1: try to infer loop's partition from a source fragment
  Buffer source_buffer, read_source_buffer;
  bool source_buffer_is_write = false;
  Buffer replicated_write_buffer; // Backup: fully replicated write buffer

  for (const auto &buffer : access_order_) {
    const auto &access = GetAccessInfo(buffer);
    if (layout_args.layout_map.count(buffer)) {
      // skip reducers with rep=ALL
      if (auto info = reducer_info_map_.Get(buffer->data);
          info && info.value()->rep == ReducerRepType::ALL)
        continue;

      bool is_fully_replicated =
          IsBufferCompletelyReplicated(buffer, layout_args.layout_map);

      if (access.is_write) {
        source_buffer = buffer;
        source_buffer_is_write = true;
      } else {
        // Keep the buffer with largest number of indices
        // (which means the inference based on that buffer is more accurate)
        // as read_source_buffer to get more accurate layout
        // if the buffer is completed replicated, we don't need to infer the
        // layout from this buffer.
        if ((!read_source_buffer.defined() ||
             access.indices.size() >
                 GetAccessInfo(read_source_buffer).indices.size())) {
          read_source_buffer = buffer;
        }
        // If the buffer is not replicated and shape is equal to the
        // source_buffer, use it as source_buffer because the layout inference
        // is more accurate
        auto frag = layout_args.layout_map[buffer].as<Fragment>();
        if (frag.has_value() && is_one(frag.value()->ReplicateExtent()) &&
            !source_buffer.defined()) {
          source_buffer = buffer;
          source_buffer_is_write = false;
        }
      }
    }
  }
  // moved to ComputeLoopLayoutFromBuffer

  // Try to infer loop layout from buffers in order of preference only if we
  // don't already have a layout (e.g., from annotations):
  // 1. Annotated loop layout
  // 2. Non-replicated write buffer (most reliable)
  // 3. Non-replicated read buffer
  // 4. Fully replicated write buffer (backup, may cause issues)
  // 5. Free inference mode (no source buffer)
  if (!loop_layout_.defined() && annotated_layout_unbound_.defined()) {
    loop_layout_ = annotated_layout_unbound_.value()->BindThreadRange(
        layout_args.thread_bounds);
    loop_layout_requires_padding_guard_ = annotated_requires_padding_guard_;
    if (annotated_predicate_.defined()) {
      predicate_ = annotated_predicate_.value();
    }
  } else if (!loop_layout_.defined() && source_buffer.defined() &&
             (allow_layout_propgate || source_buffer_is_write)) {
    loop_layout_ = ComputeLoopLayoutFromBuffer(source_buffer, layout_args);
  } else if (!loop_layout_.defined() && level == InferLevel::kFree) {
    // For free layout inference
    // In free inference, try two mechanisms and prefer the one that
    // minimizes replication while remaining compatible:
    // 1) compute_loop_layout_from_buffer (always correct but may
    // over-replicate) 2) PlanLoopPartition (often smaller replication)
    Fragment candidate_from_buffer;
    Fragment candidate_from_plan;
    bool selected_plan_candidate = false;

    if (read_source_buffer.defined() && allow_layout_propgate) {
      candidate_from_buffer =
          ComputeLoopLayoutFromBuffer(read_source_buffer, layout_args);
    }

    // try to infer loop layout with two mechanisms and choose the best one
    {
      candidate_from_plan = ComputePlanCandidate(layout_args);
    }

    // Choose the best candidate:
    if (candidate_from_buffer.defined() && candidate_from_plan.defined()) {
      loop_layout_ = ChooseBestCandidate(candidate_from_buffer,
                                         candidate_from_plan, layout_args);
    } else if (candidate_from_plan.defined()) {
      loop_layout_ = candidate_from_plan;
      selected_plan_candidate = true;
      DLOG(INFO) << "[FreeInfer] only PlanLoopPartition available, choose it.";
    } else if (candidate_from_buffer.defined()) {
      loop_layout_ = candidate_from_buffer;
      DLOG(INFO)
          << "[FreeInfer] only compute_from_buffer available, choose it.";
    }
    loop_layout_requires_padding_guard_ =
        selected_plan_candidate && indice_map_.empty();
  } else if (!loop_layout_.defined()) {
    // In non-free mode without a source buffer, if we don't have any layout
    // yet (e.g., no annotation), we have nothing to infer here.
    return {};
  }

  // Non-fragment SIMT loops may deliberately over-cover a ragged iteration
  // space; PartitionLoop emits guards for the padded points. Fragment/reducer
  // loops stay strict because padding would change per-thread ownership.
  auto injective_res =
      loop_layout_->DetectInjective(loop_layout_requires_padding_guard_);
  if (!injective_res->errors.empty()) {
    std::ostringstream oss;
    oss << "Loop layout is not injective: " << loop_layout_->DebugOutput()
        << '\n'
        << "  errors: " << injective_res->errors << '\n'
        << "  loop AST: " << root_;
    throw LoopLayoutInjectiveException(oss.str());
  }

  PrimExpr loop_thread_extent = loop_layout_->ThreadExtent();

  auto block_size = layout_args.thread_bounds->extent;
  if (loop_layout_.defined()) {
    if (loop_layout_->ThreadRange().defined()) {
      auto thread_range = loop_layout_->ThreadRange();
      block_size = thread_range->extent;
      AddPredicate(GE(InputPlaceholder(0), thread_range->min));
      AddPredicate(
          LT(InputPlaceholder(0), thread_range->min + thread_range->extent));
    }
  }

  if (!analyzer_.CanProveEqual(loop_thread_extent, block_size)) {
    AddPredicate(LT(InputPlaceholder(0),
                    loop_thread_extent + layout_args.thread_bounds->min));
  }

  // Step 2: Check that the loop's partition can correctly align with all source
  // fragment, and infer layout only when it's not yet layout-ed.
  ValidateCandidateAgainstFragments(
      loop_layout_, layout_args, /*throw_on_error=*/true,
      /*check_forward_index=*/false, source_buffer);

  // Step 3: Build replication guards
  BuildReplicationGuardsIfNeeded(
      layout_args, store_shared_global_buffers_, store_fragment_buffers_,
      has_cross_thread_access_, const_index_fragment_buffer);

  // Step 4: Collect buffer fragments
  LayoutMap results;
  for (const auto &buffer : access_order_) {
    if (!layout_args.layout_map.count(buffer)) {
      auto dst_layout = CompleteBufferFragment(buffer)->BindThreadRange(
          layout_args.thread_bounds);
      results.Set(buffer, dst_layout);
    }
  }
  loop_layout_inferred_ = true;
  return results;
}

Optional<PrimExpr> ParallelOpNode::GetPredicate(Var thread_var) const {
  if (predicate_.defined()) {
    return Substitute(predicate_.value(), {{InputPlaceholder(0), thread_var}});
  } else {
    return std::nullopt;
  }
}

Fragment ParallelOpNode::CompleteBufferFragment(const Buffer &buffer) const {
  ICHECK(loop_layout_.defined());
  if (IsCommonAccessIndice(buffer)) {
    return loop_layout_;
  }
  // Prefer a simple path: if original 2D indices form a bijective map, invert
  // them directly and avoid introducing a synthetic replicate dimension.
  {
    auto res2d =
        arith::DetectIterMap(GetAccessInfo(buffer).indices, ToVMap(loop_vars_),
                             1, arith::IterMapLevel::Bijective,
                             const_cast<arith::Analyzer *>(&analyzer_));
    if (res2d->errors.empty()) {
      Layout ind_inv2d =
          Layout(loop_vars_, GetAccessInfo(buffer).indices)->Inverse();
      PrimExpr indice_rep_extent = 1;
      PrimExpr loop_rep_extent = loop_layout_->ReplicateExtent();
      PrimExpr dest_buffer_rep_extent = indice_rep_extent * loop_rep_extent;
      Array<PrimExpr> fwd2;
      for (size_t i = 0; i < buffer->shape.size(); i++) {
        fwd2.push_back(InputPlaceholder(i));
      }
      PrimExpr thd_b2 =
          loop_layout_->ForwardThread(ind_inv2d->Forward(fwd2), std::nullopt);
      return Fragment(buffer->shape, {}, thd_b2, dest_buffer_rep_extent,
                      std::nullopt)
          ->CondenseReplicateVar();
    }
  }
  // Otherwise, infer an extra flattened iterator that captures truly-unused
  // pieces of the loop space (if any), then try inversion with it.
  PrimExpr rep_b = MakeFlattenedExpression(DivideUnusedIterators(
      GetAccessInfo(buffer).indices, loop_vars_, &analyzer_));
  auto bijective_indice = GetAccessInfo(buffer).indices;
  bijective_indice.push_back(rep_b);
  Layout ind_inv = Layout(loop_vars_, bijective_indice)->Inverse();

  PrimExpr indice_rep_extent =
      ind_inv->InputShape().back(); // this is the size of rep_b
  PrimExpr loop_rep_extent = loop_layout_->ReplicateExtent();
  PrimExpr dest_buffer_rep_extent = indice_rep_extent * loop_rep_extent;
  Array<PrimExpr> fwd;
  for (size_t i = 0; i < buffer->shape.size(); i++) {
    fwd.push_back(InputPlaceholder(i));
  }
  fwd.push_back(FloorMod(ReplicationPlaceholder(), indice_rep_extent));
  PrimExpr thd_b = loop_layout_->ForwardThread(
      ind_inv->Forward(fwd),
      FloorDiv(ReplicationPlaceholder(), indice_rep_extent));
  return Fragment(buffer->shape, {}, thd_b, dest_buffer_rep_extent,
                  std::nullopt)
      ->CondenseReplicateVar();
}

TVM_FFI_STATIC_INIT_BLOCK() { ParallelOpNode::RegisterReflection(); }

bool ParallelOpNode::ValidateCandidateAgainstFragments(
    const Fragment &candidate, const LayoutInferArgs &layout_args,
    bool throw_on_error, bool check_forward_index,
    const Buffer &source_buffer) const {
  auto vars =
      loop_vars_.Map([](const IterVar &iv) { return PrimExpr(iv->var); });
  for (const auto &buffer : access_order_) {
    const auto &access = GetAccessInfo(buffer);
    if (!layout_args.layout_map.count(buffer))
      continue;
    if (auto info = reducer_info_map_.Get(buffer->data);
        info && info.value()->rep == ReducerRepType::ALL)
      continue;
    auto fragment = layout_args.layout_map[buffer].as<Fragment>().value();
    std::ostringstream oss;
    bool success = true;
    if (access.is_read &&
        !ProveFragmentContains(candidate, fragment, vars, access.indices,
                               analyzer_, check_forward_index)) {
      if (throw_on_error) {
        oss << "Layout infer conflict between " << buffer << " and "
            << source_buffer << " in T.Parallel loop:" << '\n'
            << "    loop " << candidate->DebugOutput() << '\n'
            << "    fragment " << fragment->DebugOutput() << '\n';
      }
      success = false;
    }
    // Fragment writes need exact owner-thread compatibility. Coverage alone
    // allows extra loop threads to write local slots that belong to different
    // logical fragment elements under the buffer layout.
    if (access.is_write &&
        (!ProveFragmentContains(fragment, candidate, access.indices, vars,
                                analyzer_, check_forward_index) ||
         !ProveFragmentContains(candidate, fragment, vars, access.indices,
                                analyzer_, check_forward_index))) {
      if (throw_on_error) {
        oss << "Layout infer conflict between " << buffer << " and "
            << source_buffer << " in T.Parallel loop:" << '\n'
            << "    loop " << candidate->DebugOutput() << '\n'
            << "    fragment " << fragment->DebugOutput() << '\n';
      }
      success = false;
    }
    if (!success) {
      if (throw_on_error) {
        throw LayoutConflictException(oss.str());
      }
      return false;
    }
  }
  return true;
}

Fragment ParallelOpNode::ComputeLoopLayoutFromBuffer(
    const Buffer &buffer, const LayoutInferArgs &layout_args) const {
  Fragment src_layout = layout_args.layout_map[buffer].as<Fragment>().value();
  DLOG(INFO) << "[compute_loop_layout_from_buffer] infer from buffer `"
             << buffer << "` of layout " << src_layout->DebugOutput() << '\n';

  Fragment result;

  if (IsCommonAccessIndice(buffer)) {
    result = src_layout;
  } else {
    Var rep("_rep");
    auto rep_iter =
        IterVar({0, src_layout->ReplicateExtent()}, rep, IterVarType::kDataPar);
    PrimExpr loop_var_to_thread =
        src_layout->ForwardThread(GetAccessInfo(buffer).indices, rep);
    loop_var_to_thread = analyzer_.Simplify(loop_var_to_thread);
    PostOrderVisit(loop_var_to_thread, [&](const ObjectRef &objref) {
      if (auto opt_var = objref.as<Var>();
          opt_var && inner_vars_.count(*opt_var)) {
        std::ostringstream oss;
        oss << "Cannot lower access to thread-distributed fragment `"
            << buffer->name
            << "`: its elements are spread across CUDA threads, so the "
               "non-parallel loop variable `"
            << *opt_var
            << "` cannot be used to address them safely (thread map: "
            << loop_var_to_thread
            << "). Use T.reduce_sum for reductions, T.Parallel for elementwise "
               "access, or stage through shared memory to change the layout.";
        throw LayoutConflictException(oss.str());
      }
    });

    try {
      result = Fragment(loop_vars_, {}, loop_var_to_thread, rep_iter)
                   ->BindThreadRange(layout_args.thread_bounds);
    } catch (const Error &err) {
      std::ostringstream msg;
      msg << "Layout inference for buffer `" << buffer->name
          << "` failed inside `T.parallel` loop.";

      msg << "\nUnderlying TVM error: " << err.what();
      msg << "\nProblematic loop AST:\n " << root_;
      msg << "\nHint: ensure the loop extent divides the thread binding or "
             "adjust the fragment mapping.";
      LOG(FATAL) << msg.str();
    }
  }
  DLOG(INFO) << "[compute_loop_layout_from_buffer] ... and get "
             << result->DebugOutput() << '\n';
  // Lei: This is a tradeoff, disable it for now.
  // // Try DeReplicate first to reduce replication if possible.
  // Fragment dereplicated_layout = candidate_from_buffer->DeReplicate();
  // if (ValidateCandidateAgainstFragments(
  //         dereplicated_layout, T, /*throw_on_error=*/false,
  //         /*check_forward_index=*/false,
  //         /*source_buffer=*/read_source_buffer)) {
  //   candidate_from_buffer = dereplicated_layout;
  // }
  return result;
}

Fragment
ParallelOpNode::ComputePlanCandidate(const LayoutInferArgs &layout_args) const {
  // Vectorize Size must be aware of the buffer_remap
  // As the pass will do post processing to the layout
  auto maybe_remapped_root_ = IfBufferRemapLoopGenerator::run(
      root_, layout_args.buffer_remap, layout_args.layout_map);
  int vector_size = GetVectorizeSize(maybe_remapped_root_, layout_args.analyzer,
                                     layout_args.layout_map, reducer_info_map_);
  DLOG(INFO) << "[PlanLoopPartition] vector_size = " << vector_size << '\n';

  PrimExpr loop_total_size = 1;
  for (Stmt l = root_; l.as<For>().has_value(); l = l.as<For>().value()->body)
    loop_total_size = loop_total_size * l.as<For>().value()->extent;
  DLOG(INFO) << "[PlanLoopPartition] loop_total_size = " << loop_total_size
             << '\n';
  bool has_fragment_access = !indice_map_.empty();
  if (has_fragment_access) {
    while (!analyzer_.CanProve(
               floormod(loop_total_size, layout_args.thread_bounds->extent *
                                             vector_size) == 0) &&
           vector_size > 1) {
      vector_size /= 2;
    }
  } else if (!root_->annotations.count(attr::kCoalescedWidth)) {
    vector_size = SelectMinPaddingVectorSize(vector_size, loop_total_size,
                                             layout_args.thread_bounds->extent,
                                             &analyzer_);
  }
  DLOG(INFO) << "[PlanLoopPartition] after adjust: vector_size = "
             << vector_size << '\n';

  // Check if coalesced_width is defined
  if (auto coalesced_width = root_->annotations.Get(attr::kCoalescedWidth)) {
    if (const auto *imm = coalesced_width->as<IntImmNode>()) {
      int expected = imm->value;
      // Verify that vector_size is divisible by expected
      if (vector_size % expected != 0) {
        LOG(FATAL) << "Vector size " << vector_size
                   << " is not divisible by coalesced width " << expected;
      }
      vector_size = expected;
    } else {
      LOG(FATAL) << "coalesced_width should be an IntImmNode.";
    }
  }
  DLOG(INFO) << "[PlanLoopPartition] root_ = " << root_
             << " ############# vector_size = " << vector_size
             << ", thread_bounds = " << layout_args.thread_bounds << '\n';
  auto plan = PlanLoopPartition(root_, vector_size, layout_args.thread_bounds);
  DLOG(INFO) << "[PlanLoopPartition] candidate = " << plan->DebugOutput()
             << '\n';
  return plan;
}

void ParallelOpNode::BuildReplicationGuardsIfNeeded(
    const LayoutInferArgs &layout_args,
    const std::vector<Buffer> &store_shared_global_buffers,
    const std::vector<Buffer> &store_fragment_buffers,
    bool has_cross_thread_access,
    const std::vector<Buffer> &const_index_fragment_buffer) const {
  if (is_one(loop_layout_->ReplicateExtent()))
    return;
  if (!has_cross_thread_access)
    return;

  if (!store_fragment_buffers.empty()) {
    bool replicate_is_from_dynamic_index_fragment = false;
    for (const auto &fragment : store_fragment_buffers) {
      if (!layout_args.layout_map.count(fragment)) {
        continue;
      }

      auto fragment_layout =
          layout_args.layout_map[fragment].as<Fragment>().value();
      if (is_one(fragment_layout->ReplicateExtent()))
        continue;

      if (analyzer_.CanProveEqual(fragment_layout->ReplicateExtent(),
                                  loop_layout_->ReplicateExtent()))
        continue;
      if (std::find(const_index_fragment_buffer.begin(),
                    const_index_fragment_buffer.end(),
                    fragment) == const_index_fragment_buffer.end()) {
        replicate_is_from_dynamic_index_fragment = true;
      }
    }

    if (!replicate_is_from_dynamic_index_fragment)
      return;

    ICHECK(store_shared_global_buffers.empty())
        << "Invalid layout: cannot have both fragment and shared store buffers "
           "in replicated loop layout.";
    return;
  } else {
    auto inv = loop_layout_->Inverse();
    Array<PrimExpr> fwd;
    for (size_t i = 0; i < loop_layout_->OutputDim(); i++)
      fwd.push_back(0);
    fwd.push_back(InputPlaceholder(0) - layout_args.thread_bounds->min);
    auto rep = inv->Forward(fwd).back();
    AddPredicate(EQ(rep, 0));
  }
}
Fragment
ParallelOpNode::ChooseBestCandidate(const Fragment &candidate_from_buffer,
                                    const Fragment &candidate_from_plan,
                                    const LayoutInferArgs &layout_args) const {
  // Strategy overview:
  // 1) Validate each candidate against all known source fragments. If only one
  //    is compatible, choose it immediately.
  // 2) If both are compatible, compare their containment relation:
  //      - If buffer-based contains plan-based, prefer plan (usually smaller
  //      rep).
  //      - If plan-based contains buffer-based, prefer buffer.
  // 3) If neither contains the other, prefer the one with provably smaller or
  //    equal replication extent; otherwise fall back to buffer-based candidate.
  // Note: Final global validation happens after selection elsewhere.
  auto vars =
      loop_vars_.Map([](const IterVar &iv) { return PrimExpr(iv->var); });
  auto contains = [&](const Fragment &big, const Fragment &small) {
    // contains(A, B) means: for any loop index, the threads that access
    // B's elements are a subset of those that access A's elements.
    return ProveFragmentContains(small, big, vars, vars, analyzer_);
  };

  bool buf_ok =
      ValidateCandidateAgainstFragments(candidate_from_buffer, layout_args);
  bool plan_ok =
      ValidateCandidateAgainstFragments(candidate_from_plan, layout_args);

  if (buf_ok && !plan_ok) {
    DLOG(INFO)
        << "[FreeInfer] prefer compute_from_buffer (only valid candidate).";
    return candidate_from_buffer;
  }
  if (plan_ok && !buf_ok) {
    DLOG(INFO)
        << "[FreeInfer] prefer PlanLoopPartition (only valid candidate).";
    return candidate_from_plan;
  }
  if (!(buf_ok && plan_ok)) {
    // Both invalid here; let the caller continue to final validation/throw.
    // Returning buffer-based candidate keeps behavior deterministic.
    return candidate_from_buffer; // arbitrary; caller will catch later
  }

  // Both candidates may be semantically valid for reads from a fully
  // replicated fragment.  The plan candidate often minimizes replication by
  // assigning one logical element to each CUDA thread, but that can turn a
  // register fragment access into `fragment[threadIdx.x]`.  CUDA cannot keep a
  // register array with runtime indexing in scalar registers, so this shape
  // spills the whole fragment to local memory before the final store.
  //
  // When this parallel loop publishes data to shared/global memory, prefer the
  // valid candidate that keeps fragment physical indices independent of the
  // thread variable.  For a fully replicated readback this selects the
  // fragment-derived layout, and BuildReplicationGuardsIfNeeded then emits the
  // single-replica guard (e.g. `threadIdx.x == 0`) to avoid duplicate stores.
  // Fragment-only loops keep the old containment/replication heuristic below.
  if (!store_shared_global_buffers_.empty()) {
    bool buf_thread_index =
        HasThreadDependentFragmentIndex(candidate_from_buffer, layout_args);
    bool plan_thread_index =
        HasThreadDependentFragmentIndex(candidate_from_plan, layout_args);
    if (buf_thread_index != plan_thread_index) {
      return buf_thread_index ? candidate_from_plan : candidate_from_buffer;
    }
  }

  bool buf_contains_plan = contains(candidate_from_buffer, candidate_from_plan);
  bool plan_contains_buf = contains(candidate_from_plan, candidate_from_buffer);

  auto rep_buf = candidate_from_buffer->ReplicateExtent();
  auto rep_plan = candidate_from_plan->ReplicateExtent();

  // Prefer the contained candidate (tends to minimize replication while
  // respecting access coverage):
  if (buf_contains_plan && !plan_contains_buf) {
    return candidate_from_plan;
  }
  if (plan_contains_buf && !buf_contains_plan) {
    return candidate_from_buffer;
  }
  // Neither strictly contains the other; prefer the one with smaller/equal rep.
  if (analyzer_.CanProve(rep_plan <= rep_buf)) {
    return candidate_from_plan;
  }
  // Safe fallback: buffer-based candidate is always correct.
  return candidate_from_buffer;
}

bool ParallelOpNode::HasThreadDependentFragmentIndex(
    const Fragment &candidate, const LayoutInferArgs &layout_args) const {
  // Mirror the index substitution that PartitionLoop would apply for this
  // candidate without actually rebuilding the loop.  We create symbolic output
  // coordinates plus a symbolic thread variable, invert the candidate layout
  // back to the original loop variables, and then apply each known fragment
  // layout's Forward() map.  If any resulting physical fragment index still
  // contains the symbolic thread variable, lowering this candidate would emit a
  // runtime-thread-indexed register array access such as
  // `fragment[threadIdx.x]`.
  if (!candidate.defined()) {
    return false;
  }

  int old_loop_depth = static_cast<int>(loop_vars_.size());
  if (candidate->InputDim() != static_cast<size_t>(old_loop_depth)) {
    return false;
  }

  Var thread_var("__tl_candidate_thread");
  Array<PrimExpr> partition_vars;
  for (size_t i = 0; i < candidate->OutputDim(); ++i) {
    partition_vars.push_back(Var("__tl_candidate_i" + std::to_string(i)));
  }
  partition_vars.push_back(thread_var);

  auto inv_loop = candidate->InverseWithLevel().first;
  Array<PrimExpr> recovered_indices = inv_loop->Forward(partition_vars);

  Map<Var, PrimExpr> loop_var_map;
  for (int i = 0; i < old_loop_depth; ++i) {
    loop_var_map.Set(loop_vars_[i]->var, recovered_indices[i]);
  }

  Map<Var, PrimExpr> thread_offset_map;
  bool has_thread_offset = false;
  if (candidate->ThreadRange().defined()) {
    auto range = candidate->ThreadRange();
    thread_offset_map.Set(thread_var, thread_var - range->min);
    has_thread_offset = true;
  }

  auto uses_thread_var = [&](const PrimExpr &expr) {
    PrimExpr simplified = analyzer_.Simplify(expr);
    return tirx::UsesVar(simplified, [&](const VarNode *var_node) {
      return GetRef<Var>(var_node).same_as(thread_var);
    });
  };

  for (const auto &buffer : access_order_) {
    if (!IsFragmentBuffer(buffer) || !layout_args.layout_map.count(buffer)) {
      continue;
    }
    auto fragment_layout = layout_args.layout_map[buffer].as<Fragment>();
    if (!fragment_layout.has_value()) {
      continue;
    }

    Array<PrimExpr> indices = GetAccessInfo(buffer).indices.Map(
        [&](const PrimExpr &index) { return Substitute(index, loop_var_map); });
    if (has_thread_offset) {
      indices = Substitute(indices, thread_offset_map);
    }

    Array<PrimExpr> physical_indices =
        fragment_layout.value()->Forward(indices);
    for (const auto &index : physical_indices) {
      if (uses_thread_var(index)) {
        return true;
      }
    }
  }

  return false;
}

} // namespace tl
} // namespace tvm
