Fix support for non-type template params

include/clad/Differentiator/VisitorBase.h

@@ -320,14 +320,14 @@ namespace clad {
     /// \n Variable declaration cannot be added to code directly, instead we
     /// have to build a declaration staement.
     /// \param[in] D The declaration to build a declaration statement from.
-    /// \returns The declration statement expression corresponding to the input
+    /// \returns The declaration statement expression corresponding to the input
     /// variable declaration.
     clang::DeclStmt* BuildDeclStmt(clang::Decl* D);
     /// Wraps a set of declarations in a DeclStmt.
     /// \n This function is useful to wrap multiple variable declarations in one
     /// single declaration statement.
     /// \param[in] D The declarations to build a declaration statement from.
-    /// \returns The declration statemetn expression corresponding to the input
+    /// \returns The declaration statement expression corresponding to the input
     /// variable declaration.
     clang::DeclStmt* BuildDeclStmt(llvm::MutableArrayRef<clang::Decl*> DS);
 

lib/Differentiator/BaseForwardModeVisitor.cpp

@@ -512,9 +512,13 @@ BaseForwardModeVisitor::DerivePushforward(const FunctionDecl* FD,
 }
 
 StmtDiff BaseForwardModeVisitor::VisitStmt(const Stmt* S) {
+  // if S is a SubstNonTypeTemplateParmExpr, we need to visit its replacement
+  // expression instead.
+  if (const auto* SNTTPE = dyn_cast<SubstNonTypeTemplateParmExpr>(S))
+    return Visit(SNTTPE->getReplacement());
+  // Unknown stmt, just clone it and return a warning.
   diag(DiagnosticsEngine::Warning, S->getBeginLoc(),
        "attempted to differentiate unsupported statement, no changes applied");
-  // Unknown stmt, just clone it.
   return StmtDiff(Clone(S));
 }
 

lib/Differentiator/ReverseModeVisitor.cpp

@@ -738,11 +738,15 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     }
   }
   StmtDiff ReverseModeVisitor::VisitStmt(const Stmt* S) {
+    // if S is a SubstNonTypeTemplateParmExpr, we need to visit its replacement
+    // expression instead.
+    if (const auto* SNTTPE = dyn_cast<SubstNonTypeTemplateParmExpr>(S))
+      return Visit(SNTTPE->getReplacement());
+    // Unknown stmt, just clone it and return a warning.
     diag(
         DiagnosticsEngine::Warning,
         S->getBeginLoc(),
         "attempted to differentiate unsupported statement, no changes applied");
-    // Unknown stmt, just clone it.
     return StmtDiff(Clone(S));
   }
 

test/FirstDerivative/BasicArithmeticMulDiv.C

@@ -104,6 +104,23 @@ double m_10(double x, bool flag) {
 // CHECK-NEXT:   return flag ? (((_d_x = _d_x * 2 + x * 0) , (x *= 2)) , (_d_x * x + x * _d_x)) : (((_d_x += 0) , (x += 1)) , (_d_x * x + x * _d_x));
 // CHECK-NEXT: }
 
+template<size_t N>
+double m_11(double x) {
+  const size_t maxN = 53;
+  const size_t m = maxN < N ? maxN : N;
+  return x*m;
+}
+
+// CHECK: double m_11_darg0(double x) {
+// CHECK-NEXT:   double _d_x = 1;
+// CHECK-NEXT:   const size_t _d_maxN = 0;
+// CHECK-NEXT:   const size_t maxN = 53;
+// CHECK-NEXT:   bool _t0 = maxN < 64UL;
+// CHECK-NEXT:   const size_t _d_m = _t0 ? _d_maxN : 0UL;
+// CHECK-NEXT:   const size_t m = _t0 ? maxN : 64UL;
+// CHECK-NEXT:   return _d_x * m + x * _d_m;
+// CHECK-NEXT: }
+
 int d_1(int x) {
   int y = 4;
   return y / y; // == 0
@@ -190,6 +207,7 @@ double m_7_darg0(double x);
 double m_8_darg0(double x);
 double m_9_darg0(double x);
 double m_10_darg0(double x, bool flag);
+double m_11_darg0(double x);
 int d_1_darg0(int x);
 int d_2_darg0(int x);
 int d_3_darg0(int x);
@@ -230,6 +248,9 @@ int main () {
   printf("Result is = %f\n", m_10_darg0(1, true)); // CHECK-EXEC: Result is = 8
   printf("Result is = %f\n", m_10_darg0(1, false)); // CHECK-EXEC: Result is = 4
 
+  clad::differentiate(m_11<64>, 0);
+  printf("Result is = %f\n", m_11_darg0(1)); // CHECK-EXEC: Result is = 53
+
   clad::differentiate(d_1, 0);
   printf("Result is = %d\n", d_1_darg0(1)); // CHECK-EXEC: Result is = 0
 

test/Gradient/Gradients.C

@@ -761,6 +761,27 @@ void fn_increment_in_return_grad(double i, double j, clad::array_ref<double> _d_
 // CHECK-NEXT:     * _d_i += _d_temp;
 // CHECK-NEXT: }
 
+template<size_t N>
+double fn_template_non_type(double x) {
+  const size_t maxN = 53;
+  const size_t m = maxN < N ? maxN : N;
+  return x*m;
+}
+
+// CHECK: void fn_template_non_type_grad(double x, clad::array_ref<double> _d_x) {
+// CHECK-NEXT:     size_t _d_maxN = 0;
+// CHECK-NEXT:     bool _cond0;
+// CHECK-NEXT:     size_t _d_m = 0;
+// CHECK-NEXT:     const size_t maxN = 53;
+// CHECK-NEXT:     _cond0 = maxN < 15UL;
+// CHECK-NEXT:     const size_t m = _cond0 ? maxN : 15UL;
+// CHECK-NEXT:     goto _label0;
+// CHECK-NEXT:   _label0:
+// CHECK-NEXT:     * _d_x += 1 * m;
+// CHECK-NEXT:     if (_cond0)
+// CHECK-NEXT:         _d_maxN += _d_m;
+// CHECK-NEXT: }
+
 #define TEST(F, x, y)                                                          \
   {                                                                            \
     result[0] = 0;                                                             \
@@ -812,4 +833,9 @@ int main() {
   TEST_GRADIENT(fn_global_var_use, /*numOfDerivativeArgs=*/2, 3, 5, &d_i, &d_j);  // CHECK-EXEC: {7.00, 0.00}
 
   TEST(fn_increment_in_return, 3, 2); // CHECK-EXEC: Result is = {7.00, 0.00}
+
+  auto fn_template_non_type_dx = clad::gradient(fn_template_non_type<15>);
+  double x = 5, dx = 0;
+  fn_template_non_type_dx.execute(x, &dx);
+  printf("Result is = %.2f\n", dx); // CHECK-EXEC: Result is = 15.00
 }