Check pointer alignment in as_c_contiguous_array_generic_impl

Author: ndgrigorian

dpctl/tensor/libtensor/include/kernels/copy_as_contiguous.hpp

@@ -148,41 +148,92 @@ as_c_contiguous_array_generic_impl(sycl::queue &exec_q,
  constexpr std::size_t preferred_lws = 256;
  constexpr std::uint32_t n_vecs = 2;
  constexpr std::uint32_t vec_sz = 4;
- constexpr bool enable_sg_load = true;
- using KernelName =
- as_contig_krn<T, IndexerT, vec_sz, n_vecs, enable_sg_load>;
 
- const auto &kernel_id = sycl::get_kernel_id<KernelName>();
+ using dpctl::tensor::kernels::alignment_utils::
+ disabled_sg_loadstore_wrapper_krn;
+ using dpctl::tensor::kernels::alignment_utils::is_aligned;
+ using dpctl::tensor::kernels::alignment_utils::required_alignment;
 
- auto const &ctx = exec_q.get_context();
- auto const &dev = exec_q.get_device();
- auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(
- ctx, {dev}, {kernel_id});
+ sycl::event copy_ev;
+ if (is_aligned<required_alignment>(src_p) &&
+ is_aligned<required_alignment>(dst_p))
+ {
+ constexpr bool enable_sg_load = true;
+ using KernelName =
+ as_contig_krn<T, IndexerT, vec_sz, n_vecs, enable_sg_load>;
 
- auto krn = kb.get_kernel(kernel_id);
+  const auto &kernel_id = sycl::get_kernel_id<KernelName>();
 
- const std::uint32_t max_sg_size = krn.template get_info<
- sycl::info::kernel_device_specific::max_sub_group_size>(dev);
+ auto const &ctx = exec_q.get_context();
+ auto const &dev = exec_q.get_device();
+ auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(
+ ctx, {dev}, {kernel_id});
 
- const std::size_t lws =
- ((preferred_lws + max_sg_size - 1) / max_sg_size) * max_sg_size;
+ auto krn = kb.get_kernel(kernel_id);
 
- constexpr std::uint32_t nelems_per_wi = n_vecs * vec_sz;
- size_t n_groups =
- (nelems + nelems_per_wi * lws - 1) / (nelems_per_wi * lws);
+ const std::uint32_t max_sg_size = krn.template get_info<
+ sycl::info::kernel_device_specific::max_sub_group_size>(dev);
 
- sycl::event copy_ev = exec_q.submit([&](sycl::handler &cgh) {
- cgh.depends_on(depends);
- cgh.use_kernel_bundle(kb);
+ const std::size_t lws =
+ ((preferred_lws + max_sg_size - 1) / max_sg_size) * max_sg_size;
 
- const sycl::range<1> gRange{n_groups * lws};
- const sycl::range<1> lRange{lws};
+ constexpr std::uint32_t nelems_per_wi = n_vecs * vec_sz;
+ size_t n_groups =
+ (nelems + nelems_per_wi * lws - 1) / (nelems_per_wi * lws);
 
- cgh.parallel_for<KernelName>(
- sycl::nd_range<1>(gRange, lRange),
- CopyAsCContigFunctor<T, IndexerT, vec_sz, n_vecs, enable_sg_load>(
- nelems, src_tp, dst_tp, src_indexer));
- });
+ copy_ev = exec_q.submit([&](sycl::handler &cgh) {
+ cgh.depends_on(depends);
+ cgh.use_kernel_bundle(kb);
+
+ const sycl::range<1> gRange{n_groups * lws};
+ const sycl::range<1> lRange{lws};
+
+ cgh.parallel_for<KernelName>(
+ sycl::nd_range<1>(gRange, lRange),
+ CopyAsCContigFunctor<T, IndexerT, vec_sz, n_vecs,
+ enable_sg_load>(nelems, src_tp, dst_tp,
+ src_indexer));
+ });
+ }
+ else {
+ constexpr bool disable_sg_load = false;
+ using InnerKernelName =
+ as_contig_krn<T, IndexerT, vec_sz, n_vecs, disable_sg_load>;
+ using KernelName = disabled_sg_loadstore_wrapper_krn<InnerKernelName>;
+
+ const auto &kernel_id = sycl::get_kernel_id<KernelName>();
+
+ auto const &ctx = exec_q.get_context();
+ auto const &dev = exec_q.get_device();
+ auto kb = sycl::get_kernel_bundle<sycl::bundle_state::executable>(
+ ctx, {dev}, {kernel_id});
+
+ auto krn = kb.get_kernel(kernel_id);
+
+ const std::uint32_t max_sg_size = krn.template get_info<
+ sycl::info::kernel_device_specific::max_sub_group_size>(dev);
+
+ const std::size_t lws =
+ ((preferred_lws + max_sg_size - 1) / max_sg_size) * max_sg_size;
+
+ constexpr std::uint32_t nelems_per_wi = n_vecs * vec_sz;
+ size_t n_groups =
+ (nelems + nelems_per_wi * lws - 1) / (nelems_per_wi * lws);
+
+ copy_ev = exec_q.submit([&](sycl::handler &cgh) {
+ cgh.depends_on(depends);
+ cgh.use_kernel_bundle(kb);
+
+ const sycl::range<1> gRange{n_groups * lws};
+ const sycl::range<1> lRange{lws};
+
+ cgh.parallel_for<KernelName>(
+ sycl::nd_range<1>(gRange, lRange),
+ CopyAsCContigFunctor<T, IndexerT, vec_sz, n_vecs,
+ disable_sg_load>(nelems, src_tp, dst_tp,
+ src_indexer));
+ });
+ }
 
  return copy_ev;
 }