Skip to content

fix metal conv bounds checks and integer overflows - #4672

Open
Prudctual wants to merge 1 commit into
ml-explore:mainfrom
Prudctual:fix/metal-conv-bounds-and-overflows
Open

Prudctual wants to merge 1 commit into
ml-explore:mainfrom
Prudctual:fix/metal-conv-bounds-and-overflows

Conversation

@Prudctual

Copy link
Copy Markdown
Contributor
  • I understand it is strictly prohibited to use AI to write PR description

fixes #4666, fixes #4664.

this fixes several bounds checks and integer overflows in the metal conv path:

  1. steel_conv_general.h: widened short diff = gemm_params->N - offset_n to int diff. with more than 32767 output channels per group, diff wrapped negative and caused the store guard to evaluate false, leaving upper channels unwritten on gpu.
  2. loader_channel_n.h and loader_channel_l.h: bounded weight rows against per-group channels gemm_params->N instead of total channels params->O. because the weight base pointer is already offset to the group, comparing against total channels caused grouped and depthwise convs (e.g. conv1d with padding) to read past the end of the weight buffer.
  3. loader_channel_l.h and loader_general.h: widened short weight_h, weight_w, weight_d spatial counters to int so large kernel spatial dimensions do not wrap.
  4. conv.metal: evaluated winograd weight transform destination offset with size_t(ohw_0) * C * O to prevent 32-bit overflow when C * O is large.
  5. conv.cpp: widened (static_cast<int64_t>(implicit_M) + bm - 1) / bm to prevent signed 32-bit overflow.

tested on m-series mac with metal shader validation and the issue reproducers:

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

1 participant