Use local var to pop output arg value included in for loop so not to …

…corrupt the forw pass
kchristin22 · Dec 15, 2024 · 31be5d8 · 31be5d8
1 parent 11b7d8e
commit 31be5d8
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/benchmark/RSBench/Makefile b/benchmark/RSBench/Makefile
@@ -93,4 +93,4 @@ edit:
 	vim -p $(source) rsbench.cuh
 
 run:
-	./rsbench -m event -l 102000 > output.txt
+	./rsbench -m event -l 102000
diff --git a/benchmark/RSBench/simulation.cu b/benchmark/RSBench/simulation.cu
@@ -169,8 +169,8 @@ __device__ void body( int i, double * __restrict__ macro_xs, int mat, double E,
 __device__ void calculate_macro_xs( double * __restrict__ macro_xs, int mat, double E, Input input, int * __restrict__ num_nucs, int * __restrict__ mats, int max_num_nucs, double * __restrict__ concs, int * __restrict__ n_windows, double * __restrict__ pseudo_K0Rs, Window * __restrict__ windows, Pole * __restrict__ poles, int max_num_windows, int max_num_poles ) 
 {
 	// zero out macro vector
-	for( int i = 0; i < 4; i++ )
-		macro_xs[i] = 0;
+	// for( int i = 0; i < 4; i++ )
+	// 	macro_xs[i] = 0;
 
 	// for nuclide in mat
 	int sz = num_nucs[mat];
@@ -438,7 +438,7 @@ __device__ RSComplex fast_nuclear_W( RSComplex Z )
 		RSComplex one = {1, 0};
 		RSComplex W = c_div(c_mul(i, ( c_sub(one, fast_cexp(c_mul(t1, Z))) )) , c_mul(t2, Z));
 		RSComplex sum = {0,0};
-		#pragma unroll
+		// #pragma unroll
 		for( int n = 0; n < 10; n++ )
 		{
 			RSComplex t3 = {neg_1n[n], 0};

diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp
@@ -3407,7 +3407,14 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     assert(E && "must be provided");
     auto Type = getNonConstType(E->getType(), m_Context, m_Sema);
 
-    if (isInsideLoop) {
+    bool isOutVar = false;
+    if (auto arraySub = dyn_cast<ArraySubscriptExpr>(E))
+      if (auto declRef = dyn_cast<DeclRefExpr>(
+              arraySub->getBase()->IgnoreImpCasts()->IgnoreParens()))
+        if (m_Variables.find(declRef->getDecl()) != m_Variables.end())
+          isOutVar = true;
+
+    if (isInsideLoop && !isOutVar) {
       auto CladTape = MakeCladTapeFor(Clone(E), prefix);
       Expr* Push = CladTape.Push;
       Expr* Pop = CladTape.Pop;
@@ -3438,6 +3445,14 @@ Expr* getArraySizeExpr(const ArrayType* AT, ASTContext& context,
     if (E->isModifiableLvalue(m_Context) == Expr::MLV_Valid)
       Restore = BuildOp(BO_Assign, Clone(E), Ref);
 
+    if (isInsideLoop) {
+      auto CladTape = MakeCladTapeFor(Clone(E), prefix);
+      Expr* Push = CladTape.Push;
+      Expr* Pop = CladTape.Pop;
+      auto* popAssign = BuildOp(BinaryOperatorKind::BO_Assign, Ref, Pop);
+      return {Push, popAssign};
+    }
+
     return {Store, Restore};
   }