Add support for array arguments in vector mode

.clang-format

@@ -3,6 +3,7 @@ BasedOnStyle: LLVM
 Language: Cpp
 Standard: Cpp11
 PointerAlignment: Left
+RemoveBracesLLVM: true
 
 IncludeCategories:
   - Regex:           '^"[^/]+\"'

.github/workflows/clang-tidy-review.yml

@@ -27,7 +27,7 @@ jobs:
         id: review
         with:
           build_dir: build
-          exclude: "test/*,benchmark/*"
+          exclude: "test/*,benchmark/*,demos/*"
           split_workflow: true
           cmake_command: >
             pip install cmake lit &&

benchmark/Simple.cpp

@@ -55,5 +55,38 @@ static void BM_ReverseModeSumExecute(benchmark::State &state) {
 }
 BENCHMARK(BM_ReverseModeSumExecute);
 
+// Benchmark computing gradient using vector forward mode via
+// a forward declaration.
+inline void sum_dvec_0(double*, int, clad::array_ref<double>);
+static void BM_VectorForwardModeSumFwdDecl(benchmark::State &state) {
+  auto vm_grad = clad::differentiate<clad::opts::vector_mode>(sum, "p");
+  (void) vm_grad;
+  double inputs[] = {1, 2, 3, 4, 5};
+  double result[3] = {};
+  clad::array_ref<double> result_ref(result, 3);
+  unsigned long long sum = 0;
+  for (auto _ : state) {
+    sum_dvec_0(inputs,/*dim*/ 3, result_ref);
+    benchmark::DoNotOptimize(sum += result[0] + result[1] + result[2]);
+  }
+}
+BENCHMARK(BM_VectorForwardModeSumFwdDecl);
+
+// Benchmark computing gradient using vector forward mode via
+// CladFunction::execute.
+static void BM_VectorForwardModeSumExecute(benchmark::State &state) {
+  auto vm_grad = clad::differentiate<clad::opts::vector_mode>(sum, "p");
+  double inputs[] = {1, 2, 3, 4, 5};
+  double result[3] = {};
+  clad::array_ref<double> result_ref(result, 3);
+  unsigned long long sum = 0;
+  for (auto _ : state) {
+    vm_grad.execute(inputs,/*dim*/ 3, result_ref);
+    benchmark::DoNotOptimize(sum += result[0] + result[1] + result[2]);
+  }
+}
+BENCHMARK(BM_VectorForwardModeSumExecute);
+
+
 // Define our main.
 BENCHMARK_MAIN();

demos/VectorForwardMode.cpp

@@ -0,0 +1,51 @@
+//--------------------------------------------------------------------*- C++ -*-
+// clad - The C++ Clang-based Automatic Differentiator
+//
+// A demo, describing how to use vector forward mode AD to differentiation
+// a function with respect to multiple parameters.
+//----------------------------------------------------------------------------//
+
+// To run the demo please type:
+// path/to/clang++  -Xclang -add-plugin -Xclang clad -Xclang -load -Xclang \
+// path/to/libclad.so  -I../include/ -x c++ -std=c++11 VectorForwardMode.cpp
+//
+// A typical invocation would be:
+// ../../../../obj/Debug+Asserts/bin/clang++  -Xclang -add-plugin -Xclang clad \
+// -Xclang -load -Xclang ../../../../obj/Debug+Asserts/lib/libclad.dylib     \
+// -I../include/ -x c++ -std=c++11 VectorForwardMode.cpp
+
+// Necessary for clad to work include
+#include "clad/Differentiator/Differentiator.h"
+
+// A function for weighted sum of array elements.
+double weighted_sum(double* arr, double* weights, int n) {
+  double res = 0;
+  for (int i = 0; i < n; ++i)
+    res += weights[i] * arr[i];
+  return res;
+}
+
+int main() {
+  auto weighted_sum_grad =
+      clad::differentiate<clad::opts::vector_mode>(weighted_sum, "arr,weights");
+
+  // Initialize array and weights.
+  double arr[3] = {3.0, 4.0, 5.0};
+  double weights[3] = {0.5, 0.7, 0.9};
+
+  // Allocate memory for derivatives.
+  double d_arr[3] = {0.0, 0.0, 0.0};
+  double d_weights[3] = {0.0, 0.0, 0.0};
+  clad::array_ref<double> d_arr_ref(d_arr, 3);
+  clad::array_ref<double> d_weights_ref(d_weights, 3);
+
+  // Calculate gradient.
+  weighted_sum_grad.execute(arr, weights, 3, d_arr_ref, d_weights_ref);
+  printf("Vector forward mode w.r.t. all:\n darr = {%.2g, %.2g, %.2g}\n "
+         "dweights = "
+         "{%.2g, %.2g, %.2g}\n",
+         d_arr[0], d_arr[1], d_arr[2], d_weights[0], d_weights[1],
+         d_weights[2]);
+
+  return 0;
+}

