JuliaGPU / JuliaGPU/KernelAbstractions.jl

How can we make KA fast on CPUs?

Open
#509 12 comments 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
523
Forks
88
Avg merge
1d 11h
Merged PRs (30d)
25

Description

See https://github.com/LuxDL/LuxLib.jl/pull/136 for some background context. The main motivation for me is to avoid code duplication between CPU and GPU versions. However, if you take a look at the benchmark comment on the PR (for batchnorm and groupnorm) you see somewhere between a 10x-40x slowdown between KA and the equivalent optimized loop version (note that it is simply using @simd or @simd ivdep and nothing like LoopVectorization).

I think there are a couple of reasons for the slowdown:

  1. @simd annotations are missing (which causes slowdown even in the loop version if I remove the annotations)
  2. threading has overhead for some of the smaller problems

Potential solutions:

  1. Allow users to control threading. #507. For smaller problems, I want to opt out of threading manually.
  2. @simd annotations (#436 seems to do this. not sure what is the status for that)
  3. Alternate threading: KA is being used inside "core" operations. As such we are unlikely (if not impossible) to call other operations that make use of threading. Hence, having the option to use "cheaper threads" (Polyester.jl) would be a great addition

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reading PR #136 and its benchmark comment for batchnorm and groupnorm, then review the related issues #507 and #436. Compare the proposed threading controls, SIMD annotations, and Polyester.jl option before narrowing the scope. Done should be an agreed CPU-performance change validated against the cited benchmark results.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.