Factors checks for sufficient memory range into output_validation.hpp

ndgrigorian · ndgrigorian · commit 81d0f9e9f6e9 · 2024-02-13T14:28:02.000-08:00
diff --git a/dpctl/tensor/libtensor/include/utils/output_validation.hpp b/dpctl/tensor/libtensor/include/utils/output_validation.hpp
@@ -37,10 +37,9 @@ namespace tensor
 namespace validation
 {
 
-/*! @brief Raises a value error if a function would attempt to write
-    to an array which is read-only.
+/*! @brief Raises a value error if an array is read-only.
 
-    This should always be called on an array before it will be written to.*/
+    This should be called with an array before writing.*/
 struct CheckWritable
 {
     static void throw_if_not_writable(const dpctl::tensor::usm_ndarray &arr)
@@ -52,6 +51,26 @@ struct CheckWritable
     }
 };
 
+/*! @brief Raises a value error if an array's memory is not sufficiently ample
+    to accommodate an input number of elements.
+
+    This should be called with an array before writing.*/
+struct AmpleMemory
+{
+    template <typename T>
+    static void throw_if_not_ample(const dpctl::tensor::usm_ndarray &arr,
+                                   T nelems)
+    {
+        auto arr_offsets = arr.get_minmax_offsets();
+        T range = static_cast<T>(arr_offsets.second - arr_offsets.first);
+        if (range + 1 < nelems) {
+            throw py::value_error("Memory addressed by the output array is not "
+                                  "sufficiently ample.");
+        }
+        return;
+    }
+};
+
 } // namespace validation
 } // namespace tensor
 } // namespace dpctl
diff --git a/dpctl/tensor/libtensor/source/boolean_advanced_indexing.cpp b/dpctl/tensor/libtensor/source/boolean_advanced_indexing.cpp
@@ -174,19 +174,8 @@ py_extract(const dpctl::tensor::usm_ndarray &src,
         throw py::value_error("Inconsistent array dimensions");
     }
 