include/clad/Differentiator/Array.h

@@ -154,44 +154,44 @@ template <typename T> class array {
     return *this;
   }
   /// Initializes the clad::array to the given clad::array_ref
-  CUDA_HOST_DEVICE array<T>& operator=(array_ref<T>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator=(const array_ref<T>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] = arr[i];
     return *this;
   }
 
   template <typename U>
-  CUDA_HOST_DEVICE array<T>& operator=(array_ref<U>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator=(const array_ref<U>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] = arr[i];
     return *this;
   }
 
   /// Performs element wise addition
-  CUDA_HOST_DEVICE array<T>& operator+=(array_ref<T>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator+=(const array_ref<T>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] += arr[i];
     return *this;
   }
   /// Performs element wise subtraction
-  CUDA_HOST_DEVICE array<T>& operator-=(array_ref<T>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator-=(const array_ref<T>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] -= arr[i];
     return *this;
   }
   /// Performs element wise multiplication
-  CUDA_HOST_DEVICE array<T>& operator*=(array_ref<T>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator*=(const array_ref<T>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] *= arr[i];
     return *this;
   }
   /// Performs element wise division
-  CUDA_HOST_DEVICE array<T>& operator/=(array_ref<T>& arr) {
+  CUDA_HOST_DEVICE array<T>& operator/=(const array_ref<T>& arr) {
     assert(arr.size() == m_size);
     for (std::size_t i = 0; i < m_size; i++)
       m_arr[i] /= arr[i];

include/clad/Differentiator/ArrayRef.h

@@ -33,7 +33,8 @@ template <typename T> class array_ref {
   /// Constructor for clad::array types
   CUDA_HOST_DEVICE array_ref(array<T>& a) : m_arr(a.ptr()), m_size(a.size()) {}
 
-  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());
     for (std::size_t i = 0; i < m_size; ++i)
       m_arr[i] = a[i];
@@ -53,6 +54,7 @@ template <typename T> class array_ref {
   /// Returns the reference to the location at the index of the underlying
   /// array
   CUDA_HOST_DEVICE T& operator[](std::size_t i) { return m_arr[i]; }
+  CUDA_HOST_DEVICE const T& operator[](std::size_t i) const { return m_arr[i]; }
   /// Returns the reference to the underlying array
   CUDA_HOST_DEVICE T& operator*() { return *m_arr; }
 
@@ -156,6 +158,116 @@ template <typename T> class array_ref {
   }
 };
 
+/// Overloaded operators for clad::array_ref which returns a new clad::array
+/// object.
+
+/// Multiplies the arrays element wise
+template <typename T, typename U>
+CUDA_HOST_DEVICE array<T> operator*(const array_ref<T>& Ar,
+                                    const array_ref<U>& Br) {
+  assert(Ar.size() == Br.size() &&
+         "Size of both the array_refs must be equal for carrying out addition "
+         "assignment");
+  array<T> C(Ar);
+  C *= Br;
+  return C;
+}
+
+/// Adds the arrays element wise
+template <typename T, typename U>
+CUDA_HOST_DEVICE array<T> operator+(const array_ref<T>& Ar,
+                                    const array_ref<U>& Br) {
+  assert(Ar.size() == Br.size() &&
+         "Size of both the array_refs must be equal for carrying out addition "
+         "assignment");
+  array<T> C(Ar);
+  C += Br;
+  return C;
+}
+
+/// Subtracts the arrays element wise
+template <typename T, typename U>
+CUDA_HOST_DEVICE array<T> operator-(const array_ref<T>& Ar,
+                                    const array_ref<U>& Br) {
+  assert(Ar.size() == Br.size() &&
+         "Size of both the array_refs must be equal for carrying out addition "
+         "assignment");
+  array<T> C(Ar);
+  C -= Br;
+  return C;
+}
+
+/// Divides the arrays element wise
+template <typename T, typename U>
+CUDA_HOST_DEVICE array<T> operator/(const array_ref<T>& Ar,
+                                    const array_ref<U>& Br) {
+  assert(Ar.size() == Br.size() &&
+         "Size of both the array_refs must be equal for carrying out addition "
+         "assignment");
+  array<T> C(Ar);
+  C /= Br;
+  return C;
+}
+
+/// Multiplies array_ref by a scalar
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator*(const array_ref<T>& Ar, U a) {
+  array<T> C(Ar);
+  C *= a;
+  return C;
+}
+
+/// Multiplies array_ref by a scalar (reverse order)
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator*(U a, const array_ref<T>& Ar) {
+  return Ar * a;
+}
+
+/// Divides array_ref by a scalar
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator/(const array_ref<T>& Ar, U a) {
+  array<T> C(Ar);
+  C /= a;
+  return C;
+}
+
+/// Adds array_ref by a scalar
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator+(const array_ref<T>& Ar, U a) {
+  array<T> C(Ar);
+  C += a;
+  return C;
+}
+
+/// Adds array_ref by a scalar (reverse order)
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator+(U a, const array_ref<T>& Ar) {
+  return Ar + a;
+}
+
+/// Subtracts array_ref by a scalar
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator-(const array_ref<T>& Ar, U a) {
+  array<T> C(Ar);
+  C -= a;
+  return C;
+}
+
+/// Subtracts array_ref by a scalar (reverse order)
+template <typename T, typename U,
+          typename std::enable_if<std::is_arithmetic<U>::value, int>::type = 0>
+CUDA_HOST_DEVICE array<T> operator-(U a, const array_ref<T>& Ar) {
+  array<T> C(Ar.size(), a);
+  C -= Ar;
+  return C;
+}
+
   /// `array_ref<void>` specialisation is created to be used as a placeholder
   /// type in the overloaded derived function. All `array_ref<T>` types are
   /// implicitly convertible to `array_ref<void>` type.

include/clad/Differentiator/BaseForwardModeVisitor.h

@@ -39,7 +39,8 @@ class BaseForwardModeVisitor
 
   static bool IsDifferentiableType(clang::QualType T);
 
-  StmtDiff VisitArraySubscriptExpr(const clang::ArraySubscriptExpr* ASE);
+  virtual StmtDiff
+  VisitArraySubscriptExpr(const clang::ArraySubscriptExpr* ASE);
   StmtDiff VisitBinaryOperator(const clang::BinaryOperator* BinOp);
   StmtDiff VisitCallExpr(const clang::CallExpr* CE);
   StmtDiff VisitCompoundStmt(const clang::CompoundStmt* CS);

include/clad/Differentiator/Differentiator.h

@@ -206,7 +206,6 @@ namespace clad {
       else
         return "<invalid>";
     }
-
     void dump() const {
       printf("The code is: \n%s\n", getCode());
     }

include/clad/Differentiator/FunctionTraits.h

@@ -719,8 +719,10 @@ namespace clad {
   using ExtractDerivedFnTraitsForwMode_t =
       typename ExtractDerivedFnTraitsForwMode<F>::type;
 
+  // OutputVecParamType is used to deduce the type of derivative arguments
+  // for vector forward mode.
   template <class T, class R> struct OutputVecParamType {
-    using type = typename std::add_pointer<R>::type;
+    using type = array_ref<typename std::remove_pointer<R>::type>;
   };
 
   template <class T, class R>

include/clad/Differentiator/VectorForwardModeVisitor.h

@@ -16,6 +16,11 @@ class VectorForwardModeVisitor : public BaseForwardModeVisitor {
   /// m_Variables map because all other intermediate variables will have
   /// derivatives as vectors.
   std::unordered_map<const clang::ValueDecl*, clang::Expr*> m_ParamVariables;
+  /// Expression for total number of independent variables. This also includes
+  /// the size of array independent variables which will be inferred from the
+  /// size of the corresponding clad array they provide at runtime for storing
+  /// the derivatives.
+  clang::Expr* m_IndVarCountExpr;
 
 public:
   VectorForwardModeVisitor(DerivativeBuilder& builder);
@@ -68,6 +73,8 @@ class VectorForwardModeVisitor : public BaseForwardModeVisitor {
   /// For example: for size = 4, the returned expression is: {0, 0, 0, 0}
   clang::Expr* getZeroInitListExpr(size_t size, clang::QualType type);
 
+  StmtDiff
+  VisitArraySubscriptExpr(const clang::ArraySubscriptExpr* ASE) override;
   StmtDiff VisitReturnStmt(const clang::ReturnStmt* RS) override;
   // Decl is not Stmt, so it cannot be visited directly.
   VarDeclDiff DifferentiateVarDecl(const clang::VarDecl* VD) override;

include/clad/Differentiator/VisitorBase.h

@@ -483,6 +483,16 @@ namespace clad {
     clang::TemplateDecl* GetCladArrayDecl();
     /// Create clad::array<T> type.
     clang::QualType GetCladArrayOfType(clang::QualType T);
+    /// Find declaration of clad::matrix templated type.
+    clang::TemplateDecl* GetCladMatrixDecl();
+    /// Create clad::matrix<T> type.
+    clang::QualType GetCladMatrixOfType(clang::QualType T);
+    /// Creates the expression clad::matrix<T>::identity(Args) for the given
+    /// type and args.
+    clang::Expr*
+    BuildIdentityMatrixExpr(clang::QualType T,
+                            llvm::MutableArrayRef<clang::Expr*> Args,
+                            clang::SourceLocation Loc);
     /// Creates the expression Base.size() for the given Base expr. The Base
     /// expr must be of clad::array_ref<T> type
     clang::Expr* BuildArrayRefSizeExpr(clang::Expr* Base);