Skip to content

Feature request: Fusing sibling exprs in segmentation #3829

Open
@naoyam

Description

This fusion is segmented to 3 segments with each expr as a distinctive segment. We should fuse tv1 and tv2, but that's not what happens.

  std::unique_ptr<Fusion> fusion_ptr = std::make_unique<Fusion>();
  Fusion& fusion = *fusion_ptr.get();
  FusionGuard fg(&fusion);

  auto tv0 = makeSymbolicTensor(1);
  fusion.addInput(tv0);

  auto tv1 = sin(tv0);
  fusion.addOutput(tv1);
  auto tv2 = cos(tv0);
  fusion.addOutput(tv2);

  auto tv3 = segment_set(tv0);
  fusion.addOutput(tv3);

  auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
  at::Tensor t0 = at::randn({1000}, options);
  std::vector<c10::IValue> inputs = {t0};

  FusionExecutorCache executor_cache(std::move(fusion_ptr));
  auto outputs = executor_cache.runFusionWithInputs(inputs);
  testValidate(&fusion, outputs, inputs, __LINE__, __FILE__);
Segmented_Fusion{
groups:
  pointwise{0}
  pointwise{1}
  no_op{2}
edges:

group details:
g{(pointwise)
group id: 0
inputs:
  T0_g_float[iS0{i0}] float
outputs:
  T1_g_float[iS1{i0}] float


T1_g_float[iS1{i0}]
   = sinf(T0_g_float[iS0{i0}]);
(0)
}

g{(pointwise)
group id: 1
inputs:
  T0_g_float[iS0{i0}] float
outputs:
  T2_g_float[iS2{i0}] float


T2_g_float[iS2{i0}]
   = cosf(T0_g_float[iS0{i0}]);
(1)
}

g{(no_op)
group id: 2
inputs:
  T0_g_float[iS0{i0}] float
outputs:
  T3_g_float[iS3{i0}] float


T3_g_float[iS3{i0}]
   = SegmenterSet( T0_g_float[iS0{i0}] )
(2)
}

} //Segmented_Fusion

The fusion segmenter only considers fusion of a producer expr group with a consumer expr group. Here, there's no such relationship between the sin and cos exprs. They are both consumers of tv0, and that relationship is not currently considered for fusion.

This isn't an issue if the whole fusion can be scheduled without segmentation. The repro above has segment_set, which is there just to enforce the segmentation step.

I encountered this case when testing #3659, which further promotes the input forwarding. That helps generate better segmentation results in some cases when segmentation is unavoidable, but it also exposed the above issue, resulting in segmenting previously fused expr groups.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Assignees

No one assigned

    Labels

    SegmentationIssues related to nvFuser Segmentation

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions