DID parallelization of loop domain

[cowanmeg]

Currently, DID parallelization is only supported on the logical domain which requires us to manually add extra axes in the logical domain when parallelizing Fusions across multiple devices.

This has caused some PRs to introduce temporary hacks to work around logical DID axes or code that needs to be modified which is tracked here.

FYR, an internal design doc from @cowanmeg

[wujingyue]

As we discussed in the meeting, here's a breakdown of the task:

• Bind sharded input/output tensors with DID-parallelized allocation domains.  #3282. ETC: 3 weeks. Assigned to @wujingyue
• Fix communication lowering to support DID split. ETC: 2 weeks. Assigned to @wujingyue
• Fix ATen-evaluated nodes, e.g. linear, matmul and sdpa. ETC: 2 weeks
• Fix all schedulers. According to the nsys traces @cowanmeg collected, we'll need to fix at least pointwise, *_persistent (aka normalization), transpose and reduction. FWIW, forward needs pointwise and persistent only. ETC: 1 month
• Fix sharding propagation to propagate DID splits, e.g., using TransformPropagator. ETC: 1 month
• Fix other sharding-related pre-segmenter passes, e.g., insert_reshardings and reorder_sharded_axis. One example is to decompose a DID-involving reduction (e.g. Reduction, Matmul and Linear) into multiple ops. This is conceptually similar to the existing reduction rfactoring. ETC: 1 month
• Remove tests that rely on logical split of DID and clean up dead code. ETC: 2 weeks

These subtasks are roughly execution ordered. #3282 pretty much blocks everything else, so it's currently my top priority. After that's done, some other tasks (e.g. 2-4) can be parallelized as long as we limit the scope of testing.

cc @naoyam, @cowanmeg and @samnordmann to take a look

[naoyam]

It's probably necessary to update the allocation logic of gmem tensors. Since we always allocate the whole tensor for a single GPU execution, the logical domain is used to determine the allocation size, which would need to be changed to the loop domain.

This is actually just one example of loop vs logical. I could imagine there would be some other cases as well.

[cowanmeg]

Yes, I actually experimented a bit with that in this very old and out of date PR, #2181
When we explicitly set the allocation domain of the sharded tensors to the loop domain (before the schedulers applied intra-device parallelization) allocation worked.

[wujingyue]

It's probably necessary to update the allocation logic of gmem tensors.

I believe it has been fixed:

Fuser/csrc/runtime/allocations.cpp

Lines 737 to 746 in 85c22a2

	for (const auto id : tv->getMaybeAllocationDomain()) {
	if (id->isReduction() || id->isStride()) {
	continue;
	}
	if (id->isDeviceDim()) {
	symbolic_sizes.push_back(id->container()->oneVal());
	} else {
	symbolic_sizes.push_back(id->getMaybeExpandedExtent());
	}

. Anyhow, even if it hasn't, we will be enforced to fix it along the way of fixing the subtasks on the list.

[samnordmann]

Thanks for creating this issue so we can keep track!

the logical domain is used to determine the allocation size, which would need to be changed to the loop domain.
This is actually just one example of loop vs logical. I could imagine there would be some other cases as well.

@naoyam agreed!
If using split/merge on sharded axis, we will need some operations/logic to act on the loop domain rather than on the logical domain.