nvmath-python 1.0.0__cp314-cp314t-win_amd64.whl

nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/grid_indexer.h

@@ -0,0 +1,69 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_GRID_INDEXER_H_
+#define NVMATH_COPY_KERNEL_IMPL_GRID_INDEXER_H_
+
+#include "copy_kernel_impl/utils.h"
+#include "copy_kernel_impl/vec.h"
+
+namespace nvmath {
+
+template <typename stride_t, bool _needs_grid_stride_loop> struct element_indexer {
+  // stride_t can be 32-bit integer for tensor_volume and gridDim * blockDim up to INT_MAX,
+  // this way unsigned x < INT_MAX; x += INT_MAX cannot overflow
+  using ustride_t = typename type_traits::unsign<stride_t>::type;
+  static constexpr bool needs_grid_stride_loop = _needs_grid_stride_loop;
+
+  constexpr HOST_DEV __forceinline__ element_indexer(const stride_t tensor_volume) : tensor_volume(tensor_volume) {}
+
+  template <typename Cb> __device__ __forceinline__ void with_grid_stride_loop(Cb &&cb) const {
+    // early cast the special indexing variables to the desired integer width type
+    // to avoid arithmetic on 32-bit integers when 64-bit stride_t is used
+    const ustride_t thread_idx = threadIdx.x;
+    const ustride_t block_idx = blockIdx.x;
+    const ustride_t block_dim = blockDim.x;
+    if constexpr (!needs_grid_stride_loop) {
+      const ustride_t x = block_idx * block_dim + thread_idx;
+      if (x < tensor_volume) {
+        cb(x);
+      }
+    } else if constexpr (needs_grid_stride_loop) {
+      const ustride_t grid_dim = gridDim.x;
+      const ustride_t grid_size = grid_dim * block_dim;
+      for (ustride_t x = block_idx * block_dim + thread_idx; x < tensor_volume; x += grid_size) {
+        cb(x);
+      }
+    }
+  }
+
+  ustride_t tensor_volume;
+};
+
+template <typename stride_t, bool _needs_grid_stride_loop> struct block_indexer {
+  using ustride_t = typename type_traits::unsign<stride_t>::type;
+  static constexpr bool needs_grid_stride_loop = _needs_grid_stride_loop;
+
+  constexpr HOST_DEV __forceinline__ block_indexer(const stride_t n_blocks) : n_blocks(n_blocks) {}
+
+  template <typename Cb> __device__ __forceinline__ void with_grid_stride_loop(Cb &&cb) const {
+    // early cast the special indexing variables to the desired integer width type
+    // to avoid arithmetic on 32-bit integers when 64-bit stride_t is used
+    const ustride_t thread_idx = threadIdx.x;
+    const ustride_t block_idx = blockIdx.x;
+    if constexpr (!needs_grid_stride_loop) {
+      cb(block_idx, thread_idx);
+    } else if constexpr (needs_grid_stride_loop) {
+      const ustride_t grid_dim = gridDim.x;
+      for (ustride_t x = block_idx; x < n_blocks; x += grid_dim) {
+        cb(x, thread_idx);
+      }
+    }
+  }
+
+  ustride_t n_blocks;
+};
+
+} // namespace nvmath
+#endif // NVMATH_COPY_KERNEL_IMPL_GRID_INDEXER_H_

nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/transposed.h

