// possible, but may not be complete.
bool CoveredGroup::isEqualToOrSuperSetOf(const CoveredGroup& other) const {
  if (*this == other) {
    return true;
  }

  // When both are derived from split and correspond to either inner
  // or outer.
  if (split_in_.get() && other.split_in_.get() &&
      is_inner_ == other.is_inner_) {
    const CoveredGroups& split_in = *split_in_;
    const CoveredGroups& other_split_in = *other.split_in_;

    // When both have the same split input (and both correspond to
    // either inner or outer), they should cover the same exact
    // groups. This should only happen when broadcast is merged. For
    // example, suppose there are two tensors and they are scheduled as
    // follows;
    //
    // t0: [i0]
    // t1: [i1, b2]
    //
    // t1->merge(0, 1)->split(0, 4);
    // t0->split(0, 4)
    //
    // t0->inlineAt(t1, 1)
    //
    // In this case, t0->axis(0) and t1->axis(0) have the same
    // split input group, {i0, i1}. Note that b2 is not included as
    // it's a broadcast. Also note that both are the outer
    // output. Here, group_ of t0->axis(0) is the exact group of
    // t0->axis(0), while that of tv1->axis(0) is the exact group of
    // the t1->merge(0, 1) output. In this case, however, this merge
    // is just a merge of i1 and the b2 broadcast ID, so in terms of
    // covered exact groups, it's effectively the same as that of
    // t0->axis(0). All in all, as long as both correspond to either
    // inner or outer of the same split input, they should be
    // considered the same.
    if (split_in == other_split_in) {
      return true;
    }

    // Both are derived from a split but have differnt split input
    // groups. If the input groups of this split is a superset of the
    // input groups of the split of the other CoveredGroup, this
    // CoveredGroup is a superset
    if (std::all_of(
            other_split_in.begin(),
            other_split_in.end(),
            [&](const CoveredGroup& other_split_in_group) {
              return std::any_of(
                  split_in.begin(),
                  split_in.end(),
                  [&](const CoveredGroup& split_in_group) {
                    return split_in_group.isEqualToOrSuperSetOf(
                        other_split_in_group);
                  });
            })) {
      return true;
    }
  }

  // If the other CoveredGroup has a split input, it is sufficient to
  // satisfy that this CoveredGroup is equal to or
  // superior to the split input of other
  if (other.split_in_.get()) {
    if (std::all_of(
            other.split_in_->begin(),
            other.split_in_->end(),
            [&](const CoveredGroup& other_split_in) {
              return isEqualToOrSuperSetOf(other_split_in);
            })) {
      return true;
    }
  }

  // At this point, it does not mean it's definitely false but not
  // proven to be true either.

  return false;
}

