zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-19 23:51:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix get 2d grid dims (#2316)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-06-25 13:03:09 -07:00 committed by GitHub
parent 81bb9a2a9e
commit 656ed7f780
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mlx/backend/common/utils.cpp
View File

@ -199,12 +199,15 @@ Dims get_2d_grid_dims_common(
} }
} }
} }
if (grid_y > UINT32_MAX || grid_x > UINT32_MAX || divisor > 1) { if (grid_y > UINT32_MAX || grid_x > UINT32_MAX) {
throw std::runtime_error("Unable to safely factor shape."); throw std::runtime_error("Unable to safely factor shape.");
} }
if (grid_y > grid_x) { if (grid_y > grid_x) {
std::swap(grid_x, grid_y); std::swap(grid_x, grid_y);
} }
if (divisor > 1) {
grid_x = ((grid_x + divisor - 1) / divisor) * divisor;
}
return std::make_tuple( return std::make_tuple(
static_cast<uint32_t>(grid_x), static_cast<uint32_t>(grid_y), 1); static_cast<uint32_t>(grid_x), static_cast<uint32_t>(grid_y), 1);
} }