nvmath-python 0.9.0__cp314-cp314-win_amd64.whl

nvmath/internal/__init__.pxd

@@ -0,0 +1,3 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/internal/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+"""APIs in this module are for NVIDIA-use only.
+
+APIs in internal are not intended for third-party developers, but are shared between NVIDIA
+packages.
+"""

nvmath/internal/_bindings.pxd

@@ -0,0 +1,12 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from libc.stdint cimport intptr_t, int64_t, uint64_t
+
+cpdef memcpy_async(intptr_t dst_ptr, intptr_t src_ptr, int64_t size, intptr_t stream)
+cpdef stream_sync(intptr_t stream)
+cpdef uint64_t get_memory_pool_reserved_memory_size(pool) except? -1
+cpdef uint64_t get_memory_pool_used_memory_size(pool) except? -1
+cpdef free_memory_pool_reserved_memory(pool)
+cpdef launch_kernel(intptr_t f, intptr_t kernel_params, int gx, int gy, int gz, int bx, int by, int bz, unsigned int shared_mem_bytes, intptr_t stream_handle)

nvmath/internal/_device_utils.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+try:
+    from cuda.core import Device, system
+except ImportError:
+    from cuda.core.experimental import Device, system
+
+
+import threading
+
+thread_local = threading.local()
+
+
+def get_device(device_id: int) -> Device:
+    """
+    NOTE: the function DOES NOT change the current cuda device. If it needs
+    to call .set_current(), it immediately restores the previous device
+    before exiting. To temporarily change the current device,
+    please use utils.device_ctx().
+
+    A counterpart to cuda.core.Device(int), with extra initialization.
+    Returns initialized cuda.core.Device(device_id) object, i.e.
+    makes sure to call .set_current() on the instance once per thread.
+    Calling .set_current() is not just about a valid CUDA context in a thread,
+    but specifically cuda.core requirement to initialize the device instance.
+
+    The function keeps a thread-local cache of "initialized"
+    device instances. It mirrors cuda.core.Device(int) caching,
+    adding the set_current() call on the first access.
+    """
+    try:
+        devices = thread_local.devices
+    except AttributeError:
+        try:
+            num_devices = system.get_num_devices()
+        except AttributeError:
+            # cuda.core < 0.5.0
+            num_devices = system.num_devices
+        thread_local.devices = devices = [None] * num_devices
+
+    device = devices[device_id]
+    if device is None:
+        current_device = Device()
+        current_device_id = current_device.device_id
+        try:
+            device = Device(device_id)
+            device.set_current()
+            devices[device_id] = device
+        finally:
+            if current_device_id != device_id:
+                current_device.set_current()
+    return device

