Skip to content

Update Metal reduce kernels to remove intermediate fill! - #429

Open
PhilipFackler wants to merge 2 commits into
JuliaGPU:mainfrom
PhilipFackler:improve-metal-reduce
Open

Update Metal reduce kernels to remove intermediate fill!#429
PhilipFackler wants to merge 2 commits into
JuliaGPU:mainfrom
PhilipFackler:improve-metal-reduce

Conversation

@PhilipFackler

@PhilipFackler PhilipFackler commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Closes #420

@PhilipFackler

PhilipFackler commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@williamfgc

williamfgc commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

@PhilipFackler I am checking the GPU CI for failures

@williamfgc

williamfgc commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Test this please

@williamfgc

williamfgc commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

@PhilipFackler please check the CI failure on macOS. It's related to the number of threads.

@williamfgc

williamfgc commented Aug 27, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

@PhilipFackler I did some digging,

  1. ArgumentError: Number of threads in group (1024) should not exceed 896 happens due to Metal, see here: or Metal may lower the thread limit when creating the compute pipeline state, which can reduce runtime performance.
  2. The tests pass on M3 GPUs as-is
  3. You can reduce the max number of dimensions from 5 to 4 in the reduce-ND test for our M1 CI as it works.
  4. We're thinking to add M3 GPU to CI later and how to distinguish M1 from M3.

@PhilipFackler

PhilipFackler commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

williamfgc
williamfgc requested changes Aug 28, 2026
View reviewed changes
Comment thread ext/AMDGPUExt/AMDGPUExt.jl
Comment thread ext/MetalExt/MetalExt.jl
Comment thread ext/MetalExt/MetalExt.jl Outdated
@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

2 similar comments
@PhilipFackler

PhilipFackler commented Aug 28, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

@PhilipFackler

PhilipFackler commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

This was referenced Sep 13, 2026
Merged
Open
christiangnrd
christiangnrd reviewed Sep 18, 2026
View reviewed changes
Comment thread ext/MetalExt/MetalExt.jl

@inline function _kernel_maxthreads(kernel_function, kargs, kname)
p_kernel = _make_kernel(kernel_function, kargs, kname)
return (p_kernel, p_kernel.maxthreads)

@christiangnrd christiangnrd Sep 18, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

This is the correct way to do this since the maximum number of threads is set per pipeline, but you'll have to raise Metal.jl compat to 1.10.

You could also use p_kernel.pipeline.maxThreadsPerThreadgroup to avoid raising compat but that might add some overhead

@PhilipFackler
Update Metal maxthreads
e7d5367 
Update barrier call

Fix Metal launch
@PhilipFackler

PhilipFackler commented Sep 24, 2026

Copy link
Copy Markdown
Collaborator Author

Test this please

williamfgc
williamfgc requested changes Sep 24, 2026
View reviewed changes

@williamfgc williamfgc left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Nice job! Let's see what CI says. Just minor request to comment on oneAPI or do a separate PR. Up to you.

Comment thread ext/MetalExt/MetalExt.jl
max_pwr = JACC.ilog2(shmem_length) - 1
for p in (max_pwr:-1:0)
threadgroup_barrier()
threadgroup_barrier(Metal.MemoryFlagThreadGroup)

@williamfgc williamfgc Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Good one!

Comment thread ext/oneAPIExt/oneAPIExt.jl
kernel1 = _make_kernel(_make_kname(name, "block_reduce"),
_parallel_reduce_oneapi, Val(256), N, op, ret_inst, init, f, x...)
threads1 = oneAPI.launch_configuration(kernel1)

@williamfgc williamfgc Sep 24, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Either a separate PR or modify this PR's title :)

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

Reviewers

@williamfgc williamfgc williamfgc requested changes
+1 more reviewer
@christiangnrd christiangnrd christiangnrd left review comments
Reviewers whose approvals may not affect merge requirements

Requested changes must be addressed to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Remove per-reduce fill! for Metal and oneAPI backends

3 participants

@PhilipFackler @williamfgc @christiangnrd