LoopPromotionMapBuilder::computeCoveredGroups(
    const ValGraph& exact_graph,
    const IdModel& id_model) {
  // Map from an exact iter domain group, to all the exact iter domain groups it
  // covers
  std::unordered_map<ValGroup, std::shared_ptr<CoveredGroups>> covered_ids;

  ValGroups input_groups_of_graph =
      getInputGroupsOfExactGraph(exact_graph, id_model);

  for (const ValGroup& id_group :
       exact_graph.disjointValSets().disjointSets()) {
    // Initialize inputs
    if (input_groups_of_graph.has(id_group)) {
      auto init_groups = std::make_shared<CoveredGroups>();
      init_groups->insert(CoveredGroup(id_group));
      NVF_ERROR(covered_ids.emplace(id_group, init_groups).second);
    }

    // Initialize broadcast groups to empty since broadcast domains
    // don't matter for indexing
    if (std::any_of(id_group->begin(), id_group->end(), [&](Val* id) {
          return id->as<IterDomain>()->isBroadcast();
        })) {
      covered_ids[id_group] = std::make_shared<CoveredGroups>();
    }
  }

  ValGraphStmtSort exact_stmt_sort(exact_graph, input_groups_of_graph);

  for (const ExprGroup& exact_expr : exact_stmt_sort.exprs()) {
    // Initialize to empty group if not yet initialized
    for (const ValGroup& output_group : exact_graph.outputGroups(exact_expr)) {
      covered_ids.emplace(output_group, std::make_shared<CoveredGroups>());
    }

    const std::vector<ValGroup> input_groups =
        exact_graph.inputGroups(exact_expr);
    const std::vector<ValGroup> output_groups =
        exact_graph.outputGroups(exact_expr);

    // If this expr is a split, don't propagate the input coverage as
    // is but set the covered group of each output group by itself.
    // The input coverage info is propagated as the split input.
    if (exact_expr->front()->isA<Split>()) {
      NVF_ERROR(input_groups.size() == 1);
      const std::shared_ptr<CoveredGroups>& covered_groups =
          covered_ids.at(input_groups.at(0));

      for (const ValGroup& output_group : output_groups) {
        bool is_inner =
            output_group->has(exact_expr->front()->as<Split>()->inner());
        covered_ids[output_group]->insert(
            CoveredGroup(output_group, covered_groups, is_inner));
      }
      continue;
    }

    for (const ValGroup& output_group : output_groups) {
      // Note that an exact group may have multiple
      // exact expr groups and may have different coverage groups depending on
      // the expr groups. For example, this can happen with reshape or resize.
      // See test LoopPromotionCoverage for a concrete example.
      for (const ValGroup& inp_group : input_groups) {
        const std::shared_ptr<CoveredGroups>& inp_covered_groups =
            covered_ids.at(inp_group);
        covered_ids[output_group]->insert(
            inp_covered_groups->begin(), inp_covered_groups->end());
      }
    }
  }

  return covered_ids;
}

