Tweaks to copy_as_contiguous CContig functor

1. Save common subexpressions to variables
2. Sub-group size type changed to uint16 (from uint32)
3. sg.get_local_range() replaced with sg.get_max_local_range()

   This is safe to do since work-group size is chosen to be
   a multiple of sub-group size for all possile choices
   of sub-group size (1, 8, 16, 32, 64)

4. Simplified computation of base value in generic branch for
   complex types, or when sg_load is disabled, to avoid a
   division (and left a comment)

Author: oleksandr-pavlyk

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

@@ -44,8 +44,8 @@ namespace copy_as_contig
 
 template <typename T,
           typename IndexerT,
-          int vec_sz = 4,
-          int n_vecs = 2,
+          unsigned int vec_sz = 4,
+          unsigned int n_vecs = 2,
           bool enable_sg_loadstore = true>
 class CopyAsCContigFunctor
 {
@@ -66,52 +66,59 @@ class CopyAsCContigFunctor
 
     void operator()(sycl::nd_item<1> ndit) const
     {
+        constexpr std::uint8_t elems_per_wi =
+            static_cast<std::uint8_t>(vec_sz * n_vecs);
         using dpctl::tensor::type_utils::is_complex;
+
         if constexpr (!enable_sg_loadstore || is_complex<T>::value) {
-            const std::uint32_t sgSize =
-                ndit.get_sub_group().get_local_range()[0];
+            const std::uint16_t sgSize =
+                ndit.get_sub_group().get_max_local_range()[0];
             const std::size_t gid = ndit.get_global_linear_id();
 
-            const std::size_t base =
-                (gid / sgSize) * sgSize * n_vecs * vec_sz + (gid % sgSize);
-            for (size_t offset = base;
-                 offset < std::min(nelems, base + sgSize * (n_vecs * vec_sz));
-                 offset += sgSize)
-            {
+            // base = (gid / sgSize) * sgSize * elems_per_wi + (gid % sgSize)
+            // gid % sgSize == gid - (gid / sgSize) * sgSize
+            const std::size_t elems_per_sg = sgSize * (elems_per_wi - 1);
+            const std::size_t base = (gid / sgSize) * elems_per_sg + gid;
+            const std::size_t offset_max =
+                std::min(nelems, base + sgSize * elems_per_wi);
+
+            for (size_t offset = base; offset < offset_max; offset += sgSize) {
                 auto src_offset = src_indexer(offset);
                 dst_p[offset] = src_p[src_offset];
             }
         }
         else {
             auto sg = ndit.get_sub_group();
-            const std::uint32_t sgSize = sg.get_local_range()[0];
-            const size_t base = n_vecs * vec_sz *
-                                (ndit.get_group(0) * ndit.get_local_range(0) +
-                                 sg.get_group_id()[0] * sgSize);
+            const std::uint16_t sgSize = sg.get_max_local_range()[0];
+            const size_t base =
+                elems_per_wi * (ndit.get_group(0) * ndit.get_local_range(0) +
+                                sg.get_group_id()[0] * sgSize);
 
-            if (base + n_vecs * vec_sz * sgSize < nelems) {
+            if (base + elems_per_wi * sgSize < nelems) {
                 sycl::vec<T, vec_sz> dst_vec;
 
 #pragma unroll
-                for (std::uint32_t it = 0; it < n_vecs * vec_sz; it += vec_sz) {
+                for (std::uint8_t it = 0; it < elems_per_wi; it += vec_sz) {
+                    const size_t block_start_id = base + it * sgSize;
+
                     auto dst_multi_ptr = sycl::address_space_cast<
                         sycl::access::address_space::global_space,
-                        sycl::access::decorated::yes>(
-                        &dst_p[base + it * sgSize]);
+                        sycl::access::decorated::yes>(&dst_p[block_start_id]);
 
+                    const size_t elem_id0 = block_start_id + sg.get_local_id();
 #pragma unroll
-                    for (std::uint32_t k = 0; k < vec_sz; k++) {
-                        ssize_t src_offset = src_indexer(
-                            base + (it + k) * sgSize + sg.get_local_id());
+                    for (std::uint8_t k = 0; k < vec_sz; k++) {
+                        const size_t elem_id = elem_id0 + k * sgSize;
+                        const ssize_t src_offset = src_indexer(elem_id);
                         dst_vec[k] = src_p[src_offset];
                     }
                     sg.store<vec_sz>(dst_multi_ptr, dst_vec);
                 }
             }
             else {
-                for (size_t k = base + sg.get_local_id()[0]; k < nelems;
-                     k += sgSize)
-                {
+                const size_t lane_id = sg.get_local_id()[0];
+                const size_t k0 = base + lane_id;
+                for (size_t k = k0; k < nelems; k += sgSize) {
                     ssize_t src_offset = src_indexer(k);
                     dst_p[k] = src_p[src_offset];
                 }
@@ -122,8 +129,8 @@ class CopyAsCContigFunctor
 
 template <typename T,
           typename IndexT,
-          int vec_sz,
-          int n_vecs,
+          std::uint32_t vec_sz,
+          std::uint32_t n_vecs,
           bool enable_sgload>
 class as_contig_krn;
 
@@ -154,6 +161,9 @@ as_c_contiguous_array_generic_impl(sycl::queue &exec_q,
     using dpctl::tensor::kernels::alignment_utils::is_aligned;
     using dpctl::tensor::kernels::alignment_utils::required_alignment;
 
+    constexpr std::uint8_t nelems_per_wi =
+        static_cast<std::uint8_t>(n_vecs * vec_sz);
+
     sycl::event copy_ev;
     if (is_aligned<required_alignment>(src_p) &&
         is_aligned<required_alignment>(dst_p))
@@ -177,7 +187,6 @@ as_c_contiguous_array_generic_impl(sycl::queue &exec_q,
         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);
 
@@ -216,7 +225,6 @@ as_c_contiguous_array_generic_impl(sycl::queue &exec_q,
         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);