Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Feature request: Fusing sibling exprs in segmentation #3829

Open
naoyam opened this issue Feb 5, 2025 · 0 comments
Open

Feature request: Fusing sibling exprs in segmentation #3829

naoyam opened this issue Feb 5, 2025 · 0 comments
Labels
Segmentation Issues related to nvFuser Segmentation

Comments

@naoyam
Copy link
Collaborator

naoyam commented Feb 5, 2025

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.

@naoyam naoyam added the Segmentation Issues related to nvFuser Segmentation label Feb 5, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Segmentation Issues related to nvFuser Segmentation
Projects
None yet
Development

No branches or pull requests

1 participant