@@ -0,0 +1,242 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_TRANSPOSED_H
+#define NVMATH_COPY_KERNEL_IMPL_TRANSPOSED_H
+
+#include "copy_kernel_impl/array_view.h"
+#include "copy_kernel_impl/grid_indexer.h"
+#include "copy_kernel_impl/type_utils.h"
+#include "copy_kernel_impl/utils.h"
+#include "copy_kernel_impl/vec.h"
+
+#if defined(_MSC_VER)
+    // For Visual Studio, use __restrict
+    #define RESTRICT __restrict
+#elif defined(__GNUC__) || defined(__clang__)
+    // For GCC and Clang, use __restrict__
+    #define RESTRICT __restrict__
+#else
+    // Fallback for other compilers, or if restrict is not supported
+    #define RESTRICT
+#endif
+
+namespace nvmath {
+
+namespace detail {
+
+template <typename _T, typename _stride_t, int _transposed_dim, int _block_height, int _block_width,
+          char _reading_order>
+struct transpose2d {
+  using T = _T;
+  using stride_t = _stride_t;
+
+  // the thread position in the tile is represented as a 2d index
+  using thread_coords_t = vec<2, stride_t>;
+
+  // Consecutive flat thread index can be unraveled into 2d coordinates
+  // so that either the left or right coordinate changes faster.
+  // Reading order tells which of the two ways to pick for global memory reads.
+  static constexpr char reading_order = _reading_order;
+  static_assert(reading_order == 'C' || reading_order == 'F', "reading_order must be 'C' or 'F'");
+
+  // Describes splitting of the src/dst shape into two parts
+  // [0, transposed_dim] and [transposed_dim + 1, ndim - 1] to be traversed
+  // with the thread_coords.
+  static constexpr int transposed_dim = _transposed_dim;
+
+  // to unravel flat cuda thread index to 2d index in the tile fast,
+  // we require the tile extents (_block_height, _block_height)
+  // to be powers of 2, and use compile-time helpers that turn
+  // the division/modulo into a bit shift/mask.
+  using warp_size_t = mod_div<32>;
+  using block_height_t = mod_div<_block_height>;
+  using block_width_t = mod_div<_block_width>;
+  using swizzle_stride_t = mod_div<(_block_height > 32 ? 1 : 32 / _block_height)>;
+  using warps_in_width_t = mod_div<(_block_width < 32 ? 1 : _block_width / 32)>;
+  static constexpr block_height_t block_height = block_height_t{};
+  static constexpr block_width_t block_width = block_width_t{};
+  static constexpr swizzle_stride_t swizzle_stride = swizzle_stride_t{};
+  static constexpr warps_in_width_t warps_in_width = warps_in_width_t{};
+  static constexpr warp_size_t warp_size = warp_size_t{};
+  static constexpr int tile_num_elements = _block_height * _block_width;
+
+  HOST_DEV __forceinline__ transpose2d(T *RESTRICT data, stride_t *RESTRICT dst_offset)
+      : data(data), dst_offset(dst_offset) {}
+
+  // 2d index of the thread in the tile in C-like order (i.e. second index changing faster)
+  HOST_DEV __forceinline__ constexpr thread_coords_t unravel_thread_idx_c(const stride_t thread_idx) const {
+    return {thread_idx / block_width, thread_idx % block_width};
+  }
+
+  // 2d index of the thread in the tile in Fortran-like order (i.e. first index changing faster)
+  HOST_DEV __forceinline__ constexpr thread_coords_t unravel_thread_idx_f(const stride_t thread_idx) const {
+    return {thread_idx % block_height, thread_idx / block_height};
+  }
+
+  HOST_DEV __forceinline__ constexpr thread_coords_t unravel_thread_idx_reading(const stride_t thread_idx) const {
+    if constexpr (reading_order == 'C') {
+      return unravel_thread_idx_c(thread_idx);
+    } else {
+      return unravel_thread_idx_f(thread_idx);
+    }
+  }
+
+  HOST_DEV __forceinline__ constexpr thread_coords_t unravel_thread_idx_writing(const stride_t thread_idx) const {
+    if constexpr (reading_order == 'C') {
+      return unravel_thread_idx_f(thread_idx);
+    } else {
+      return unravel_thread_idx_c(thread_idx);
+    }
+  }
+
+  HOST_DEV __forceinline__ constexpr stride_t shm_offset(const thread_coords_t idx) const {
+    const stride_t y = idx[0];
+    const stride_t x = idx[1];
+    // Note, offset(unravel_thread_idx_c(thread_idx)) = idx
+    const stride_t offset = y * block_width + x;
+    const stride_t offset_y = offset / warp_size;
+    const stride_t offset_x = offset % warp_size;
+    // In the simplest case of 32x32 tile, we need to rotate elements
+    // by one every 32 elements to make sure that accessing with
+    // unravel_thread_idx_f(thread_idx)=((0, x), ..., (31, x)) does not
+    // introduce bank conflicts. If block_height is smaller than 32,
+    // there will be 32/block_height different xs in the
+    // unravel_thread_idx_f(thread_idx) warp, so we rotate by
+    // swizzle_stride to make sure that different xs do not land in the same bank.
+    // If block_width is bigger than 32, we want to make sure that
+    // we make one rotation per one y, hence the division by warps_in_width.
+    const stride_t swizzle = (offset_y / warps_in_width) * swizzle_stride;
+    return offset_y * warp_size + ((offset_x + swizzle) % warp_size);
+  }
+
+  // shared memory array to store the elements read from the src tensor
+  T *RESTRICT data;
+  // for data[i], the dst_offset[i] is the offset of the element data[i] in the dst tensor.
+  stride_t *RESTRICT dst_offset;
+};
+} // namespace detail
+
+template <typename dst_array_view_t, typename src_array_view_t, typename grid_indexer_t, typename copy_helper_t>
+struct transpose_copy_impl {
+  // 32 or 64 bit signed integer
+  using stride_t = typename dst_array_view_t::stride_t;
+  // ndim vector of stride_t integers
+  using coords_t = typename dst_array_view_t::coords_t;
+  static constexpr int ndim = dst_array_view_t::ndim;
+  static_assert(ndim == src_array_view_t::ndim, "src and dst must have the same number of dimensions");
+  static constexpr int transposed_dim = copy_helper_t::transposed_dim;
+  static_assert(0 <= transposed_dim && transposed_dim < ndim - 1, "transposed_dim must be between 0 and ndim - 2");
+
+  // use min possible offset to indicate that the element index is out of tensor bounds
+  static constexpr stride_t out_of_bounds_sentinel = type_traits::min_val<stride_t>::value;
+
+  /**
+   * @brief Unravel flat block index and 2d thread index to ndim index of the element in the `shape`.
+   *
+   * E. g. Given tensor of shape 2x3x4 and transposed_dim = 1
+   * [[[0, 1, 2, 3,],
+   *   [4, 5, 6, 7,],
+   *   [8, 9, 10, 11,]],
+
+   *  [[12, 13, 14, 15,],
+   *   [16, 17, 18, 19,],
+   *   [20, 21, 22, 23,]]]
+   *
+   * and a 2x3 tile:
+   * [[(0,0), (0, 1), (0, 2)],
+   *  [(1,0), (1, 1), (1, 2)]]
+   *
+   * maps to the 3D indices of the following elements:
+   * block_idx = 0        block_idx = 1        block_idx = 2        block_idx = 3
+   * [[0, 1, 2],          [[3, 8, 9],          [[10, 11, 16],       [[17, 18, 19],
+   *  [4, 5, 6]]           [7, 12, 13]]         [14, 15, 20]]        [21, 22, 23]]
+   *
+   *
+   * Note, if the tiled extents are not divisible by the tile dimensions,
+   * the parts of the tile are carried over to the next position. Potential
+   * uncoalesced memory accesses are preferred over threads with no work to do
+   * (when mapped to invalid positions in the tensors).
+   *
+   * @param block_idx The threadblock index in the grid
+   * @param shape The shape of the tensor
+   * @param thread_idx The 2d index of the thread in the tile
+   * @return The ndim index of the element in the tensor
+   */
+  __device__ __forceinline__ coords_t unravel_tiled_idx(const stride_t block_idx, const coords_t shape,
+                                                        const vec<2, stride_t> thread_idx) {
+
+    static_assert(ndim >= 2, "ndim must be at least 2");
+    // the extents cannot be negative and the arithmetic on unsigned integer
+    // is noticeably faster
+    using u_stride_t = typename type_traits::unsign<stride_t>::type;
+    u_stride_t flat_idx;
+    coords_t unraveled_idx;
+    auto unravel_extent = [&flat_idx, &unraveled_idx, &shape](int i) {
+      u_stride_t extent = shape[i];
+      if (extent & (extent - 1)) {
+        u_stride_t next_flat_idx = flat_idx / extent;
+        unraveled_idx[i] = flat_idx - next_flat_idx * extent;
+        flat_idx = next_flat_idx;
+      } else {
+        unraveled_idx[i] = flat_idx & (extent - 1);
+        flat_idx >>= ffs(extent) - 1;
+      }
+    };
+    flat_idx = block_idx * copy_helper_t::block_width + thread_idx[1];
+#pragma unroll
+    for (int i = ndim - 1; i > transposed_dim; i--) {
+      unravel_extent(i);
+    }
+    flat_idx = flat_idx * copy_helper_t::block_height + thread_idx[0];
+#pragma unroll
+    for (int i = transposed_dim; i > 0; i--) {
+      unravel_extent(i);
+    }
+    unraveled_idx[0] = flat_idx;
+    return unraveled_idx;
+  }
+
+  void __forceinline__ __device__ operator()(const dst_array_view_t dst_view, const src_array_view_t src_view,
+                                             const grid_indexer_t grid_helper, const copy_helper_t transpose_helper) {
+    grid_helper.with_grid_stride_loop([=](const stride_t flat_block_idx, const stride_t flat_thread_idx) {
+      {
+        // 2d index of the thread
+        const auto thread_coords = transpose_helper.unravel_thread_idx_reading(flat_thread_idx);
+        // ndim index of the element in the source tensor
+        const auto src_coords = unravel_tiled_idx(flat_block_idx, src_view.shape(), thread_coords);
+        // note, all other idxs are 0 < ... < shape[i] by the unravel_tiled_idx,
+        // there's no need to check all of them.
+        const bool is_in_bounds = src_coords[0] < src_view.shape()[0];
+        if constexpr (grid_indexer_t::needs_grid_stride_loop) {
+          __syncthreads();
+        }
+        // flat index in the shared memory arrays where thread_coords should store its data
+        const auto shm_offset = transpose_helper.shm_offset(thread_coords);
+        if (is_in_bounds) {
+          // The data is copied from src to dst through shared memory in two steps:
+          // 1. src[unravel_tiled_idx(.., thread_coords)] -> shm[shm_offset(thread_coords)]
+          // 2. shm[shm_offset(thread_coords_transposed)] -> dst[unravel_tiled_idx(.., thread_coords_transposed)]
+          // To avoid computing the unravel_tiled_idx twice, we compute the dst stride here
+          // and store it together with the data.
+          transpose_helper.data[shm_offset] = src_view[src_coords];
+          transpose_helper.dst_offset[shm_offset] = dst_view.offset(src_coords);
+        } else {
+          transpose_helper.dst_offset[shm_offset] = out_of_bounds_sentinel;
+        }
+        __syncthreads();
+      }
+      // 2d index of the thread with flipped order (i.e. the other index is changing faster)
+      const auto thread_coords = transpose_helper.unravel_thread_idx_writing(flat_thread_idx);
+      const auto shm_offset = transpose_helper.shm_offset(thread_coords);
+      const auto dst_offset = transpose_helper.dst_offset[shm_offset];
+      if (dst_offset != out_of_bounds_sentinel) {
+        dst_view[dst_offset] = transpose_helper.data[shm_offset];
+      }
+    });
+  }
+};
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_TRANSPOSED_H

nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/type_utils.h

@@ -0,0 +1,39 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_TYPE_UTILS_H_
+#define NVMATH_COPY_KERNEL_IMPL_TYPE_UTILS_H_
+
+namespace nvmath {
+using int32_t = int;
+using uint32_t = unsigned int;
+using int64_t = long long int;
+using uint64_t = unsigned long long int;
+static_assert(sizeof(int32_t) == 4, "int32_t must be 4 bytes");
+static_assert(sizeof(uint32_t) == 4, "uint32_t must be 4 bytes");
+static_assert(sizeof(int64_t) == 8, "int64_t must be 8 bytes");
+static_assert(sizeof(uint64_t) == 8, "uint64_t must be 8 bytes");
+
+// Use a struct to represent type of element so that we don't rely
+// on actual representation of the type, available arithmetic etc.
+template <int n_bytes> struct alignas(n_bytes) opaque_t {
+  char data[n_bytes];
+};
+
+static_assert(sizeof(opaque_t<1>) == 1, "opaque_t<1> must be 1 byte");
+static_assert(sizeof(opaque_t<2>) == 2, "opaque_t<2> must be 2 bytes");
+static_assert(sizeof(opaque_t<4>) == 4, "opaque_t<4> must be 4 bytes");
+static_assert(sizeof(opaque_t<8>) == 8, "opaque_t<8> must be 8 bytes");
+static_assert(sizeof(opaque_t<16>) == 16, "opaque_t<16> must be 16 bytes");
+
+static_assert(alignof(opaque_t<1>) == alignof(unsigned char), "opaque_t<1> must be 1 byte");
+static_assert(alignof(opaque_t<2>) == alignof(unsigned short), "opaque_t<2> must be 2 bytes");
+static_assert(alignof(opaque_t<4>) == alignof(unsigned int), "opaque_t<4> must be 4 bytes");
+static_assert(alignof(opaque_t<8>) == alignof(unsigned long long int), "opaque_t<8> must be 8 bytes");
+#ifdef __CUDA_ARCH__
+static_assert(alignof(opaque_t<16>) == 16, "opaque_t<16> must be 16 bytes");
+#endif
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_TYPE_UTILS_H_

nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/utils.h

@@ -0,0 +1,132 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_UTILS_H_
+#define NVMATH_COPY_KERNEL_IMPL_UTILS_H_
+
+#include "copy_kernel_impl/type_utils.h"
+
+#if defined(__CUDACC__)
+#define HOST_DEV __host__ __device__
+#else
+#define HOST_DEV
+#endif
+
+namespace nvmath {
+
+// Some of stl type traits are not available with nvrtc
+namespace type_traits {
+template <bool B, class T, class F> struct conditional {};
+
+template <class T, class F> struct conditional<true, T, F> {
+  using type = T;
+};
+
+template <class T, class F> struct conditional<false, T, F> {
+  using type = F;
+};
+
+template <bool B, class T = void> struct enable_if {};
+template <class T> struct enable_if<true, T> {
+  typedef T type;
+};
+
+template <typename T> struct unsign {};
+
+template <> struct unsign<int64_t> {
+  using type = uint64_t;
+};
+
+template <> struct unsign<int32_t> {
+  using type = uint32_t;
+};
+
+template <typename T> struct is_32_or_64_int {
+  static constexpr bool value = false;
+};
+
+template <> struct is_32_or_64_int<int32_t> {
+  static constexpr bool value = true;
+};
+
+template <> struct is_32_or_64_int<int64_t> {
+  static constexpr bool value = true;
+};
+
+template <typename T> struct min_val {};
+
+template <> struct min_val<int32_t> {
+  static constexpr int32_t value = -2147483648;
+};
+
+template <> struct min_val<int64_t> {
+  static constexpr int64_t value = -9223372036854775808LL;
+};
+} // namespace type_traits
+
+template <typename T, T val> struct const_val {
+  using type = T;
+  static constexpr T value = val;
+};
+
+template <int N> using iconst = const_val<int, N>;
+
+template <bool B> using bconst = const_val<bool, B>;
+
+template <typename true_val_t, typename false_val_t>
+auto HOST_DEV __forceinline__ constexpr cond_val(bconst<false>, true_val_t &&true_val, false_val_t &&false_val) {
+  return false_val;
+}
+
+template <typename true_val_t, typename false_val_t>
+auto HOST_DEV __forceinline__ constexpr cond_val(bconst<true>, true_val_t &&true_val, false_val_t &&false_val) {
+  return true_val;
+}
+
+#if defined(__CUDACC__)
+
+__device__ __forceinline__ int ffs(uint32_t x) { return __ffs(x); }
+
+__device__ __forceinline__ int ffs(int32_t x) { return __ffs(x); }
+
+__device__ __forceinline__ int ffs(uint64_t x) { return __ffsll(x); }
+
+__device__ __forceinline__ int ffs(int64_t x) { return __ffsll(x); }
+
+#endif
+
+constexpr int log2_floor(const int k) {
+  return k == 1 ? 0 : 1 + log2_floor(k >> 1);
+}
+
+template <int k>
+struct mod_div {
+  static_assert(k > 0, "k must be positive");
+  static_assert((k & (k - 1)) == 0, "k must be a power of 2");
+  static constexpr int value = k;
+  static constexpr int log2 = log2_floor(k);
+  static constexpr int mask = k - 1;
+  constexpr int operator()() const {
+    return k;
+  }
+};
+
+template <typename T, int k>
+HOST_DEV __forceinline__ constexpr T operator/(const T a, const mod_div<k>) {
+  return a >> mod_div<k>::log2;
+}
+
+template <typename T, int k>
+HOST_DEV __forceinline__ constexpr T operator%(const T a, const mod_div<k>) {
+  return a & mod_div<k>::mask;
+}
+
+template <typename T, int k>
+HOST_DEV __forceinline__ constexpr T operator*(const T a, const mod_div<k>) {
+  return a << mod_div<k>::log2;
+}
+
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_UTILS_H_

nvmath/internal/ndbuffer/copy_kernel/copy_kernel_impl/vec.h

@@ -0,0 +1,159 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_VEC_H_
+#define NVMATH_COPY_KERNEL_IMPL_VEC_H_
+
+#include "copy_kernel_impl/utils.h"
+
+namespace nvmath {
+
+template <typename T, int N> struct vec_base {
+  T v[N];
+
+  template <typename... Components,
+            typename = typename type_traits::enable_if<sizeof...(Components) <= N &&
+                                                       (type_traits::is_32_or_64_int<Components>::value && ...)>::type>
+  HOST_DEV constexpr vec_base(Components... components) : v{T(components)...} {}
+
+  template <typename U, typename = typename type_traits::enable_if<type_traits::is_32_or_64_int<U>::value>::type>
+  HOST_DEV constexpr vec_base(const U *ptr) {
+    for (int i = 0; i < N; i++) {
+      v[i] = ptr[i];
+    }
+  }
+
+  HOST_DEV __forceinline__ constexpr T &operator[](int i) { return v[i]; }
+  HOST_DEV __forceinline__ constexpr const T &operator[](int i) const { return v[i]; }
+};
+
+template <typename T> struct vec_base<T, 0> {};
+
+template <int N, typename T> struct vec : vec_base<T, N> {
+  using base_t = vec_base<T, N>;
+  using type = T;
+  constexpr static int ndim = N;
+
+  constexpr vec() = default;
+
+  template <typename... Components> HOST_DEV constexpr vec(Components... components) : base_t{components...} {}
+
+  template <typename U> HOST_DEV constexpr vec(const U *ptr, int ndim) : base_t(ptr, ndim) {}
+
+  HOST_DEV __forceinline__ constexpr int size() const { return ndim; }
+
+  template <int K> HOST_DEV __forceinline__ constexpr auto last(const iconst<K>) const {
+    static_assert(K <= ndim);
+    return slice(iconst<ndim - K>(), iconst<ndim>());
+  }
+
+  template <int K> HOST_DEV __forceinline__ constexpr auto first(const iconst<K>) const {
+    static_assert(K <= ndim);
+    return slice(iconst<0>(), iconst<K>());
+  }
+
+  template <int start, int end>
+  HOST_DEV __forceinline__ constexpr vec<end - start, T> slice(const iconst<start>, const iconst<end>) const {
+    static_assert(start >= 0 && end <= ndim);
+    constexpr int slice_ndim = end - start;
+    static_assert(slice_ndim >= 0);
+    if constexpr (slice_ndim != 0) {
+      vec<slice_ndim, T> result;
+#pragma unroll
+      for (int i = 0; i < slice_ndim; i++) {
+        result[i] = this->operator[](start + i);
+      }
+      return result;
+    }
+    return {};
+  }
+};
+
+template <int N, int M, typename T>
+HOST_DEV __forceinline__ constexpr vec<N + M, T> cat(const vec<N, T> a, const vec<M, T> b) {
+  constexpr int ndim = N + M;
+  if constexpr (ndim != 0) {
+    vec<ndim, T> result;
+    if constexpr (N > 0) {
+#pragma unroll
+      for (int i = 0; i < N; i++) {
+        result[i] = a[i];
+      }
+    }
+    if constexpr (M > 0) {
+#pragma unroll
+      for (int i = 0; i < M; i++) {
+        result[N + i] = b[i];
+      }
+    }
+    return result;
+  }
+  return {};
+}
+
+template <int N, typename T, typename Op>
+HOST_DEV __forceinline__ constexpr auto vector_bin_op(const vec<N, T> a, const vec<N, T> b, Op &&op) {
+  if constexpr (N == 0) {
+    return {};
+  } else if constexpr (N != 0) {
+    using result_t = decltype(op(a[0], b[0]));
+    vec<N, result_t> result;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      result[i] = op(a[i], b[i]);
+    }
+    return result;
+  }
+}
+
+template <int N, typename T> HOST_DEV __forceinline__ constexpr auto operator+(const vec<N, T> a, const vec<N, T> b) {
+  return vector_bin_op(a, b, [](T a, T b) { return a + b; });
+}
+
+template <int N, typename T> HOST_DEV __forceinline__ constexpr auto operator*(const vec<N, T> a, const vec<N, T> b) {
+  return vector_bin_op(a, b, [](T a, T b) { return a * b; });
+}
+
+template <int N, typename T> HOST_DEV __forceinline__ constexpr auto operator-(const vec<N, T> a, const vec<N, T> b) {
+  return vector_bin_op(a, b, [](T a, T b) { return a - b; });
+}
+
+template <int N, typename T> HOST_DEV __forceinline__ constexpr auto operator/(const vec<N, T> a, const vec<N, T> b) {
+  return vector_bin_op(a, b, [](T a, T b) { return a / b; });
+}
+
+template <int N, typename Pred, typename T, typename... Ts>
+HOST_DEV __forceinline__ constexpr bool any(Pred &&pred, const vec<N, T> a, const vec<N, Ts>... vs) {
+  for (int i = 0; i < N; i++) {
+    if (pred(a[i], vs[i]...))
+      return true;
+  }
+  return false;
+}
+
+template <int N, typename Pred, typename T, typename... Ts>
+HOST_DEV __forceinline__ constexpr bool all(Pred &&pred, const vec<N, T> a, const vec<N, Ts>... vs) {
+  for (int i = 0; i < N; i++) {
+    if (!pred(a[i], vs[i]...))
+      return false;
+  }
+  return true;
+}
+
+template <int N, typename T> HOST_DEV __forceinline__ constexpr T dot(const vec<N, T> a, const vec<N, T> b) {
+  if constexpr (N == 0) {
+    return 0;
+  } else if constexpr (N != 0) {
+    T sum = a[0] * b[0];
+#pragma unroll
+    for (int i = 1; i < N; i++) {
+      sum += a[i] * b[i];
+    }
+    return sum;
+  }
+}
+
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_VEC_H_