Skip to content
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.

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

An experimental kernel dispatcher for numba_dpex.kernel decorator #1178

Merged
merged 16 commits into from
Nov 3, 2023

Conversation

diptorupd
Copy link
Collaborator

@diptorupd diptorupd commented Oct 18, 2023

  • Have you provided a meaningful PR description?

Numba-dpex uses the numba_dpex.core.kernel_interface.JitKernel class a custom dispatcher class for running kernel functions. The JitKernel dispatcher has a different API than numba.dispatcher.Dispatcher making it hard to use JitKernel as a target dispatcher inside Numba. The limitation means we currently cannot do the following:

  • add overloads in the DpexKernelTarget
  • call kernels from numba_dpex.dpjit
  • use Numba's unbox/box APIs (numba-dpex uses ctypes to implement the functionality)

There are also potential performance overheads as kernel dispatching fully happens in Python.

The PR introduces a new KernelDispatcher class that is a sub-class of numba.core.dispatcher.Dispatcher. The KernelDispatcher is for now included as an experimental feature, as to be useful all existing lower functions in numba_dpex.ocl.oclimpl would have to be added to DpexKernelTargetContext. The overloads will be done as a follow up PR.

The experimental KernelDispacther compiles a @kernel decorated function to SPIRV and stores it as an overload.

To call an experimental kernel, a new call_kernel function has been added to numba_dpex.experimental. The function is decorated using dpjit and can be called from both CPython or another dpjit function.

The call_kernel function does the following:
- create an array to store the kernel arguments,
- create an array to store the type ids for kernel arguments,
- compile the SPIR-V binary to a sycl::kernel_bundle.
- Extract the sycl::kernel from the sycl::kernel_bundle
- Submit the kernel to the sycl::queue extracted from a DpnpNdArray kernel argument
- Wait on the queue

  • Have you added a test, reproducer or referred to an issue with a reproducer?
  • Have you tested your changes locally for CPU and GPU devices?
  • Have you made sure that new changes do not introduce compiler warnings?
  • If this PR is a work in progress, are you filing the PR as a draft?

@diptorupd diptorupd changed the title An experimental kernel dispatcher that generates a native LLVM wrapper An experimental kernel dispatcher for numba_dpex.kernel decorator Oct 18, 2023
@diptorupd diptorupd self-assigned this Oct 18, 2023
@diptorupd diptorupd marked this pull request as draft October 18, 2023 17:00
@diptorupd
Copy link
Collaborator Author

diptorupd commented Oct 18, 2023

TODOs

  • Add a new debug option to generate printf calls from inside launcher function.
  • Port the NUMBA_DPEX_DUMP_KERNEL_LLVM config option to experimental KernelDispatcher
  • Add a new config option to dump the LLVM IR for the optimized launcher and cpython wrapper functions
  • Why is the compilation failing when NUMBA_CAPTURED_ERRORS is changed to new_style
  • Check if overloads are stored based on device and not queue
  • Move queue equivalence irrespective of whether kernel overload exists or not.

@diptorupd diptorupd force-pushed the feature/KernelDispatcher branch 4 times, most recently from 81f631a to 52b356a Compare October 24, 2023 17:40
@diptorupd diptorupd marked this pull request as ready for review October 24, 2023 17:45
Copy link
Collaborator

@ZzEeKkAa ZzEeKkAa left a comment

Choose a reason for hiding this comment

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

Changes look awesome!
It takes some time for me to understand what's going, so my review is still in progress.

numba_dpex/experimental/launcher.py Outdated Show resolved Hide resolved
numba_dpex/experimental/kernel_dispatcher.py Show resolved Hide resolved
numba_dpex/experimental/kernel_dispatcher.py Outdated Show resolved Hide resolved
@diptorupd diptorupd force-pushed the feature/KernelDispatcher branch 3 times, most recently from 6656dcd to b14fdec Compare October 26, 2023 04:07
@ZzEeKkAa ZzEeKkAa force-pushed the feature/KernelDispatcher branch 3 times, most recently from edd7c04 to aad1184 Compare October 26, 2023 16:07
    - The helper function was renamed and can now optionally
      return a DpctlSyclEventRef object to allow waiting at callsite.
    - Other changes to the API of the kernel_launcher module.
    - The numba_dpex.experimental module adds a new dispatcher
      class for numba_dpex kernels. The new dispatcher is a
      numba.dispatcher.Dispathcer subclass.
    - Introduce a new compiler class that is used to compile a
      numba_dpex.kernel decorated function to spirv and then
      store the spirv module as the compiled "overload".
    - Adds an experimental `call_kernel` dpjit function that
      will be used to submit or launch kernels. The `call_kernel`
      function generates LLVM IR code for all the functionality
      currenty done in pure Python in JitKernel.__call__.
@diptorupd
Copy link
Collaborator Author

Merging the PR as it is an experimental preview feature. I have opened issues (refer: #1120) to track all pending work and further polishing that needs to be done to make it ready for moving to core.

@diptorupd diptorupd merged commit ad2cde1 into main Nov 3, 2023
35 of 36 checks passed
@diptorupd diptorupd deleted the feature/KernelDispatcher branch November 3, 2023 15:50
github-actions bot added a commit that referenced this pull request Nov 3, 2023
An experimental kernel dispatcher for numba_dpex.kernel decorator ad2cde1
@ZzEeKkAa ZzEeKkAa mentioned this pull request Dec 1, 2023
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants