Add a test to verify a segment can hold multiple segment_sets.

This was a question I asked myself when working on #2639. Knowing this
allows me to simplify some segment-set-inserting logic there.

tests/cpp/test_segmentation.cpp

@@ -15,6 +15,8 @@
 
 namespace nvfuser {
 
+using testing::SizeIs;
+
 using SegmentationTest = NVFuserTest;
 
 TEST_F(SegmentationTest, Issue1284_Repro1) {
@@ -637,8 +639,7 @@ TEST_F(SegmentationTest, EraseReductionsInSegmentationEdges) {
     std::unique_ptr<Fusion> segment_fusion =
         segmented_fusion->makeFusion(group).second;
 
-    TensorView* segment_input =
-        segment_fusion->inputs().at(0)->as<TensorView>();
+    auto* segment_input = segment_fusion->inputs().at(0)->as<TensorView>();
 
     EXPECT_FALSE(segment_input->domain()->hasReduction());
   }
@@ -678,4 +679,27 @@ TEST_F(SegmentationTest, AliasedOutputOnSegmentation) {
       __FILE__);
 }
 
+TEST_F(SegmentationTest, MultipleSegmentSetsInOneSegment) {
+  auto fusion = std::make_unique<Fusion>();
+  FusionGuard fg(fusion.get());
+
+  auto* in = makeContigTensor(1);
+  auto* t = add(in, in);
+  auto* seg_out_0 = segment_set(t);
+  auto* seg_out_1 = segment_set(t);
+  auto* out = add(seg_out_0, seg_out_1);
+  fusion->addInput(in);
+  fusion->addOutput(out);
+
+  FusionExecutorCache fec(std::move(fusion));
+  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
+  at::Tensor in_tensor = at::randn({10}, options);
+  at::Tensor out_tensor = fec.runFusionWithInputs({in_tensor})[0];
+
+  testValidate(fec.fusion(), {out_tensor}, {in_tensor}, __LINE__, __FILE__);
+
+  FusionKernelRuntime* runtime = fec.getMostRecentKernelRuntime();
+  EXPECT_THAT(runtime->fusionSegments()->groups(), SizeIs(2));
+}
+
 } // namespace nvfuser