Combine for RaggedIterDomain (#5716)

This PR introduces the combine operation as discussed in the
RaggedIterDomain [design
doc](https://github.com/NVIDIA/Fuser/blob/main/doc/dev/ragged_iter_domain_design_doc.md).

One design decision that I changed from the original design doc is about
detecting and validating component iter domains. Previously, I was
thinking about using the exact graph to find the corresponding component
iter domain for a given ragged iter domain (e.g.,
#5550 (comment)).
However, it won't work, for example, when a fusion is segmented and a
segment does not have the corresponding `Partition` expr for a
`RaggedIterDomain`. For example, when a tensor is used as an input for
`asNested`, followed by some other operations, if the fusion is
segmented after some operations, the latter segment won't be able to see
the `asNested` and the `Partition` operations as they don't exist in the
segment. This could be alleviated by providing an exact graph for the
whole complete fusion, but more fundamentally, if a fusion has a nested
tensor as an input, there doesn't seem to be any reasonable way to
attach a `Partition` expr.

See doc/dev/ragged_iter_domain_combine_design_doc.md‎ for detailed
discussions. At this moment, I decided to not worry too much about the
validation and assume the correctness is guaranteed by the user.

Note that partitioning is still limited to 1D extents. Multi-dim offsets
will be the next step of this series of RPs.

* Update *
Tracking which iter domains correspond to which extent iter domains
seems to be actually necessary for supporting combine with ragged iter
domains produced by multi-dim extent tensors. I'll revisit this as part
of multi-dim combine work, but my current plan is to take Option 4 as
described in the design doc.

Author: naoyam

csrc/dispatch.h

@@ -116,6 +116,7 @@ class Val;
   f(ScanOp);                         \
   f(Merge);                          \
   f(Partition);                      \
+  f(Combine);                        \
   f(Swizzle);                        \
   f(Swizzle2D);                      \
   f(Resize);                         \

csrc/ir/internal_base_nodes.cpp

@@ -1052,6 +1052,107 @@ std::pair<IterDomain*, RaggedIterDomain*> RaggedIterDomain::partition(
   return {component_id, ragged_id};
 }
 
+IterDomain* RaggedIterDomain::combine(
+    IterDomain* component,
+    RaggedIterDomain* ragged) {
+  NVF_ERROR(component != nullptr, "combine: component IterDomain is null");
+  NVF_ERROR(ragged != nullptr, "combine: ragged IterDomain is null");
+
+  NVF_ERROR(
+      !component->isA<RaggedIterDomain>(),
+      "combine: component must be a regular IterDomain, got RaggedIterDomain: ",
+      component->toString());
+
+  // Validate that component and ragged have compatible properties
+  NVF_ERROR_EQ(
+      component->getParallelType(),
+      ParallelType::Serial,
+      "Combining parallelized IterDomain not supported: ",
+      component->toString());
+
+  NVF_ERROR_EQ(
+      ragged->getParallelType(),
+      ParallelType::Serial,
+      "Combining parallelized RaggedIterDomain not supported: ",
+      ragged->toString());
+
+  NVF_ERROR_EQ(
+      component->getIterType(),
+      IterType::Iteration,
+      "combine: only IterType::Iteration is supported for component, got ",
+      component->getIterType(),
+      " for IterDomain: ",
+      component->toString());
+
+  NVF_ERROR_EQ(
+      ragged->getIterType(),
+      IterType::Iteration,
+      "combine: only IterType::Iteration is supported for ragged, got ",
+      ragged->getIterType(),
+      " for RaggedIterDomain: ",
+      ragged->toString());
+
+  // Validate component-ragged pairing when Partition definition is available
+  // (Option 3 of doc/dev/ragged_iter_domain_combine_design_doc.md).
+  // Only validate when the RaggedIterDomain has a direct Partition definition.
+  // After propagation (e.g., set() operations), the definition may be nullptr,
+  // in which case we trust the user to provide the correct component.
+  if (ragged->definition() != nullptr &&
+      ragged->definition()->isA<Partition>()) {
+    auto* partition = ragged->definition()->as<Partition>();
+    IterDomain* expected_component = partition->component();
+
+    NVF_ERROR(
+        component == expected_component,
+        "combine: component mismatch. The provided component does not match ",
+        "the component from the Partition that created this "
+        "RaggedIterDomain.\n",
+        "  Provided component: ",
+        component->toString(),
+        "\n",
+        "  Expected component: ",
+        expected_component->toString());
+  }
+  // If no Partition definition (after set, in segmented fusion, or external
+  // input), trust the user and proceed without validation
+
+  // The combined extent is the sum of all extents in the ragged dimension
+  // For a 1D extents tensor [e0, e1, ..., en-1], the total is sum(extents)
+  TensorView* extents_tv = ragged->extents();
+  NVF_ERROR(extents_tv != nullptr, "combine: ragged extents tensor is null");
+
+  // It is still assumed the extents tensor is just 1D
+  NVF_ERROR_EQ(
+      std::ranges::distance(
+          extents_tv->getLogicalDomain() | TensorDomain::kNoReductions),
+      1,
+      "Unexpected rank of extent tensor: ",
+      extents_tv->toString());
+
+  auto container = component->container();
+  auto zero = container->zeroVal(DataType::Index);
+
+  // Create a symbolic extent for the combined IterDomain
+  // This represents the sum of all ragged extents, i.e.,
+  // sum(extents_tv, {0}). We could use the sum output as the extent
+  // but we would need to extract the scalar value out of the 0-dim
+  // tensor. For now, we leave it as a symbolic Val.
+  Val* combined_extent =
+      IrBuilder::createInContainer<Val>(container, DataType::Index);
+
+  // Create the combined IterDomain with the symbolic extent
+  IterDomain* combined_id = IterDomainBuilder(zero, combined_extent)
+                                .parallel_type(ParallelType::Serial)
+                                .iter_type(IterType::Iteration)
+                                .build();
+
+  // Create the Combine expression linking component + ragged -> combined
+  IrBuilder::createInContainer<Combine>(
+      container, combined_id, component, ragged);
+
+  return combined_id;
+}
+
 TensorDomain::TensorDomain(
     IrBuilderPasskey passkey,
     std::vector<IterDomain*> logical_domain,

csrc/ir/internal_base_nodes.h

@@ -499,6 +499,22 @@ class NVF_API RaggedIterDomain : public IterDomain {
       IterDomain* in,
       TensorView* extents);
 
+  //! Combine a component IterDomain with a RaggedIterDomain to flatten
+  //! This is the inverse of partition, creating a regular IterDomain
+  //!
+  //! \param component Component IterDomain (extent = num_components)
+  //! \param ragged RaggedIterDomain with variable extents per component
+  //! \return Regular IterDomain with extent = sum of all component extents
+  //!
+  //! This operation flattens the ragged structure back into a single dimension.
+  //! Example: component extent=3, ragged extents=[127, 0, 198]
+  //!          -> output extent = 325 (= 127 + 0 + 198)
+  //!
+  //! Note: We use "combine" instead of "merge" to differentiate from the
+  //! regular IterDomain::merge operation which only works with regular
+  //! IterDomains.
+  static IterDomain* combine(IterDomain* component, RaggedIterDomain* ragged);
+
   //! Override cloneWithoutRFactor to preserve RaggedIterDomain type
   IterDomain* cloneWithoutRFactor(bool map_with_original = false) override;
 

csrc/ir/internal_nodes.cpp

@@ -2708,6 +2708,33 @@ std::string Partition::toInlineString(int indent_size) const {
 
 NVFUSER_DEFINE_CLONE_AND_CREATE(Partition)
 
+Combine::Combine(
+    IrBuilderPasskey passkey,
+    IterDomain* out,
+    IterDomain* component,
+    RaggedIterDomain* ragged)
+    : Expr(passkey) {
+  addOutput(out);
+  addInput(component);
+  addInput(ragged);
+}
+
+std::string Combine::toString(int indent_size) const {
+  std::stringstream ss;
+  ss << "Combine: ";
+  ss << "component: " << component()->toString();
+  ss << " + ragged: " << ragged()->toString();
+  ss << " -> " << out()->toString();
+  ss << "\n";
+  return ss.str();
+}
+
+std::string Combine::toInlineString(int indent_size) const {
+  NVF_CHECK(false, "Combine can not be printed inline");
+}
+
+NVFUSER_DEFINE_CLONE_AND_CREATE(Combine)
+
 Swizzle::Swizzle(
     IrBuilderPasskey passkey,
     IterDomain* out_x,

csrc/ir/internal_nodes.h

@@ -1910,6 +1910,44 @@ class NVF_API Partition : public Expr {
   }
 };
 
+//! Combine a component IterDomain with a RaggedIterDomain to flatten
+//! This is the inverse of Partition, merging component and ragged dimensions
+//! into a single regular IterDomain
+class NVF_API Combine : public Expr {
+ public:
+  using Expr::Expr;
+
+  Combine(
+      IrBuilderPasskey,
+      IterDomain* out,
+      IterDomain* component,
+      RaggedIterDomain* ragged);
+
+  NVFUSER_DECLARE_CLONE_AND_CREATE
+
+  const char* getOpString() const override {
+    return "Combine";
+  }
+
+  std::string toString(int indent_size = 0) const override;
+  std::string toInlineString(int indent_size = 0) const override;
+
+  //! Output IterDomain (combined/flattened dimension)
+  IterDomain* out() const {
+    return output(0)->as<IterDomain>();
+  }
+
+  //! Component dimension input (extent = num_components)
+  IterDomain* component() const {
+    return input(0)->as<IterDomain>();
+  }
+
+  //! Ragged dimension input (variable extents per component)
+  RaggedIterDomain* ragged() const {
+    return input(1)->as<RaggedIterDomain>();
+  }
+};
+
 class Swizzle : public Expr {
  public:
   using Expr::Expr;