squall/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix grouped gemm invalid memory access to problem shapes (#1543)

Browse Source
This commit is contained in:
Chengquan Jiang 2024-07-10 23:55:22 +08:00 committed by GitHub
parent 52fb43f30f
commit 56b46e2d13
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
include/cutlass/gemm/kernel/sm90_gemm_array_tma_warpspecialized_cooperative.hpp
View File

@ -465,6 +465,9 @@ public:
TileScheduler scheduler{params.scheduler}; TileScheduler scheduler{params.scheduler};
auto work_tile_info = scheduler.get_current_work(); auto work_tile_info = scheduler.get_current_work();
if (not work_tile_info.is_valid()) {
return;
}
// Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK) // Optionally append 1s until problem shape is rank-4 in case it is only rank-3 (MNK)
auto problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{}); auto problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
@ -623,8 +626,10 @@ public:
// Get next work tile // Get next work tile
work_tile_info = fetch_next_work(work_tile_info, scheduler); work_tile_info = fetch_next_work(work_tile_info, scheduler);
if constexpr (IsGroupedGemmKernel) { if constexpr (IsGroupedGemmKernel) {
if (work_tile_info.is_valid()) {
problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{}); problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
} }
}
} // Scheduler work fetch loop } // Scheduler work fetch loop
// Make sure all Consumer Warp Groups have been waited upon // Make sure all Consumer Warp Groups have been waited upon
@ -702,8 +707,10 @@ public:
// Get next work tile // Get next work tile
work_tile_info = fetch_next_work(work_tile_info, scheduler); work_tile_info = fetch_next_work(work_tile_info, scheduler);
if constexpr (IsGroupedGemmKernel) { if constexpr (IsGroupedGemmKernel) {
if (work_tile_info.is_valid()) {
problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{}); problem_shape_MNKL = append<4>(params.problem_shape.get_problem_shape(work_tile_info.L_idx), Int<1>{});
} }
}
} // Scheduler work fetch loop } // Scheduler work fetch loop
if (do_store_tail) { if (do_store_tail) {