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