Improve support for operations between clad::array and clad::array_ref

PetroZarytskyi · Nov 14, 2024 · b30a742 · b30a742
1 parent 997c8f9
commit b30a742
Showing 1 changed file with 56 additions and 8 deletions.
diff --git a/include/clad/Differentiator/ArrayRef.h b/include/clad/Differentiator/ArrayRef.h
@@ -52,6 +52,14 @@ template <typename T> class array_ref {
     m_size = a.size();
     return *this;
   }
+  template <typename L, typename BinaryOp, typename R>
+  CUDA_HOST_DEVICE array_ref<T>&
+  operator=(const array_expression<L, BinaryOp, R>& arr_exp) {
+    assert(arr_exp.size() == m_size);
+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] = arr_exp[i];
+    return *this;
+  }
   /// Returns the size of the underlying array
   constexpr CUDA_HOST_DEVICE std::size_t size() const { return m_size; }
   constexpr CUDA_HOST_DEVICE PUREFUNC T* ptr() const { return m_arr; }
@@ -71,7 +79,7 @@ template <typename T> class array_ref {
   // Arithmetic overloads
   /// Divides the arrays element wise
   template <typename U>
-  CUDA_HOST_DEVICE array_ref<T>& operator/=(array_ref<U>& Ar) {
+  CUDA_HOST_DEVICE array_ref<T>& operator/=(const array_ref<U>& Ar) {
     assert(m_size == Ar.size() && "Size of both the array_refs must be equal "
                                   "for carrying out addition assignment");
     for (std::size_t i = 0; i < m_size; i++)
@@ -80,7 +88,7 @@ template <typename T> class array_ref {
   }
   /// Multiplies the arrays element wise
   template <typename U>
-  CUDA_HOST_DEVICE array_ref<T>& operator*=(array_ref<U>& Ar) {
+  CUDA_HOST_DEVICE array_ref<T>& operator*=(const array_ref<U>& Ar) {
     assert(m_size == Ar.size() && "Size of both the array_refs must be equal "
                                   "for carrying out addition assignment");
     for (std::size_t i = 0; i < m_size; i++)
@@ -89,7 +97,7 @@ template <typename T> class array_ref {
   }
   /// Adds the arrays element wise
   template <typename U>
-  CUDA_HOST_DEVICE array_ref<T>& operator+=(array_ref<U>& Ar) {
+  CUDA_HOST_DEVICE array_ref<T>& operator+=(const array_ref<U>& Ar) {
     assert(m_size == Ar.size() && "Size of both the array_refs must be equal "
                                   "for carrying out addition assignment");
     for (std::size_t i = 0; i < m_size; i++)
@@ -98,36 +106,76 @@ template <typename T> class array_ref {
   }
   /// Subtracts the arrays element wise
   template <typename U>
-  CUDA_HOST_DEVICE array_ref<T>& operator-=(array_ref<U>& Ar) {
+  CUDA_HOST_DEVICE array_ref<T>& operator-=(const array_ref<U>& Ar) {
     assert(m_size == Ar.size() && "Size of both the array_refs must be equal "
                                   "for carrying out addition assignment");
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] -= Ar[i];
     return *this;
   }
   /// Divides the elements of the array_ref by elements of the array
-  template <typename U> CUDA_HOST_DEVICE array_ref<T>& operator/=(array<U>& A) {
+  template <typename U>
+  CUDA_HOST_DEVICE array_ref<T>& operator/=(const array<U>& A) {
     assert(m_size == A.size() && "Size of arrays must be equal");
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] /= A[i];
     return *this;
   }
   /// Multiplies the elements of the array_ref by elements of the array
-  template <typename U> CUDA_HOST_DEVICE array_ref<T>& operator*=(array<U>& A) {
+  template <typename L, typename BinaryOp, typename R>
+  CUDA_HOST_DEVICE array_ref<T>&
+  operator*=(const array_expression<L, BinaryOp, R>& arr_exp) {
+    assert(arr_exp.size() == m_size);
+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] *= arr_exp[i];
+    return *this;
+  }
+  /// Adds the elements of the array_ref by elements of the array
+  template <typename L, typename BinaryOp, typename R>
+  CUDA_HOST_DEVICE array_ref<T>&
+  operator+=(const array_expression<L, BinaryOp, R>& arr_exp) {
+    assert(arr_exp.size() == m_size);
+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] += arr_exp[i];
+    return *this;
+  }
+  /// Subtracts the elements of the array_ref by elements of the array
+  template <typename L, typename BinaryOp, typename R>
+  CUDA_HOST_DEVICE array_ref<T>&
+  operator-=(const array_expression<L, BinaryOp, R>& arr_exp) {
+    assert(arr_exp.size() == m_size);
+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] -= arr_exp[i];
+    return *this;
+  }
+  /// Divides the elements of the array_ref by elements of the array
+  template <typename L, typename BinaryOp, typename R>
+  CUDA_HOST_DEVICE array_ref<T>&
+  operator/=(const array_expression<L, BinaryOp, R>& arr_exp) {
+    assert(arr_exp.size() == m_size);
+    for (std::size_t i = 0; i < m_size; i++)
+      m_arr[i] /= arr_exp[i];
+    return *this;
+  }
+  /// Multiplies the elements of the array_ref by elements of the array
+  template <typename U>
+  CUDA_HOST_DEVICE array_ref<T>& operator*=(const array<U>& A) {
     assert(m_size == A.size() && "Size of arrays must be equal");
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] *= A[i];
     return *this;
   }
   /// Adds the elements of the array_ref by elements of the array
-  template <typename U> CUDA_HOST_DEVICE array_ref<T>& operator+=(array<U>& A) {
+  template <typename U>
+  CUDA_HOST_DEVICE array_ref<T>& operator+=(const array<U>& A) {
     assert(m_size == A.size() && "Size of arrays must be equal");
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] += A[i];
     return *this;
   }
   /// Subtracts the elements of the array_ref by elements of the array
-  template <typename U> CUDA_HOST_DEVICE array_ref<T>& operator-=(array<U>& A) {
+  template <typename U>
+  CUDA_HOST_DEVICE array_ref<T>& operator-=(const array<U>& A) {
     assert(m_size == A.size() && "Size of arrays must be equal");
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] -= A[i];