Enable same padding for non-register sharing warp specialization

[rdspring1]

This PR enforces same padding rules for non-register sharing warp specialization.

• Replaced std::unordered_set<ParallelType> warp_specialized_types_ with std::optional<ParallelType> warp_specialized_parallel_type_ because we only support a single ParallelType.

[github-actions]

Review updated until commit bc4e084

Description

• Replaced std::unordered_set<ParallelType> with std::optional<ParallelType> for warp specialization.

• Added checks for multiple warp specialized axes and updated padding logic.

• Updated test cases to reflect changes in warp specialization handling.

---

Changes walkthrough 📝

	Relevant files
Enhancement	parallel_dimension_map.cpp Refactor warp specialization handling                                        csrc/parallel_dimension_map.cpp Replaced warp_specialized_types_ with warp_specialized_parallel_type_. Added error handling for multiple warp specialized axes. Updated padding logic and variable names. +19/-37  parallel_dimension_map.h Update header file for warp specialization                              csrc/parallel_dimension_map.h Updated hasWarpSpecialization and added isWarpSpecialized method. Renamed variables related to warp specialization padding. +9/-8
Tests	test_circular_buffering.cpp Update test cases for warp specialization                                tests/cpp/test_circular_buffering.cpp Added testEnablesWarpSpecialization and testEnablesTIDx methods. Updated test skips to check for warp specialization. +16/-6    test_combined_inner_outer_reduction.cpp Update test cases for warp specialization                                tests/cpp/test_combined_inner_outer_reduction.cpp Updated test skips to check for warp specialization. +13/-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 check if (!warp_specialized_parallel_type_.has_value()) in getWarpSpecializationPaddedVal might return 1, which could be incorrect if no warp specialization is intended. Ensure this behavior is correct or handle the case appropriately. return 1; } Test Coverage The tests skip scenarios where warp specialization is enabled. Verify that these scenarios are covered elsewhere or that the skips are justified. TEST_P(TmaCircularBufferingTest, Persistent) { NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0); if (testEnablesWarpSpecialization()) { GTEST_SKIP() << "Bdimx is dynamic, Warp Specialization is disabled."; return; Test Coverage The tests skip scenarios where warp specialization is enabled. Verify that these scenarios are covered elsewhere or that the skips are justified. TEST_P(TmaWarpSpecializedTest, SimpleFusion) { NVFUSER_TEST_CUDA_ARCH_GUARD(9, 0); auto [contig, ws_enabled, dtype, dim0, dim1] = GetParam(); if (ws_enabled) { GTEST_SKIP() << "Bdimx is dynamic, Warp Specialization is disabled."; return; }

[rdspring1]

!test

[zasdfgbnm]

Nice cleanup!

[rdspring1]

!test