Skip to content

vulkan: optimize operations in the IM2COL shader #22685

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
0cc4m merged 2 commits into from
May 20, 2026
Merged

vulkan: optimize operations in the IM2COL shader #22685

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c6860fd
vulkan: optimize operations in the IM2COL shader
daniandtheweb May 4, 2026
52336b2
Add comments and improve the code formatting
daniandtheweb May 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 59 additions & 14 deletions ggml/src/ggml-vulkan/vulkan-shaders/im2col.comp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,36 +44,81 @@ void im2col(const uint ow, const uint z_idx) {

const uint KHKW = p.KH * p.KW;

// Precompute base input coordinates
const int base_iw = int(ow * p.s0) - p.p0;
const int base_ih = int(oh * p.s1) - p.p1;

// Precompute step deltas
const uint delta_ic = BLOCK_SIZE / KHKW;
const uint delta_rem = BLOCK_SIZE % KHKW;

const uint delta_ky = delta_rem / p.KW;
const uint delta_kx = delta_rem % p.KW;

@jeffbolznv jeffbolznv May 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not totally following this. In general it seems unsafe to precompute divs/mods and add them, as sometimes you would wrap to the next value and need to do a fixup. Maybe that's what the fixup logic is doing, but it's not clear.

I wonder if it might be better to pass KW as a spec constant and let the compiler transform it into something faster.

@daniandtheweb daniandtheweb May 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

kx_wrap and ky_wrap should take care of the wrapping. Moreover it shouldn't be possible for the values to wrap around twice as:

  • delta_kx by using a modulo is always less than p.KW, kx is always less than p.KW because of the modulo and the maximum value that it can reach is 2 * p.KW - 2 which is always less than 2 * p.KW.
  • delta_ky's max value is p.KH - 1 so always less than p.KH, ky is at most p.KH - 1 and the maximum value that it can reach is 2 * p.KH - 1 so less than 2 * p.KH.

@daniandtheweb daniandtheweb May 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not that experienced on vulkan shaders so the most I could confidently achieve was this as it should (in theory) be mathematically correct. If there are better approaches I'll gladly look into them.

@daniandtheweb daniandtheweb May 4, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe it's better if I add some comments on the wrap values to make it more clear?

@jeffbolznv jeffbolznv May 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Comments would help, but I think spec constants would keep the code more clear. I'll defer to @0cc4m on what to do.

@daniandtheweb daniandtheweb May 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Got it, I'll start by adding some comments on the most confusing parts for now.

@0cc4m 0cc4m May 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think a spec constant version might be easier to read on the shader side, but would add complexity on the host side. Either is fine with me, but we already have this now, so I think it is okay to keep it.

@0cc4m 0cc4m May 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jeffbolznv Any concerns? Otherwise I'll merge it.

@jeffbolznv jeffbolznv May 20, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm OK with it.


const uint delta_ic_offset = delta_ic * p.offset_delta;

// If using BDA mode, precompute the base pointer and step size
#if BDA
const BDA_STORAGE_T base_dst_addr = p.dst_addr + D_SIZE * dst_row;
const uint bda_step = D_SIZE * BLOCK_SIZE;
#endif

uint wg_x = gl_WorkGroupID.x;
do {
const uint wg_offset = wg_x * 512;

[[unroll]] for (uint i = 0; i < NUM_ITER; ++i) {
const uint chw_idx = wg_offset + gidx + i * BLOCK_SIZE;
uint chw_idx = wg_offset + gidx;

uint ic = chw_idx / KHKW;
uint rem = chw_idx % KHKW;

uint ky = rem / p.KW;
uint kx = rem % p.KW;

uint ic_offset = src_batch + ic * p.offset_delta;

// Initialize running pointer/index for the destination buffer
#if BDA
BDA_STORAGE_T current_dst_addr = base_dst_addr + D_SIZE * chw_idx;
#else
uint current_dst_idx = dst_row + chw_idx;
#endif

[[unroll]] for (uint i = 0; i < NUM_ITER; ++i) {
if (chw_idx >= p.CHW) {
return;
}

const uint ic = chw_idx / KHKW;
const uint rem = chw_idx - ic * KHKW;
const uint ky = rem / p.KW;
const uint kx = rem - ky * p.KW;

const uint iiw = ow * p.s0 + kx * p.d0 - p.p0;
const uint iih = oh * p.s1 + ky * p.d1 - p.p1;
const int iiw = base_iw + int(kx * p.d0);
const int iih = base_ih + int(ky * p.d1);

A_TYPE val = A_TYPE(0);
if (iih < p.IH && iiw < p.IW) {
val = data_a[src_batch + ic * p.offset_delta + iih * p.IW + iiw];
if (uint(iih) < p.IH && uint(iiw) < p.IW) {
val = data_a[ic_offset + uint(iih) * p.IW + uint(iiw)];
}

#if BDA
D_ptr out_ptr = D_ptr(p.dst_addr + D_SIZE * (dst_row + chw_idx));
out_ptr.d = D_TYPE(val);
D_ptr(current_dst_addr).d = D_TYPE(val);
current_dst_addr += bda_step;
#else
data_d[dst_row + chw_idx] = D_TYPE(val);
data_d[current_dst_idx] = D_TYPE(val);
current_dst_idx += BLOCK_SIZE;
#endif

chw_idx += BLOCK_SIZE;
ic_offset += delta_ic_offset;
kx += delta_kx;
ky += delta_ky;

// Handle X axis wrap
uint kx_wrap = uint(kx >= p.KW);
kx -= kx_wrap * p.KW;
ky += kx_wrap;

// Handle Y axis wrap
uint ky_wrap = uint(ky >= p.KH);
ky -= ky_wrap * p.KH;
ic_offset += ky_wrap * p.offset_delta;
}

wg_x += gl_NumWorkGroups.x;
Expand Down
Loading