-    // ensure that dst is sufficiently ample
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < static_cast<size_t>(ortho_nelems * masked_dst_nelems)) {
-            throw py::value_error(
-                "Memory addressed by the destination array can not "
-                "accommodate all the "
-                "array elements.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(
+        dst, ortho_nelems * masked_dst_nelems);
 
     auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
     // check that dst does not intersect with src, not with cumsum.
@@ -507,19 +496,8 @@ py_place(const dpctl::tensor::usm_ndarray &dst,
         throw py::value_error("Inconsistent array dimensions");
     }
 
-    // ensure that dst is sufficiently ample
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < static_cast<size_t>(ortho_nelems * masked_dst_nelems)) {
-            throw py::value_error(
-                "Memory addressed by the destination array can not "
-                "accommodate all the "
-                "array elements.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(
+        dst, ortho_nelems * masked_dst_nelems);
 
     auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
     // check that dst does not intersect with src, not with cumsum.
@@ -794,18 +772,8 @@ py_nonzero(const dpctl::tensor::usm_ndarray
         throw py::value_error("Arrays are expected to ave no memory overlap");
     }
 
-    // ensure that dst is sufficiently ample
-    auto indexes_offsets = indexes.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(indexes_offsets.second - indexes_offsets.first);
-        if (range + 1 < static_cast<size_t>(nz_elems * _ndim)) {
-            throw py::value_error(
-                "Memory addressed by the destination array can not "
-                "accommodate all the array elements.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(
+        indexes, nz_elems * _ndim);
 
     std::vector<sycl::event> host_task_events;
     host_task_events.reserve(2);
diff --git a/dpctl/tensor/libtensor/source/clip.cpp b/dpctl/tensor/libtensor/source/clip.cpp
@@ -155,19 +155,7 @@ py_clip(const dpctl::tensor::usm_ndarray &src,
                               "have the same data type");
     }
 
-    // ensure that dst is sufficiently ample
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < static_cast<size_t>(nelems)) {
-            throw py::value_error(
-                "Memory addressed by the destination array can not "
-                "accommodate all the "
-                "array elements.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, nelems);
 
     char *src_data = src.get_data();
     char *min_data = min.get_data();
diff --git a/dpctl/tensor/libtensor/source/copy_and_cast_usm_to_usm.cpp b/dpctl/tensor/libtensor/source/copy_and_cast_usm_to_usm.cpp
@@ -101,17 +101,7 @@ copy_usm_ndarray_into_usm_ndarray(const dpctl::tensor::usm_ndarray &src,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // destination must be ample enough to accommodate all elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     // check compatibility of execution queue and allocation queue
     if (!dpctl::utils::queues_are_compatible(exec_q, {src, dst})) {
diff --git a/dpctl/tensor/libtensor/source/copy_for_reshape.cpp b/dpctl/tensor/libtensor/source/copy_for_reshape.cpp
@@ -88,17 +88,7 @@ copy_usm_ndarray_for_reshape(const dpctl::tensor::usm_ndarray &src,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // destination must be ample enough to accommodate all elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        py::ssize_t range =
-            static_cast<py::ssize_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     // check same contexts
     if (!dpctl::utils::queues_are_compatible(exec_q, {src, dst})) {
diff --git a/dpctl/tensor/libtensor/source/copy_for_roll.cpp b/dpctl/tensor/libtensor/source/copy_for_roll.cpp
@@ -111,17 +111,7 @@ copy_usm_ndarray_for_roll_1d(const dpctl::tensor::usm_ndarray &src,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // destination must be ample enough to accommodate all elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        py::ssize_t range =
-            static_cast<py::ssize_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     // check same contexts
     if (!dpctl::utils::queues_are_compatible(exec_q, {src, dst})) {
@@ -301,17 +291,7 @@ copy_usm_ndarray_for_roll_nd(const dpctl::tensor::usm_ndarray &src,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // destination must be ample enough to accommodate all elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        py::ssize_t range =
-            static_cast<py::ssize_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     // check for compatible queues
     if (!dpctl::utils::queues_are_compatible(exec_q, {src, dst})) {
diff --git a/dpctl/tensor/libtensor/source/copy_numpy_ndarray_into_usm_ndarray.cpp b/dpctl/tensor/libtensor/source/copy_numpy_ndarray_into_usm_ndarray.cpp
@@ -89,18 +89,7 @@ void copy_numpy_ndarray_into_usm_ndarray(
         return;
     }
 
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements of source
-    // array
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     if (!dpctl::utils::queues_are_compatible(exec_q, {dst})) {
         throw py::value_error("Execution queue is not compatible with the "
diff --git a/dpctl/tensor/libtensor/source/elementwise_functions/elementwise_functions.hpp b/dpctl/tensor/libtensor/source/elementwise_functions/elementwise_functions.hpp
@@ -118,18 +118,7 @@ py_unary_ufunc(const dpctl::tensor::usm_ndarray &src,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // ensure that output is ample enough to accommodate all elements
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     // check memory overlap
     auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
@@ -376,18 +365,7 @@ std::pair<sycl::event, sycl::event> py_binary_ufunc(
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // ensure that output is ample enough to accommodate all elements
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < src_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, src_nelems);
 
     auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
     auto const &same_logical_tensors =
@@ -702,18 +680,7 @@ py_binary_inplace_ufunc(const dpctl::tensor::usm_ndarray &lhs,
         return std::make_pair(sycl::event(), sycl::event());
     }
 
-    // ensure that output is ample enough to accommodate all elements
-    auto lhs_offsets = lhs.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(lhs_offsets.second - lhs_offsets.first);
-        if (range + 1 < rhs_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(rhs, rhs_nelems);
 
     // check memory overlap
     auto const &same_logical_tensors =
diff --git a/dpctl/tensor/libtensor/source/integer_advanced_indexing.cpp b/dpctl/tensor/libtensor/source/integer_advanced_indexing.cpp
@@ -352,17 +352,8 @@ usm_ndarray_take(const dpctl::tensor::usm_ndarray &src,
         }
     }
 
-    // destination must be ample enough to accommodate all elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if ((range + 1) < (orthog_nelems * ind_nelems)) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(
+        dst, orthog_nelems * ind_nelems);
 
     int ind_sh_elems = std::max<int>(ind_nd, 1);
 
@@ -641,17 +632,7 @@ usm_ndarray_put(const dpctl::tensor::usm_ndarray &dst,
     py::ssize_t dst_offset = py::ssize_t(0);
     py::ssize_t val_offset = py::ssize_t(0);
 
-    // destination must be ample enough to accommodate all possible elements
-    {
-        auto dst_offsets = dst.get_minmax_offsets();
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if ((range + 1) < dst_nelems) {
-            throw py::value_error(
-                "Destination array can not accommodate all the "
-                "elements of source array.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, dst_nelems);
 
     int dst_typenum = dst.get_typenum();
     int val_typenum = val.get_typenum();
diff --git a/dpctl/tensor/libtensor/source/linalg_functions/dot.cpp b/dpctl/tensor/libtensor/source/linalg_functions/dot.cpp
@@ -248,19 +248,7 @@ py_dot(const dpctl::tensor::usm_ndarray &x1,
         throw py::value_error("dst shape and size mismatch");
     }
 
-    // ensure that dst is sufficiently ample
-    auto dst_offsets = dst.get_minmax_offsets();
-    // destination must be ample enough to accommodate all elements
-    {
-        size_t range =
-            static_cast<size_t>(dst_offsets.second - dst_offsets.first);
-        if (range + 1 < dst_nelems) {
-            throw py::value_error(
-                "Memory addressed by the destination array can not "
-                "accommodate all the "
-                "array elements.");
-        }
-    }
+    dpctl::tensor::validation::AmpleMemory::throw_if_not_ample(dst, dst_nelems);
 
     auto const &overlap = dpctl::tensor::overlap::MemoryOverlap();
     // check that dst does not intersect with x1 or x2
diff --git a/dpctl/tensor/libtensor/source/reductions/reduction_over_axis.hpp b/dpctl/tensor/libtensor/source/reductions/reduction_over_axis.hpp
diff --git a/dpctl/tensor/libtensor/source/repeat.cpp b/dpctl/tensor/libtensor/source/repeat.cpp
diff --git a/dpctl/tensor/libtensor/source/sorting/argsort.cpp b/dpctl/tensor/libtensor/source/sorting/argsort.cpp
diff --git a/dpctl/tensor/libtensor/source/sorting/sort.cpp b/dpctl/tensor/libtensor/source/sorting/sort.cpp
diff --git a/dpctl/tensor/libtensor/source/where.cpp b/dpctl/tensor/libtensor/source/where.cpp