nvmath/internal/enum_utils.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Factories for create options dataclasses, as well as utilities to add docstring to enum
+classes.
+"""
+
+import dataclasses
+import re
+from collections.abc import Callable
+from enum import IntEnum
+from typing import Any, ClassVar
+
+import numpy
+
+
+def create_options_class_from_enum(
+    options_class_name: str,
+    enum_class: IntEnum,
+    get_attr_dtype: Callable,
+    description: str,
+    filter_re: str = r"(?P<option_name>.*)",
+):
+    """
+    Create an options dataclass from a Python enum class. Names can be filtered if desired.
+
+    Args:
+       options_class_name: Name of the dataclass that will be created.
+
+       enum_class: The IntEnum class that contains the options for the dataclass.
+
+       get_attr_dtype: A callable that takes in an enum value as the argument and returns
+            the size in bytes of the cuTensorNet.
+
+       filter_re: A re definition that defines the match named 'option_name'.
+    """
+    if r"(?P<option_name>" not in filter_re:
+        message = """Incorrect re.
+The re for the filter must contain the named group 'option_name'."""
+        raise ValueError(message)
+
+    # Helper vars for creating attribute docstring.
+    doc = f"""A data class for capturing the {description} options.
+
+    Attributes:
+"""
+    indent = " " * 8
+    prefix = determine_enum_prefix(enum_class, "_ATTRIBUTE")
+
+    filter_pattern = re.compile(filter_re)
+    option_to_enum: dict[str, IntEnum] = {}
+    option_to_dtype: dict[str, Any] = {}
+    # mypy claims that IntEnum is not iterable
+    for e in enum_class:  # type: ignore
+        m = filter_pattern.match(e.name)
+        if not m:
+            continue
+        option_name = m.group("option_name").lower()
+        option_to_enum[option_name] = e
+        option_to_dtype[option_name] = get_attr_dtype(e)
+
+        # Add docstring for this attribute.
+        doc += indent + option_name + ":" + f" See `{prefix + '_' + m.group(0)}`.\n"
+
+    field: tuple[str, Any, None | dict]
+    fields: list[tuple[str, Any, None | dict]] = []
+    for option_name, dtype in option_to_dtype.items():
+        if numpy.issubdtype(dtype, numpy.integer):
+            field = option_name, int | None, dataclasses.field(default=None)
+        else:
+            field = option_name, Any | None, dataclasses.field(default=None)
+        fields.append(field)
+
+    # Add class attributes.
+
+    field = "option_to_enum", ClassVar[dict], dataclasses.field(default=option_to_enum)
+    fields.append(field)
+
+    field = "option_to_dtype", ClassVar[dict], dataclasses.field(default=option_to_dtype)
+    fields.append(field)
+
+    # Create the options class.
+    options_class = dataclasses.make_dataclass(options_class_name, fields, slots=True)
+    options_class.__doc__ = doc
+
+    return options_class
+
+
+def snake_to_camel(names):
+    name = ""
+    for i, sub_name in enumerate(names):
+        if i == 0:
+            name += sub_name.lower()
+        else:
+            name += sub_name[0].upper() + sub_name[1:]
+    name += "_t"
+    return name
+
+
+def camel_to_snake(name, upper=True):
+    """
+    Convert string from camel case to snake style.
+    """
+
+    def transformer(m):
+        prefix = iter(("", "_", ""))
+        for i in 1, 3, 5:
+            first, second = i, i + 1
+            s = next(prefix)
+            if m.group(second):
+                if m.group(first):  # If the second group doesn't exist, the first won't either by the design of the RE.
+                    s += m.group(first).lower() + "_"
+                s += m.group(second).lower()
+                break
+        return s
+
+    name = re.sub(r"^([A-Z]*)([A-Z])|(?<!_)([A-Z]*)([A-Z])|([A-Z]*)([A-Z])", transformer, name)
+    if upper:
+        name = name.upper()
+    return name
+
+
+def determine_enum_prefix(enum_class, chomp):
+    """
+    This function assumes that the convention used to translate C enumerators to Python enum
+    names holds.
+    """
+
+    prefix = enum_class.__module__.split(".")[-1].upper()
+    prefix += "_" + camel_to_snake(enum_class.__name__)
+    prefix = re.sub(chomp, "", prefix)
+    return prefix
+
+
+def add_enum_class_doc(enum_class, chomp):
+    """
+    Add docstring to enum classes.
+    """
+    for e in enum_class:
+        e.__doc__ = f"See `{determine_enum_prefix(enum_class, chomp) + '_' + e.name.upper()}`."

nvmath/internal/formatters.py

@@ -0,0 +1,87 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Formatters for printing data.
+"""
+
+import numpy as np
+
+
+class MemoryStr:
+    """
+    A simple type to pretty-print memory-like values.
+    """
+
+    def __init__(self, memory, base_unit="B"):
+        self.memory = memory
+        self.base_unit = base_unit
+        self.base = 1024
+
+    def __str__(self):
+        """
+        Convert large values to powers of 1024 for readability.
+        """
+
+        base, base_unit, memory = self.base, self.base_unit, self.memory
+
+        if memory < base:
+            value, unit = memory, base_unit
+        elif memory < base**2:
+            value, unit = memory / base, f"Ki{base_unit}"
+        elif memory < base**3:
+            value, unit = memory / base**2, f"Mi{base_unit}"
+        else:
+            value, unit = memory / base**3, f"Gi{base_unit}"
+
+        return f"{value:0.2f} {unit}"
+
+
+class FLOPSStr:
+    """
+    A simple type to pretty-print FLOP count and FLOP/s-like values.
+    """
+
+    def __init__(self, flops, base_unit):
+        self.flops = flops
+        self.base_unit = base_unit
+        self.base = 1000
+
+    def __str__(self):
+        """
+        Convert large values to powers of 1000 for readability.
+        """
+
+        base, base_unit, flops = self.base, self.base_unit, self.flops
+
+        if flops < base:
+            value, unit = flops, base_unit
+        elif flops < base**2:
+            value, unit = flops / base, f"K{base_unit}"
+        elif flops < base**3:
+            value, unit = flops / base**2, f"M{base_unit}"
+        elif flops < base**4:
+            value, unit = flops / base**3, f"G{base_unit}"
+        elif flops < base**5:
+            value, unit = flops / base**4, f"T{base_unit}"
+        else:
+            value, unit = flops / base**5, f"P{base_unit}"
+
+        return f"{value:0.3f} {unit}"
+
+
+def array2string(array_like):
+    """
+    String representation of an array-like object with possible truncation of "interior"
+    values to limit string size.
+
+    The NumPy function "set_printoptions" can be used to control the display of the array.
+    """
+
+    return np.array2string(
+        np.asanyarray(array_like, dtype=object),
+        separator=", ",
+        # NumPy hates empty strings so we print 'None' instead.
+        formatter={"object": lambda s: s if s != "" else "None"},
+    )

