Resize scheduler: enable vectorization

[naoyam]

Stacked on #3693

This PR adds a preliminary vectorization support to the resize scheduler. It currently only considers vectorization of the innermost dimension, just because that's good enough for the RoPE cases. It should eventually be extended to support vectorizing multiple innermost dimensions.

[naoyam]

!test

[jjsjann123]

LGTM

[naoyam]

@jjsjann123 Forgot to include this WAR. This is what I mentioned to you that I saw an error due to #3640.

[naoyam]

!test

[naoyam]

!test

[naoyam]

!test

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

⚡ Recommended focus areas for review Logic Change The ResizeScheduler class has been modified to support vectorization. The computeHeuristics function now calculates a vectorization_factor based on the innermost dimension of the largest input. The schedule function has been updated to apply vectorization to the innermost dimension of the reference tensor view. This change may have performance implications and should be reviewed carefully. auto ref_tv_entry = HeuristicDataCacheEntry<HeuristicCompileTime::ReferenceTensors>( data_cache, [fusion]() { std::vector<TensorView*> data{getReferenceTensor(fusion)}; return std::make_unique<std::vector<TensorView*>>(std::move(data)); }); TensorView* ref_tv = ref_tv_entry.get()[0]; // Before applying the vectorization split, any reshape transform of // the largest input will be cancelled whenever possible, so the // largest input is used as the reference of vectorization. auto vec_ref_tv = largest_input != nullptr ? largest_input : ref_tv; // Only consider the innermost dimension to vectorize for now. // TODO: Consider vectorizing merged IDs, not just the innermost params->vectorization_factor = vectorize_helper::getVectorizationFactor( runtime_info, vec_ref_tv, data_cache, (int64_t)vec_ref_tv->getLogicalDomain().size() - 1, {}); return params; Potential Bug In the schedule function, the next_innermost_pos variable is used to keep track of the next innermost position after splitting the reference tensor view. However, if the vectorization_factor is greater than 1, the next_innermost_pos variable is decremented twice, which may lead to incorrect results. int64_t next_innermost_pos = -1; // [..., ...] // ^ // +--- next_innermost_pos if (vec_factor > 1) { ref_tv->split(-1, vec_factor); --next_innermost_pos; // [..., vec_factor] // ^ // +--- next_innermost_pos } Test Modification The SliceRotateCatResidual test has been modified to use a shape that is divisible by 8 to avoid an error due to issue #3640. This change may not be related to the vectorization support added in this PR, but it should be reviewed to ensure that the test is still valid. // Due to #3640, the vectorization analysis may return 4 for this // fusion since there's the use of the input without // slicing. However, the correct factor needs to consider the // slicing paths as well. For now, in order to avoid the error due // to issue #3640, use a size that is divisible by 8. // std::vector<int64_t> shape({16, 100}); std::vector<int64_t> shape({16, 96});

[naoyam]

!test