Examine the broadcast size for aliased input/output when identifying RW race

[Priya2698]

PR #2999 adds a presegmentation pass to force segmentation when inplace update can cause RW race.
This occurs when an intermediate tensorview is aliased to a fusion input, a RW race occurs, when the intermediate tensorview or the aliased input is in path of a broadcast.

This preseg pass currently does not consider how the size of the broadcasted tv differs from aliased input/output, in which case segmentation may not be required.

Consider the test:

Fuser/tests/cpp/test_alias.cpp

Lines 980 to 1029 in 2ec2b92

	TEST_F(AliasTest, ReuseBuffer_AliasAcrossSegments) {
	auto fusion = std::make_unique<Fusion>();
	FusionGuard fg(fusion.get());
	TensorView* tv0 = makeSymbolicTensor(2);
	TensorView* tv1 = makeSymbolicTensor(1);
	TensorView* tv2 = makeSymbolicTensor(2);
	fusion->addInput(tv0);
	fusion->addInput(tv1);
	fusion->addInput(tv2);
	TensorView* tv3 = add(tv0, IrBuilder::create<Val>(1.0)); // Group 0
	TensorView* tv4 =
	max(tv3, {0}); // Group 0 (use max instead to avoid numerical issues)
	TensorView* tv5 = add(tv4, tv1); // Group 0 (Non Broadcast after reduce,
	// keeps normalization scheduler away)
	TensorView* tv6 = add(tv5, tv2); // Group 1 (Broadcast after reduce)
	// Note: test alias;
	fusion->aliasOutputToInput(tv6, tv0, AllocationType::ReuseBuffer);
	// TODO: support output on aliased fusion #1488
	// remove tv7 after #1488
	// fusion->addOutput(tv6);
	TensorView* tv7 = add(tv6, IrBuilder::create<Val>(1.0)); // Group 0
	fusion->addOutput(tv7);
	auto options = at::TensorOptions().dtype(at::kFloat).device(at::kCUDA, 0);
	at::Tensor t0 = at::randn({128, 65}, options);
	at::Tensor t1 = at::randn({65}, options);
	at::Tensor t2 = at::randn({128, 65}, options);
	FusionExecutorCache fec(std::move(fusion));
	// Make a copy of `t0` because `t0` will be in-place updated.
	at::Tensor original_t0 = t0.clone();
	std::vector<at::Tensor> outputs = fec.runFusionWithInputs({t0, t1, t2});
	testValidate(
	fec.fusion(), outputs, {original_t0, t1, t2}, __LINE__, __FILE__);
	EXPECT_EQ(
	fec.getMostRecentKernelRuntime()->fusionSegments()->groups().size(), 2)
	<< "segmentation didn't happen as expected";
	auto t3 = original_t0.add(1.0);
	auto t4 = std::get<0>(at::max(t3, 0));
	auto t5 = t4.add(t1);
	auto t6 = t5.add(t2);
	EXPECT_TRUE(t0.allclose(t6))
	<< "`t0` should have been in-place updated to the same value as `t6`.";
	}

, the fusion is segmented such that the RW race does not occur. Additionally, the broadcasted size is not different from the aliased input/output. However, this preseg pass will still insert a segment set and split the fusion into 3 segments, even though 2 segments is functionally correct.

} // {Re-written complete fusion}
Segmented_Fusion Dump: -- fusion segments:
Segmented_Fusion{ 
groups: 
  reduction{0, 1, 2}
  pointwise{3, 4, 5}
edges: 
  e{ reduction{0, 1, 2} -> pointwise{3, 4, 5}(T5_g_float[ iS9{i2} ]) }

group details:
g{(reduction)
group id: 0
inputs:
  T0_g_float[ iS0{i0}, iS1{i2} ] float
  T1_g_float[ iS16{i2} ] float
outputs:
  T5_g_float[ iS9{i2} ] float


T3_l_float[ iS5{i0}, iS6{i2} ]
   = T0_g_float[ iS0{i0}, iS1{i2} ]
   + double(1);
(0)
T4_g_float[ rS7{i0}, iS8{i2} ]
   = reduction( T3_l_float[ iS5{i0}, iS6{i2} ], op = fmax, initial value = double(-inf), allreduce = false )
(1)
T5_g_float[ iS9{i2} ]
   = T4_g_float[ rS7{i0}, iS8{i2} ]
   + T1_g_float[ iS16{i2} ];
(2)
}

g{(pointwise)
group id: 1
inputs:
  T0_g_float[ iS0{i0}, iS1{i2} ] float
  T2_g_float[ iS3{i4}, iS17{i2} ] float
  T5_g_float[ iS9{i2} ] float
outputs:
  T7_g_float[ iS12{i4}, iS13{i2} ] float
  T8_g_float[ iS14{i4}, iS15{i2} ] float


T6_g_float[ bS10{1}, iS11{i2} ]
   = broadcast( T5_g_float[ iS9{i2} ] )
(3)
T7_g_float[ iS12{i4}, iS13{i2} ]
   = T6_g_float[ bS10{1}, iS11{i2} ]
   + T2_g_float[ iS3{i4}, iS17{i2} ];
(4)
T8_g_float[ iS14{i4}, iS15{i2} ]
   = T7_g_float[ iS12{i4}, iS13{i2} ]
   + double(1);
(5)
}

This issue is to track how the preseg pass needs to be modified to identify such cases.

[kevinstephano]

@Priya2698 can this issue be closed?

[Priya2698]

No, I need to investigate this issue and understand if any changes are required to the current implementation of the 'SegmentInplaceUpdate' preset pass.