Fix the transpose scheduler for DID loop split.

[wujingyue]

For #2563

[wujingyue]

!test

[github-actions]

Review updated until commit c43d5e1

Description

• Rename getShapeInReference to getLoopDomainSizes

• Update loop domain handling in transpose scheduler

• Improve error handling in vectorization factor calculation

• Fix tensor view string representation and allocation domain setting

---

Changes walkthrough 📝

	Relevant files
Enhancement	transpose.cpp Update transpose scheduler functions                                          csrc/scheduler/transpose.cpp Renamed getShapeInReference to getLoopDomainSizes Updated loop domain handling in multiple functions Improved error handling in getVectorizationFactorTransposeGroup +16/-15  test_multidevice_sharding.cpp Update transpose test case                                                              tests/cpp/test_multidevice_sharding.cpp Updated Transpose test to use symbolic tensors Corrected tensor splitting and parallelization Simplified expected output validation +12/-12
Bug fix	vectorize_helper.cpp Improve vectorization factor calculation                                  csrc/scheduler/vectorize_helper.cpp Added ir/printer.h include Updated getVectorizationFactorTransposeGroup to handle parallelized axes +12/-4    tensor_view.cpp Fix tensor view string representation                                        csrc/tensor_view.cpp Updated toString method to use indent Corrected allocation domain setting in toString +2/-2

---

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🧪 PR contains tests

⚡ Recommended focus areas for review Function Renaming The function getShapeInReference has been renamed to getLoopDomainSizes. Ensure that this change is intentional and that all references to getShapeInReference have been updated accordingly. std::pair<std::vector<int64_t>, int64_t> getLoopDomainSizes( HeuristicDataCache* data_cache, SchedulerRuntimeInfo& runtime_info, TensorView* reference, scheduler_tools::TransposeDomainMap& domain_map) { auto ref_loop = reference->getLoopDomain(); Vectorization Factor Check The code now includes a check to ensure that the innermost dimension is not parallelized before vectorization. Verify that this check is necessary and that it does not introduce any unintended behavior. } ///////////////////////////// // Step 2: global schedule // ///////////////////////////// // make tile Test Case Update The test case for transpose has been updated to use a symbolic tensor and specific split operations. Ensure that these changes do not alter the intended behavior of the test and that it still effectively validates the transpose scheduler. TEST_F(MultiDeviceTest, Transpose) { auto fusion = std::make_unique<Fusion>(); FusionGuard fg(fusion.get()); const auto num_devices = communicator_->size(); auto mesh = DeviceMesh::createForNumDevices(num_devices); TensorView* in = makeSymbolicTensor(2); TensorView* out = transpose(in, 0, 1); in->split(0, num_devices, /*inner_split=*/false); in->axis(0)->parallelize(ParallelType::DIDx); out->split(1, num_devices, /*inner_split=*/false); out->axis(1)->parallelize(ParallelType::DIDx); out->reorder({1, 0}); for (auto* tv : {in, out}) { tv->setDeviceMesh(mesh); tv->setAllocationDomain(tv->getLoopDomain(), true); } fusion->addInput(in);

[wujingyue]

!test

[wujingyue]

!test --diff

[wujingyue]

The codegen diff is expected -- this PR changed the test.

[naoyam]

Is this change a general bug fix not specific to DID?

[wujingyue]

I think it's a bug triggered by DID loop split. According to #3927 (comment), inner_most_pos1_in_ref1 and inner_most_pos2_in_ref2 ought to be loop axis but is misused to access getLogicalDomain() in getVectorizationFactorTransposeGroup. Therefore, I also made this change: https://github.com/NVIDIA/Fuser/pull/3927/files#diff-364f51b47b1cf80c14b75d95ac2882238a65aa0594d588cff2fc54526307ca55L941

[naoyam]

I see. Looks like it was fine before because logical == loop. The change here makes sense.

[naoyam]

The change looks good to me. I'll let @zasdfgbnm to confirm.

@zasdfgbnm Could you take a look as well?