-
Notifications
You must be signed in to change notification settings - Fork 33
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
Conversation
TODOs
|
81f631a
to
52b356a
Compare
52b356a
to
aeb58f2
Compare
There was a problem hiding this 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.
6656dcd
to
b14fdec
Compare
edd7c04
to
aad1184
Compare
- 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__.
- Added checks for compute follows data compliance to kernel compilation. - Removed support for __get_item__ in KernelDispatcher - Address review comments.
2e483ab
to
81ecaef
Compare
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 |
An experimental kernel dispatcher for numba_dpex.kernel decorator ad2cde1
Numba-dpex uses the
numba_dpex.core.kernel_interface.JitKernel
class a custom dispatcher class for running kernel functions. TheJitKernel
dispatcher has a different API thannumba.dispatcher.Dispatcher
making it hard to useJitKernel
as a target dispatcher inside Numba. The limitation means we currently cannot do the following:DpexKernelTarget
numba_dpex.dpjit
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 ofnumba.core.dispatcher.Dispatcher
. TheKernelDispatcher
is for now included as an experimental feature, as to be useful all existinglower
functions innumba_dpex.ocl.oclimpl
would have to be added toDpexKernelTargetContext
. The overloads will be done as a follow up PR.The experimental
KernelDispacther
compiles a@kernel
decorated function to SPIRV and stores it as anoverload
.To call an experimental kernel, a new
call_kernel
function has been added to numba_dpex.experimental. The function is decorated usingdpjit
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 thesycl::kernel_bundle
- Submit the kernel to the
sycl::queue
extracted from a DpnpNdArray kernel argument- Wait on the queue