Fix rfactor replay for DID loop split

[wujingyue]

For #2563

[naoyam]

What does the actual fusion that is passed to the reduction scheduler look like?

[github-actions]

Review updated until commit f45bf68

Description

• Update rfactor replay logic for DID loop split

• Enhance test coverage with dynamic sizes

• Correct tutorial comments and remove redundant checks

• Improve allreduce test with dynamic tensor sizes

---

Changes walkthrough 📝

	Relevant files
Bug fix	reduction.cpp Update 3D reduction error check                                                    csrc/scheduler/reduction.cpp Update error check for 3D reduction schedules with DIDx sharding +2/-1      tensor_view.cpp Remove redundant rfactor check                                                      csrc/tensor_view.cpp Remove redundant check for rfactor on the same view +0/-6
Enhancement	transform_rfactor.cpp Simplify IterType assignment                                                          csrc/transform_rfactor.cpp Simplify IterType assignment in ReplayRFactor +2/-2      test_communication.py Enhance allreduce test                                                                      tests/python/test_communication.py Enhance allreduce test with dynamic tensor sizes and rfactor Add device mesh and parallelization for rfactor output +18/-4
Documentation	test_tutorial.cpp Update tutorial comments                                                                  tests/cpp/test_tutorial.cpp Update comments and print method in reduction rfactor test +4/-4

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Possible Issue The new error check might be too strict. The original check was isSharded(reduction_tv), which checks if the tensor view is sharded. The new check getShardedLoopAxis(reduction_tv, ParallelType::DIDx) >= 0 checks for a specific sharded loop axis. This might not cover all cases where the tensor view is sharded. !(rparams->schedule_3D && getShardedLoopAxis(reduction_tv, ParallelType::DIDx) >= 0), Removed Check The check for domain()->hasRoot() was removed. This check ensures that rfactor is not called on the same view twice, which might be important for maintaining the correctness of the tensor view transformations. definition() != nullptr && (definition() ->isStrictlyOneOf<ReductionOp, MmaOp, MatmulOp, LinearOp>()), "Error rfactoring ", this, " its definition is either a nullptr or not a reduction."); NVF_CHECK( !definition()->isA<GroupedReductionOp>(), "For GroupedReductionOp, use TensorView::rFactor(const std::vector<int64_t>& axes, const std::vector<TensorView*>& tvs)"); Hardcoded Values The test test_allreduce uses hardcoded values for the tensor dimensions (m = d * 2, n = 3). This might limit the test's ability to catch issues with different input sizes. Consider using more dynamic or parameterized test cases. def test_allreduce(multidevice_test): d = multidevice_test.size mesh = nvfuser.DeviceMesh(range(d)) class Model(FusionDefinition): def definition(self): self.inp = self.define_tensor( (-1, -1), contiguity=True, dtype=DataType.Float ) self.out = self.ops.sum(self.inp, [0]) self.add_output(self.out) def multidevice_schedule(self): self.sched.split(self.inp, 0, d, False) self.sched.split(self.out, 0, d, False) out_local = self.sched.rfactor(self.out, [1]) self.sched._set_device_mesh(self.inp, mesh) self.sched._set_device_mesh(self.out, mesh) self.sched._set_device_mesh(out_local, mesh) self.sched.parallelize(self.inp, 0, nvfuser.ParallelType.mesh_x) self.sched.parallelize(out_local, 0, nvfuser.ParallelType.mesh_x) self.sched.set_allocation_as_loop(self.inp) self.sched.set_allocation_as_loop(out_local) self.sched.set_allocation_as_loop(self.out) m = d * 2 n = 3 unsharded = torch.randn(m, n) sharded = multidevice_test.shard_tensor(unsharded, 0, mesh) fd = Model() outputs = fd.execute([sharded]) torch.testing.assert_close(outputs[0].local.cpu(), unsharded.sum(0))

[wujingyue]

!test

[wujingyue]

I removed this check because we now expect rFactor to be called by both inter- and intra-GPU schedulers.

[wujingyue]

Without this, I ran into an error with the following local reduction:

in: root/logical=[i{n}], loop=[iDIDx{d}, i{n/d}]
out = reduction(in): root=[r{n}], logical/loop=[iDIDx{d}, r{n/d}]

The reduction scheduler tries to schedule out on TIDx

out: root=[r{n}], logical=[iDIDx{d}, r{n/d}], loop=[iDIDx{d}, r{n/d/blockDim.x}, rTIDx{blockDim.x}]

and then rFactor axis 1, i.e., r{n/d/blockDim.x}.

rFactor tries to replay all transforms using ReplayRFactor on a new, identical root domain [r{n}]. Without this change, the outer-split by d produced rDIDx{d} instead of iDIDx{d}.

[wujingyue]

!test

[wujingyue]

What does the actual fusion that is passed to the reduction scheduler look like?

I finally debugged this through. PTAL!