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;
 }