text: # Copyright 2021 Intel Corporation
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  from numba.core import cgutils, types
  from numba.core.ir_utils import legalize_names
  
  from numba_dppy import numpy_usm_shared as nus
  from numba_dppy import utils
  from numba_dppy.driver import DpctlCAPIFnBuilder
  from numba_dppy.driver._helpers import numba_type_to_dpctl_typenum
  
  
  class KernelLaunchOps:
      """
      KernelLaunchOps(lowerer, cres, num_inputs)
      Defines a set of functions to launch a SYCL kernel on the "current queue"
      as defined in the dpctl queue manager.
      """
  
      def _form_kernel_arg_and_arg_ty(self, val, ty):
          kernel_arg_dst = self.builder.gep(
              self.kernel_arg_array,
              [self.context.get_constant(types.int32, self.cur_arg)],
          )
          kernel_arg_ty_dst = self.builder.gep(
              self.kernel_arg_ty_array,
              [self.context.get_constant(types.int32, self.cur_arg)],
          )
          self.cur_arg += 1
          self.builder.store(val, kernel_arg_dst)
          self.builder.store(ty, kernel_arg_ty_dst)
  
      def __init__(self, lowerer, kernel, num_inputs):
          """Create an instance of KernelLaunchOps.
  
          Args:
              lowerer: The Numba Lowerer that will be used to generate the code.
              kernel: The SYCL kernel for which we are generating the code.
              num_inputs: The number of arguments to the kernels.
          """
          self.lowerer = lowerer
          self.context = self.lowerer.context
          self.builder = self.lowerer.builder
          self.kernel = kernel
          self.kernel_addr = self.kernel.addressof_ref()
          self.total_kernel_args = 0
          self.cur_arg = 0
          self.num_inputs = num_inputs
          # list of buffer that needs to comeback to host
          self.write_buffs = []
          # list of buffer that does not need to comeback to host
          self.read_only_buffs = []
  
      def get_current_queue(self):
          """Allocates memory on the stack to store the current queue from dpctl.
  
          A SYCL queue is needed to allocate USM memory and submit a kernel. This
          function gets the queue returned by ``DPCTLQueueMgr_GetCurrentQueue``
          function and stores it on the stack. The queue should be freed properly
          after returning from the kernel.
  
          Return: A LLVM Value storing the pointer to the SYCL queue returned
                  by ``DPCTLQueueMgr_GetCurrentQueue``.
  
          """
          sycl_queue_val = cgutils.alloca_once(
              self.builder, utils.get_llvm_type(context=self.context, type=types.voidptr)
          )
          fn = DpctlCAPIFnBuilder.get_dpctl_queuemgr_get_current_queue(
              builder=self.builder, context=self.context
          )
          self.builder.store(self.builder.call(fn, []), sycl_queue_val)
          return sycl_queue_val
  
      def free_queue(self, sycl_queue_val):
          """
          Frees the ``DPCTLSyclQueueRef`` pointer that was used to launch the
          kernels.
  
          Args:
              sycl_queue_val: The SYCL queue pointer to be freed.
          """
          fn = DpctlCAPIFnBuilder.get_dpctl_queue_delete(
              builder=self.builder, context=self.context
          )
          self.builder.call(fn, [self.builder.load(sycl_queue_val)])
  
      def allocate_kernel_arg_array(self, num_kernel_args):
          """Allocates an array to store the LLVM Value for every kernel argument.
  
          Args:
              num_kernel_args : The number of kernel arguments for the kernel.
  
          """
  
          self.total_kernel_args = num_kernel_args
  
          # we need a kernel arg array to enqueue
          self.kernel_arg_array = cgutils.alloca_once(
              self.builder,
              utils.get_llvm_type(context=self.context, type=types.voidptr),
              size=self.context.get_constant(types.uintp, num_kernel_args),
              name="kernel_arg_array",
          )
  
          self.kernel_arg_ty_array = cgutils.alloca_once(
              self.builder,
              utils.LLVMTypes.int32_t,
              size=self.context.get_constant(types.uintp, num_kernel_args),
              name="kernel_arg_ty_array",
          )
  
      def process_kernel_arg(
          self, var, llvm_arg, arg_type, index, modified_arrays, sycl_queue_val
      ):
          """
          process_kernel_arg(var, llvm_arg, arg_type, index, modified_arrays,
          sycl_queue_val)
          Creates an LLVM Value for each kernel argument.
  
          Args:
              var : A kernel argument represented as a Numba type.
              llvm_arg : Only used for array arguments and points to the LLVM
                         value previously allocated to store the array arg.
              arg_type : The Numba type for the argument.
              index : The poisition of the argument in the list of arguments.
              modified_arrays : The list of array arguments that are written to
                                inside the kernel. The list is used to check if
                                the argument is read-only or not.
  
          Raises:
              NotImplementedError: If an unsupported type of kernel argument is
                                   encountered.
          """
          if isinstance(arg_type, types.npytypes.Array):
              if llvm_arg is None:
                  raise NotImplementedError(arg_type, var)
  
              storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)
              self.builder.store(self.context.get_constant(types.int64, 0), storage)
              ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
              self._form_kernel_arg_and_arg_ty(
                  self.builder.bitcast(
                      storage,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  ty,
              )
  
              storage = cgutils.alloca_once(self.builder, utils.LLVMTypes.int64_t)
              self.builder.store(self.context.get_constant(types.int64, 0), storage)
              ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
              self._form_kernel_arg_and_arg_ty(
                  self.builder.bitcast(
                      storage,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  ty,
              )
  
              # Handle array size
              array_size_member = self.builder.gep(
                  llvm_arg,
                  [
                      self.context.get_constant(types.int32, 0),
                      self.context.get_constant(types.int32, 2),
                  ],
              )
  
              ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
              self._form_kernel_arg_and_arg_ty(
                  self.builder.bitcast(
                      array_size_member,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  ty,
              )
  
              # Handle itemsize
              item_size_member = self.builder.gep(
                  llvm_arg,
                  [
                      self.context.get_constant(types.int32, 0),
                      self.context.get_constant(types.int32, 3),
                  ],
              )
  
              ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
              self._form_kernel_arg_and_arg_ty(
                  self.builder.bitcast(
                      item_size_member,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  ty,
              )
  
              # Calculate total buffer size
              total_size = cgutils.alloca_once(
                  self.builder,
                  utils.get_llvm_type(context=self.context, type=types.intp),
                  size=utils.get_one(context=self.context),
                  name="total_size" + str(self.cur_arg),
              )
              self.builder.store(
                  self.builder.sext(
                      self.builder.mul(
                          self.builder.load(array_size_member),
                          self.builder.load(item_size_member),
                      ),
                      utils.get_llvm_type(context=self.context, type=types.intp),
                  ),
                  total_size,
              )
  
              # Handle data
              data_member = self.builder.gep(
                  llvm_arg,
                  [
                      self.context.get_constant(types.int32, 0),
                      self.context.get_constant(types.int32, 4),
                  ],
              )
  
              # names are replaced using legalize names, we have to do the same
              # here for them to match.
              legal_names = legalize_names([var])
              ty = numba_type_to_dpctl_typenum(context=self.context, type=types.voidptr)
  
              if isinstance(arg_type, nus.UsmSharedArrayType):
                  self._form_kernel_arg_and_arg_ty(
                      self.builder.bitcast(
                          self.builder.load(data_member),
                          utils.get_llvm_type(context=self.context, type=types.voidptr),
                      ),
                      ty,
                  )
              else:
                  malloc_fn = DpctlCAPIFnBuilder.get_dpctl_malloc_shared(
                      builder=self.builder, context=self.context
                  )
                  memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(
                      builder=self.builder, context=self.context
                  )
                  event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(
                      builder=self.builder, context=self.context
                  )
                  event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(
                      builder=self.builder, context=self.context
                  )
  
                  # Not known to be USM so we need to copy to USM.
                  buffer_name = "buffer_ptr" + str(self.cur_arg)
                  # Create void * to hold new USM buffer.
                  buffer_ptr = cgutils.alloca_once(
                      self.builder,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                      name=buffer_name,
                  )
                  # Setup the args to the USM allocator, size and SYCL queue.
                  args = [
                      self.builder.load(total_size),
                      self.builder.load(sycl_queue_val),
                  ]
                  # Call USM shared allocator and store in buffer_ptr.
                  self.builder.store(self.builder.call(malloc_fn, args), buffer_ptr)
  
                  if legal_names[var] in modified_arrays:
                      self.write_buffs.append((buffer_ptr, total_size, data_member))
                  else:
                      self.read_only_buffs.append((buffer_ptr, total_size, data_member))
  
                  # We really need to detect when an array needs to be copied over
                  if index < self.num_inputs:
                      args = [
                          self.builder.load(sycl_queue_val),
                          self.builder.load(buffer_ptr),
                          self.builder.bitcast(
                              self.builder.load(data_member),
                              utils.get_llvm_type(
                                  context=self.context, type=types.voidptr
                              ),
                          ),
                          self.builder.load(total_size),
                      ]
                      event_ref = self.builder.call(memcpy_fn, args)
                      self.builder.call(event_wait_fn, [event_ref])
                      self.builder.call(event_del_fn, [event_ref])
  
                  self._form_kernel_arg_and_arg_ty(self.builder.load(buffer_ptr), ty)
  
              # Handle shape
              shape_member = self.builder.gep(
                  llvm_arg,
                  [
                      self.context.get_constant(types.int32, 0),
                      self.context.get_constant(types.int32, 5),
                  ],
              )
  
              for this_dim in range(arg_type.ndim):
                  shape_entry = self.builder.gep(
                      shape_member,
                      [
                          self.context.get_constant(types.int32, 0),
                          self.context.get_constant(types.int32, this_dim),
                      ],
                  )
                  ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
                  self._form_kernel_arg_and_arg_ty(
                      self.builder.bitcast(
                          shape_entry,
                          utils.get_llvm_type(context=self.context, type=types.voidptr),
                      ),
                      ty,
                  )
  
              # Handle strides
              stride_member = self.builder.gep(
                  llvm_arg,
                  [
                      self.context.get_constant(types.int32, 0),
                      self.context.get_constant(types.int32, 6),
                  ],
              )
  
              for this_stride in range(arg_type.ndim):
                  stride_entry = self.builder.gep(
                      stride_member,
                      [
                          self.context.get_constant(types.int32, 0),
                          self.context.get_constant(types.int32, this_stride),
                      ],
                  )
  
                  ty = numba_type_to_dpctl_typenum(context=self.context, type=types.int64)
                  self._form_kernel_arg_and_arg_ty(
                      self.builder.bitcast(
                          stride_entry,
                          utils.get_llvm_type(context=self.context, type=types.voidptr),
                      ),
                      ty,
                  )
  
          else:
              ty = numba_type_to_dpctl_typenum(context=self.context, type=arg_type)
              self._form_kernel_arg_and_arg_ty(
                  self.builder.bitcast(
                      llvm_arg,
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  ty,
              )
  
      def enqueue_kernel_and_copy_back(self, dim_bounds, sycl_queue_val):
          """
          enqueue_kernel_and_copy_back(dim_bounds, sycl_queue_val)
          Submits the kernel to the specified queue, waits and then copies
          back any results to the host.
  
          Args:
              dim_bounds : An array of three tuple representing the starting
                           offset, end offset and the stride (step) for each
                           dimension of the input arrays. Every array in a parfor
                           is of the same dimensionality and shape, thus ensuring
                           the bounds are the same.
              sycl_queue_val : The SYCL queue on which the kernel is
                               submitted.
          """
          submit_fn = DpctlCAPIFnBuilder.get_dpctl_queue_submit_range(
              builder=self.builder, context=self.context
          )
          queue_wait_fn = DpctlCAPIFnBuilder.get_dpctl_queue_wait(
              builder=self.builder, context=self.context
          )
          event_del_fn = DpctlCAPIFnBuilder.get_dpctl_event_delete(
              builder=self.builder, context=self.context
          )
          memcpy_fn = DpctlCAPIFnBuilder.get_dpctl_queue_memcpy(
              builder=self.builder, context=self.context
          )
          free_fn = DpctlCAPIFnBuilder.get_dpctl_free_with_queue(
              builder=self.builder, context=self.context
          )
          event_wait_fn = DpctlCAPIFnBuilder.get_dpctl_event_wait(
              builder=self.builder, context=self.context
          )
  
          # the assumption is loop_ranges will always be less than or equal to 3
          # dimensions
          num_dim = len(dim_bounds) if len(dim_bounds) < 4 else 3
  
          # form the global range
          global_range = cgutils.alloca_once(
              self.builder,
              utils.get_llvm_type(context=self.context, type=types.uintp),
              size=self.context.get_constant(types.uintp, num_dim),
              name="global_range",
          )
  
          intp_t = utils.get_llvm_type(context=self.context, type=types.intp)
          intp_ptr_t = utils.get_llvm_ptr_type(intp_t)
  
          for i in range(num_dim):
              start, stop, step = dim_bounds[i]
              if stop.type != utils.LLVMTypes.int64_t:
                  stop = self.builder.sext(stop, utils.LLVMTypes.int64_t)
  
              # we reverse the global range to account for how sycl and opencl
              # range differs
              self.builder.store(
                  stop,
                  self.builder.gep(
                      global_range,
                      [self.context.get_constant(types.uintp, (num_dim - 1) - i)],
                  ),
              )
  
          args = [
              self.builder.inttoptr(
                  self.context.get_constant(types.uintp, self.kernel_addr),
                  utils.get_llvm_type(context=self.context, type=types.voidptr),
              ),
              self.builder.load(sycl_queue_val),
              self.kernel_arg_array,
              self.kernel_arg_ty_array,
              self.context.get_constant(types.uintp, self.total_kernel_args),
              self.builder.bitcast(global_range, intp_ptr_t),
              self.context.get_constant(types.uintp, num_dim),
              self.builder.bitcast(
                  utils.create_null_ptr(builder=self.builder, context=self.context),
                  utils.get_llvm_type(context=self.context, type=types.voidptr),
              ),
              self.context.get_constant(types.uintp, 0),
          ]
  
          # Submit the kernel
          event_ref = self.builder.call(submit_fn, args)
  
          # Add a wait on the queue
          self.builder.call(queue_wait_fn, [self.builder.load(sycl_queue_val)])
  
          # Note that the dpctl_queue_wait call waits on the event and then
          # decrements the ref count of the sycl::event C++ object. However, the
          # event object returned by the get_dpctl_queue_submit_range call still
          # needs to be explicitly deleted to free up the event object properly.
          self.builder.call(event_del_fn, [event_ref])
  
          # read buffers back to host
          for write_buff in self.write_buffs:
              buffer_ptr, total_size, data_member = write_buff
              args = [
                  self.builder.load(sycl_queue_val),
                  self.builder.bitcast(
                      self.builder.load(data_member),
                      utils.get_llvm_type(context=self.context, type=types.voidptr),
                  ),
                  self.builder.load(buffer_ptr),
                  self.builder.load(total_size),
              ]
              # FIXME: In future, when the DctlQueue_Memcpy is made non-blocking
              # the returned event should be explicitly freed by calling
              # get_dpctl_event_delete.
              event_ref = self.builder.call(memcpy_fn, args)
              self.builder.call(event_wait_fn, [event_ref])
              self.builder.call(event_del_fn, [event_ref])
  
              self.builder.call(
                  free_fn,
                  [
                      self.builder.load(buffer_ptr),
                      self.builder.load(sycl_queue_val),
                  ],
              )
  
          for read_buff in self.read_only_buffs:
              buffer_ptr, total_size, data_member = read_buff
              self.builder.call(
                  free_fn,
                  [
                      self.builder.load(buffer_ptr),
                      self.builder.load(sycl_queue_val),
                  ],
              )
  