nvmath/internal/mem_limit.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Memory specification regular expression.
+"""
+
+__all__ = ["MEM_LIMIT_RE_PCT", "MEM_LIMIT_RE_VAL", "MEM_LIMIT_DOC", "check_memory_str"]
+
+import re
+
+MEM_LIMIT_RE_PCT = re.compile(r"(?P<value>[+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)\s*%\s*$")
+MEM_LIMIT_RE_VAL = re.compile(
+    r"(?P<value>[+]?[0-9]*\.?[0-9]+(?:[eE][-+]?[0-9]+)?)\s*(?P<units>[kmg])?(?P<binary>(?<=[kmg])i)?b\s*$",
+    re.IGNORECASE,
+)
+MEM_LIMIT_DOC = """The {kind} must be specified in one of the following forms:
+  (1) A number (int or float). If the number is a float between 0 and 1 inclusive, the
+      {kind} is interpreted as a fraction of the total device memory; otherwise it is
+      interpreted as the number of bytes of memory, with float value being cast to int.
+      Examples: 0.75, 50E6, 50000000, ...
+  (2) A string containing a positive value followed by B, kB, MB, or GB for powers of 1000.
+      Examples: "0.05 GB", "50 MB", "50000000 B" ...
+  (3) A string containing a positive value followed by kiB, MiB, or GiB for powers of 1024.
+      Examples:  "0.05 GiB", "51.2 MiB", "53687091 B" ...
+  (4) A string with value in the range [0, 100] followed by a % symbol.
+      Examples: "26%","82%", ...
+
+  Whitespace between values and units is optional.
+
+The provided {kind} is "{value}".
+"""
+
+
+def check_memory_str(value, kind):
+    """
+    Check if the memory specification string is valid.
+
+    value = the memory specification string.
+    kind  = a string denoting the type of memory being checked, used in error messages.
+    """
+    if not isinstance(value, int | float):
+        m1 = MEM_LIMIT_RE_PCT.match(value)
+        if m1:
+            factor = float(m1.group("value"))
+            if factor < 0 or factor > 100:
+                raise ValueError(f"The {kind} percentage must be in the range [0, 100].")
+        m2 = MEM_LIMIT_RE_VAL.match(value)
+        if not (m1 or m2):
+            raise ValueError(MEM_LIMIT_DOC.format(kind=kind, value=value))

nvmath/internal/memory.pxd

@@ -0,0 +1,10 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+cimport cython
+from libc.stdint cimport int64_t, intptr_t, uint64_t
+
+cpdef get_device_memory_resource(int device_id)
+cpdef allocate_from_mr(mr, int64_t size, stream, int device_id, logger=*)
+cpdef free_reserved_memory(bint sync=*)

nvmath/internal/memory.pyi

@@ -0,0 +1,46 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import _cython_3_2_4
+from typing import Any, ClassVar
+
+__pyx_capi__: dict
+__test__: dict
+allocate_from_mr: _cython_3_2_4.cython_function_or_method
+free_reserved_memory: _cython_3_2_4.cython_function_or_method
+get_device_memory_resource: _cython_3_2_4.cython_function_or_method
+
+class _MemoryPointer:
+    """
+    Temporary internal NDBuffer allocation adapter class. NDBuffer expects
+    custom allocator to return a cuda.core.Buffer instance. Until all supported
+    cuda.core versions have unified support for wrapping external allocations
+    with Buffer.from_handle, this class servers as an adapter/workaround.
+
+    WARNING: This is internal tool subject to change/removal without notice.
+
+    Internally, it is used conditionally if any of the following is needed:
+        * wrap external allocations that don't come as cuda.core.Buffer (e.g. from cupy).
+        This is needed because prior to cuda.core 0.5.0, it's not possible to
+        pass reference to externall RAII object as ``owner`` parameter to Buffer.from_handle.
+        * provide debug logging on deallocation. In the future, we can use
+        Buffer.from_handle(owner=...) to inject deallocation callback (cuda.core >= 0.5.0)
+        or weakref.finalize for Buffer (cuda.core >= 0.6.0).
+
+    The only publicly exposed field is the handle - a base pointer to the allocated memory.
+    """
+    from_handle: ClassVar[method] = ...
+    __pyx_vtable__: ClassVar[PyCapsule] = ...
+    handle: handle
+    owner: owner
+    @classmethod
+    def __init__(cls, *args, **kwargs) -> None:
+        """Create and return a new object.  See help(type) for accurate signature."""
+    def __del__(self, *args, **kwargs) -> None: ...
+    def __reduce__(self):
+        """_MemoryPointer.__reduce_cython__(self)"""
+    def __reduce_cython__(self) -> Any:
+        """_MemoryPointer.__reduce_cython__(self)"""
+    def __setstate_cython__(self, __pyx_state) -> Any:
+        """_MemoryPointer.__setstate_cython__(self, __pyx_state)"""

nvmath/internal/ndbuffer/__init__.pxd

@@ -0,0 +1,3 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0

nvmath/internal/ndbuffer/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from . import ndbuffer, package_utils
+
+__all__ = ["ndbuffer", "package_utils"]

nvmath/internal/ndbuffer/copy_kernel/args.h

@@ -0,0 +1,34 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_ARGS_H
+#define NVMATH_COPY_KERNEL_IMPL_ARGS_H
+
+#include "copy_kernel_impl/type_utils.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 {
+template <int N>
+struct KernelArgs {
+    void * RESTRICT dst_ptr;
+    const void * RESTRICT src_ptr;
+    int64_t dst_shape[N];
+    int64_t src_shape[N];
+    int64_t dst_strides[N];
+    int64_t src_strides[N];
+    int64_t grid_arg;
+};
+}
+
+#endif // NVMATH_COPY_KERNEL_IMPL_ARGS_H

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

@@ -0,0 +1,52 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_ARRAY_VIEW_H_
+#define NVMATH_COPY_KERNEL_IMPL_ARRAY_VIEW_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 {
+
+template <typename T, typename _coords_t> struct array_view {
+  // While indices cannot be negative (only strides can),
+  // we're using the same 32- or 64-bit signed type to represent both
+  // indices and strides for simplicity. In the end we need to convert
+  // both to the same signed type when computing the offset.
+  using coords_t = _coords_t;
+  using stride_t = typename coords_t::type;
+  using dtype_t = T;
+  static constexpr int ndim = coords_t::ndim;
+
+  HOST_DEV constexpr array_view(T *__restrict__ data, coords_t shape, coords_t strides)
+      : shape_(shape), strides_(strides), data_(data) {}
+
+  HOST_DEV __forceinline__ T &operator[](const coords_t idx) const { return data_[offset(idx)]; }
+  HOST_DEV __forceinline__ T &operator[](const stride_t offset) const { return data_[offset]; }
+  HOST_DEV __forceinline__ stride_t offset(const coords_t idx) const { return dot(idx, strides()); }
+  HOST_DEV __forceinline__ coords_t shape() const { return shape_; }
+  HOST_DEV __forceinline__ coords_t strides() const { return strides_; }
+  HOST_DEV __forceinline__ T *data() const { return data_; }
+
+protected:
+  coords_t shape_;
+  coords_t strides_;
+  T *RESTRICT data_;
+};
+
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_ARRAY_VIEW_H_

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

@@ -0,0 +1,68 @@
+// Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef NVMATH_COPY_KERNEL_IMPL_ELEMENTWISE_H
+#define NVMATH_COPY_KERNEL_IMPL_ELEMENTWISE_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"
+
+namespace nvmath {
+
+namespace detail {
+
+template <typename coords_t, typename stride_t = typename coords_t::type>
+__device__ __forceinline__ coords_t unravel_idx(const stride_t flat_idx, const coords_t shape) {
+  constexpr int ndim = coords_t::ndim;
+  if constexpr (ndim <= 0) {
+    return {};
+  } else if constexpr (ndim == 1) {
+    return {flat_idx};
+  } else if constexpr (ndim > 1) {
+
+    // 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 u_flat_idx = flat_idx;
+    coords_t unraveled_coords;
+#pragma unroll
+    for (int i = ndim - 1; i >= 1; i--) {
+      u_stride_t extent = shape[i];
+      if (extent & (extent - 1)) {
+        u_stride_t next_flat_idx = u_flat_idx / extent;
+        unraveled_coords[i] = u_flat_idx - next_flat_idx * extent;
+        u_flat_idx = next_flat_idx;
+      } else {
+        unraveled_coords[i] = u_flat_idx & (extent - 1);
+        u_flat_idx >>= ffs(extent) - 1;
+      }
+    }
+    unraveled_coords[0] = u_flat_idx;
+    return unraveled_coords;
+  }
+}
+
+} // namespace detail
+
+template <bool has_equal_shapes, typename dst_array_view_t, typename src_array_view_t, typename grid_indexer_t>
+struct elementwise_copy_impl {
+  using stride_t = typename dst_array_view_t::stride_t;
+
+  void __forceinline__ __device__ operator()(const dst_array_view_t &&dst_view, const src_array_view_t &&src_view,
+                                             const grid_indexer_t &&grid_helper) {
+    grid_helper.with_grid_stride_loop([=](const stride_t flat_element_idx) {
+      const auto dst_coords = detail::unravel_idx(flat_element_idx, dst_view.shape());
+      const auto src_coords =
+          cond_val(bconst<has_equal_shapes>(), dst_coords, detail::unravel_idx(flat_element_idx, src_view.shape()));
+      dst_view[dst_coords] = src_view[src_coords];
+    });
+  }
+};
+
+} // namespace nvmath
+
+#endif // NVMATH_COPY_KERNEL_IMPL_ELEMENTWISE_H

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_