Fix the normalization scheduler to accept DID loop split.

[wujingyue]

I'm sure we'll need more tests to be confident, but this incremental PR feels good!

For #2563

[wujingyue]

Not used

[github-actions]

Review updated until commit 5723c20

Description

• Added support for DID loop split in normalization scheduler.

• Introduced getShardedLoopAxis function for loop axis retrieval.

• Enhanced scheduleReductionTV to handle DID loop split.

• Added a new test case DivideBySum for multidevice sharding.

---

Changes walkthrough 📝

	Relevant files
Enhancement	utils.cpp Add getShardedLoopAxis and remove getShardedIterDomains    csrc/multidevice/utils.cpp Removed getShardedIterDomains function. Added getShardedLoopAxis function. +15/-10  reduction_utils.cpp Update scheduleReductionTV for DID loop split                        csrc/scheduler/reduction_utils.cpp Updated scheduleReductionTV to use getShardedLoopAxis. Added error checks for DID loop split. +11/-4    utils.h Declare getShardedLoopAxis                                                              csrc/multidevice/utils.h Added declaration for getShardedLoopAxis. +4/-0
Tests	test_multidevice_sharding.cpp Add DivideBySum test case                                                                tests/cpp/test_multidevice_sharding.cpp Added DivideBySum test case for multidevice sharding. +42/-0

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Assumption Check The PR assumes that the DIDx domain is always the outermost domain in the loop. This assumption should be validated with more test cases to ensure correctness. int64_t sharded_axis = getShardedLoopAxis(reduction_tv, ParallelType::DIDx); if (sharded_axis >= 0) { NVF_ERROR( sharded_axis == 0, "Expect 1D mesh and DIDx only appear outermost in loop, but found: ", reduction_tv->getLoopDomain()); } Error Handling The error handling in getShardedLoopAxis could be improved by providing more context in the error message, such as the specific parallel type that caused the failure. isParallelTypeDeviceDim(parallel_type), "Expect a DID but found: ", parallel_type); Test Coverage While a new test DivideBySum is added, it would be beneficial to add more test cases to cover different scenarios and edge cases for the new functionality. TEST_F(MultiDeviceTest, DivideBySum) { auto fusion = std::make_unique<Fusion>(); FusionGuard fg(fusion.get()); const int64_t d = communicator_->size(); // [b, h, s, s] TensorView* x = makeContigTensor(4); TensorView* sum_x = sum(x, {-1}); TensorView* sum_x_broadcasted = broadcast(sum_x, {false, false, false, true}); TensorView* y = div(x, sum_x_broadcasted); fusion->addInput(x); fusion->addOutput(y); auto mesh = DeviceMesh::createForNumDevices(d); for (auto* tv : {x, sum_x, sum_x_broadcasted, y}) { tv->setDeviceMesh(mesh); tv->split(1, d, /*inner_split=*/false); tv->axis(1)->parallelize(ParallelType::DIDx); tv->reorder({{1, 0}}); } for (auto* tv : {x, y}) { tv->setAllocationDomain(tv->getLoopDomain(), true); } const int64_t b = 2; const int64_t h = d * 3; const int64_t s = 5; at::Tensor unsharded_x_tensor = at::randint(5, {b, h, s, s}, tensor_options); at::Tensor x_tensor = shardTensor(unsharded_x_tensor, x); FusionExecutorCache executor_cache(std::move(fusion)); at::Tensor y_tensor = executor_cache.runFusionWithInputs({x_tensor})[0]; testValidate( executor_cache.fusion(), {y_tensor}, {x_tensor}, {x_tensor / x_tensor.sum(-1, true)}, __LINE__, __FILE__); }

[wujingyue]

!test

[wujingyue]

!test

[naoyam]

LGTM

[wujingyue]

!test

[Priya2698]

LGTM.

[wujingyue]

CI failures are due to http://nv/exg