std::unordered_map<ValGroup, IterDomain*> LoopPromotionMapBuilder::
    projectIELPromotionToLoopGraph(
        const ValGraph& iel_graph,
        const std::unordered_map<ValGroup, IterDomain*>& iel_promotion_map,
        const ValGraph& loop_graph,
        const StatefulInliningInfo& inlining_info) const {
  const std::unordered_map<ValGroup, std::shared_ptr<CoveredGroups>>
      exact_covered_ids =
          computeCoveredGroups(idGraph(IdMappingMode::EXACT), id_model_);

  // Grab terminal iter domain in the loop groups.
  const VectorOfUniqueEntries<IterDomain*> terminal_loop_ids =
      computeTerminalLoopIds(inlining_info);

  std::unordered_map<ValGroup, IterDomain*> loop_promotion_map;

  for (const ValGroup& loop_group :
       loop_graph.disjointValSets().disjointSets()) {
    IterDomain* promotion_id = findPromotionOfLoopGroup(
        loop_group,
        iel_graph,
        iel_promotion_map,
        exact_covered_ids,
        terminal_loop_ids);
    if (promotion_id) {

// Test to verify the split-aware covered group analysis. See
// also https://github.com/NVIDIA/Fuser/pull/3877.
TEST_F(IdModelTest, CoveredGroups) {
  auto fusion_ptr = std::make_unique<Fusion>();
  auto& fusion = *fusion_ptr;
  FusionGuard fg(fusion_ptr.get());

  auto tv0 = makeContigConcreteTensor({-1, 1});
  fusion.addInput(tv0);
  auto tv1 = makeContigConcreteTensor({-1});
  fusion.addInput(tv1);

  auto tv2 = set(tv0);
  auto tv3 = reshape(tv1, {8}, {2, 4});
  auto tv4 = add(tv2, tv3);
  fusion.addOutput(tv4);

  for (auto tv : fusion.allTvs()) {
    tv->flatten();
  }

  inlineMost();

  IdModel id_model(&fusion, true);
  const auto& exact_graph = id_model.idGraph(IdMappingMode::EXACT);

  // The exact
  // group of the tv3 and tv4 IDs should cover both the inner and
  // outer split groups of the input group of the tv1 logical ID.
  const auto covered_groups =
      LoopPromotionMapBuilder::computeCoveredGroups(exact_graph, id_model);

  const auto& input_group = exact_graph.toGroup(tv1->getLogicalDomain().at(0));
  auto input_covered_group_it = covered_groups.find(input_group);
  ASSERT_NE(input_covered_group_it, covered_groups.end());
  const auto& input_covered_groups = input_covered_group_it->second;

  const auto& tv4_exact_group = exact_graph.toGroup(tv4->axis(0));
  auto tv4_exact_group_it = covered_groups.find(tv4_exact_group);
  ASSERT_NE(tv4_exact_group_it, covered_groups.end());
  const auto& tv4_covered_groups = tv4_exact_group_it->second;
  // It should consist of two CoveredGroups, both of which inheriths
  // from the logical ID of tv1 through a split
  EXPECT_EQ(tv4_covered_groups->size(), 2);
  for (const CoveredGroup& covered_group : *tv4_covered_groups) {
    EXPECT_EQ(covered_group.splitIn(), input_covered_groups);
    if (covered_group.isInner()) {
      EXPECT_EQ(
          covered_group.group(),
          exact_graph.toGroup(tv4->getLogicalDomain().at(1)));
    } else {
      EXPECT_EQ(
          covered_group.group(),
          exact_graph.toGroup(tv4->getLogicalDomain().at(0)));
    }
  }
}

// Repro of issue #3702
// https://github.com/NVIDIA/Fuser/issues/3702. Indexing traversal
// faied due to invalid loop promotion.
TEST_F(IdModelTest, InvalidLoopPromotion) {
  auto fusion_ptr = std::make_unique<Fusion>();
  auto& fusion = *fusion_ptr;
  FusionGuard fg(fusion_ptr.get());

  auto T0 = makeContigConcreteTensor({1, 32, 6});
  fusion.addInput(T0);
  auto T32 = makeContigConcreteTensor({1, 6, 2048}, DataType::BFloat16);
  fusion.addInput(T32);

  auto T6 = transpose(T0, 1, 2);
  auto T98 = broadcast(T6, {false, false, true, false});
  auto T99 = expand(
      T98,
      {IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(2L),
       IrBuilder::create<Val>(-1L)});
  auto T100 = reshape(
      T99,
      {IrBuilder::create<Val>(1L),
       IrBuilder::create<Val>(6L),
       IrBuilder::create<Val>(-1)});
  auto T11 = sin(T100);
  auto T13 = mul(T11, IrBuilder::create<Val>(1.0));
  auto T15 = castOp(DataType::BFloat16, T13);
  auto T43 = broadcast(T15, {false, true, false, false});
  auto T59 = expand(
      T43,
      {IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(32L),
       IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(-1L)});
  auto T61 = castOp(DataType::Float, T59);
  auto T10 = cos(T100);
  auto T12 = mul(T10, IrBuilder::create<Val>(1.0));
  auto T14 = castOp(DataType::BFloat16, T12);
  auto T41 = broadcast(T14, {false, true, false, false});
  auto T66 = expand(
      T41,
      {IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(8L),
       IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(-1L)});
  auto T79 = expand(
      T43,
      {IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(8L),
       IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(-1L)});
  auto T81 = castOp(DataType::Float, T79);
  auto T35 = reshape(
      T32,
      {IrBuilder::create<Val>(1L),
       IrBuilder::create<Val>(6L),
       IrBuilder::create<Val>(32L),
       IrBuilder::create<Val>(64L)});
  auto T36 = transpose(T35, 1, 2);
  auto T47 = castOp(DataType::Float, T36);
  auto T46 = expand(
      T41,
      {IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(32L),
       IrBuilder::create<Val>(-1L),
       IrBuilder::create<Val>(-1L)});
  auto T48 = castOp(DataType::Float, T46);
  auto T49 = mul(T47, T48);
  fusion.addOutput(T61);
  fusion.addOutput(T66);
  fusion.addOutput(T81);
  fusion.addOutput(T36);
  fusion.addOutput(T49);

  auto options_bf16 =
      at::TensorOptions().dtype(at::kBFloat16).device(at::kCUDA, 0);
  auto options_fp32 =
      at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  auto t0 = at::randn({1, 32, 6}, options_fp32);
  auto t32 = at::randn({1, 6, 2048}, options_bf16);
  std::vector<c10::IValue> inputs({t0, t32});

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs(inputs);
  testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
}

} // namespace nvfuser