Excessive memory allocation about enzyme

[wsmoses]

What happens if you mark out and x as restrict. Without having looked deeply at it, my guess is that without other information (i.e. restrict) the compiler has to assume that out and x could point to a location in the same memory block, and thus has to cache all of x.

[wsmoses]

Since you're using the clang plugin, you may also find adding the flag -Rpass=enzyme useful as it will (sometimes excessively) spit out information about when Enzyme may need to cache something, and why.

[unrealwill]

Good guess (I had forgotten about __restrict__ ) , but unfortunately the problem persist after I've updated the code to use __restrict__ (and edited my original post to reflect it)

I have added the option to obtain more information, but I'm not sure what to decipher from it.
Maybe some memory alignment due to vectorization.

clang bugStackAlloc.cpp  -lstdc++ -lm  -fno-exceptions -Rpass=enzyme -Xclang -load -Xclang /usr/local/lib/ClangEnzyme-11.so  -O2  -o bin/bugStackAlloc
remark: Load may need caching   %out.promoted = load double, double* %out, align 8, !tbaa !10 due to   store double %add18, double* %out, align 8, !dbg !26, !tbaa !10 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub = fsub double %2, %3, !dbg !29 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
       *out += (buf[i]-buf[i+1] )*(buf[i]-buf[i+1] ) ;
                      ^
remark: Load may need caching   %out.promoted = load double, double* %out, align 8, !tbaa !11, !alias.scope !18, !noalias !15 due to   store double %add18.i.27, double* %out, align 8, !tbaa !11, !alias.scope !18, !noalias !15 [-Rpass=enzyme]
bugStackAlloc.cpp:27:17: remark: Load may need caching   %.pre.i = load double, double* %arrayidx8.phi.trans.insert.i, align 8, !dbg !26, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
       *out += (buf[i]-buf[i+1] )*(buf[i]-buf[i+1] ) ;
                ^
bugStackAlloc.cpp:27:24: remark: Load may need caching   %3 = load double, double* %arrayidx10.i, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
       *out += (buf[i]-buf[i+1] )*(buf[i]-buf[i+1] ) ;
                       ^
bugStackAlloc.cpp:27:24: remark: Load may need caching   %4 = load double, double* %arrayidx10.i.1, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %5 = load double, double* %arrayidx10.i.2, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %6 = load double, double* %arrayidx10.i.3, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %7 = load double, double* %arrayidx10.i.4, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %8 = load double, double* %arrayidx10.i.5, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %9 = load double, double* %arrayidx10.i.6, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %10 = load double, double* %arrayidx10.i.7, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %11 = load double, double* %arrayidx10.i.8, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %12 = load double, double* %arrayidx10.i.9, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %13 = load double, double* %arrayidx10.i.10, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %14 = load double, double* %arrayidx10.i.11, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %15 = load double, double* %arrayidx10.i.12, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %16 = load double, double* %arrayidx10.i.13, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %17 = load double, double* %arrayidx10.i.14, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %18 = load double, double* %arrayidx10.i.15, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %19 = load double, double* %arrayidx10.i.16, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %20 = load double, double* %arrayidx10.i.17, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %21 = load double, double* %arrayidx10.i.18, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %22 = load double, double* %arrayidx10.i.19, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %23 = load double, double* %arrayidx10.i.20, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %24 = load double, double* %arrayidx10.i.21, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %25 = load double, double* %arrayidx10.i.22, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %26 = load double, double* %arrayidx10.i.23, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %27 = load double, double* %arrayidx10.i.24, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %28 = load double, double* %arrayidx10.i.25, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %29 = load double, double* %arrayidx10.i.26, align 16, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:24: remark: Load may need caching   %30 = load double, double* %arrayidx10.i.27, align 8, !dbg !20, !tbaa !11, !noalias !24 due to   store double %mul.i, double* %arrayidx1.i, align 8, !dbg !38, !tbaa !11, !noalias !24 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.27 = fsub double %28, %29, !dbg !27 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
       *out += (buf[i]-buf[i+1] )*(buf[i]-buf[i+1] ) ;
                      ^
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.26 = fsub double %27, %28, !dbg !28 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.25 = fsub double %26, %27, !dbg !29 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.24 = fsub double %25, %26, !dbg !30 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.23 = fsub double %24, %25, !dbg !31 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.22 = fsub double %23, %24, !dbg !32 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.21 = fsub double %22, %23, !dbg !33 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.20 = fsub double %21, %22, !dbg !34 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.19 = fsub double %20, %21, !dbg !35 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.18 = fsub double %19, %20, !dbg !36 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.17 = fsub double %18, %19, !dbg !37 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.16 = fsub double %17, %18, !dbg !38 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.15 = fsub double %16, %17, !dbg !39 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.14 = fsub double %15, %16, !dbg !40 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.13 = fsub double %14, %15, !dbg !41 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.12 = fsub double %13, %14, !dbg !42 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.11 = fsub double %12, %13, !dbg !43 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.10 = fsub double %11, %12, !dbg !44 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.9 = fsub double %10, %11, !dbg !45 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.8 = fsub double %9, %10, !dbg !46 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.7 = fsub double %8, %9, !dbg !47 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.6 = fsub double %7, %8, !dbg !48 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.5 = fsub double %6, %7, !dbg !49 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.4 = fsub double %5, %6, !dbg !50 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.3 = fsub double %4, %5, !dbg !51 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.2 = fsub double %3, %4, !dbg !52 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i.1 = fsub double %2, %3, !dbg !53 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
bugStackAlloc.cpp:27:23: remark: Caching instruction   %sub.i = fsub double %.pre.i, %2, !dbg !54 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
remark: Load may need caching   %out.promoted = load double, double* %out, align 8, !tbaa !9 due to   store double %add26, double* %out, align 8, !dbg !15, !tbaa !9 [-Rpass=enzyme]
bugStackAlloc.cpp:53:20: remark: Load may need caching   %.pre = load double, double* %arrayidx12.phi.trans.insert, align 8, !dbg !27, !tbaa !9 due to   store double %mul, double* %arrayidx5, align 8, !dbg !22, !tbaa !9 [-Rpass=enzyme]
         temp +=  (buf[i]-buf[i+1] ) * (buf[i]-buf[i+1] ) ;
                   ^
bugStackAlloc.cpp:53:27: remark: Load may need caching   %5 = load double, double* %arrayidx14, align 8, !dbg !34, !tbaa !9 due to   store double %mul, double* %arrayidx5, align 8, !dbg !22, !tbaa !9 [-Rpass=enzyme]
         temp +=  (buf[i]-buf[i+1] ) * (buf[i]-buf[i+1] ) ;
                          ^
bugStackAlloc.cpp:53:26: remark: Caching instruction   %sub = fsub double %3, %4, !dbg !34 legalRecompute: 1 shouldRecompute: 0 tryLegalRecomputeCheck: 1 [-Rpass=enzyme]
         temp +=  (buf[i]-buf[i+1] ) * (buf[i]-buf[i+1] ) ;
                         ^
testMemoryAllocator 
demofun1() 
out without enzyme
252
out with enzyme
252
dx with enzyme
168
demofun2() 
out without enzyme
2.52e+10
Killed

[unrealwill]

After more thinking, I think it's just a problem with enzyme's recomputing heuristic and picking the wrong choice.

In this case I'd like enzyme to cache buf, but only temporary while computing Fun1, and Fun1Adj.
It is possible if Enzyme takes notice that x is not modified in the loop.

Here is the manual code I'd like enzyme to do automatically in this specific case :

#include <iostream>
using namespace std;

extern int enzyme_dup;
extern int enzyme_dupnoneed;
extern int enzyme_out;
extern int enzyme_const;

void __enzyme_autodiff(...);

template<typename T>
void comp( double* __restrict__ x,  double*  __restrict__ out);

const int nbiter = 100000000;
const int d = 30;

class Fun1{};template<> void comp<Fun1>( double* __restrict__ x, double*__restrict__ out )
{
    double buf[d];
    double xinp = x[0];
    for( int i = 0 ; i < d ; i++)
    {
        buf[i] = xinp * i;
    }

    for( int i =1 ; i < d-1 ; i++)
    {
          //One way of doing it would be to recompute bufi but that's not the best choice here
          //double bufi = xinp*i;
          //double bufip = xinp*(i+1);
          //*out+= (bufi-bufip )*(bufi-bufip ) ;
          *out+= (buf[i]-buf[i+1] )*(buf[i]-buf[i+1] ) ;
    }

}

void fun2adjoint( double* x, double* dx, double * out, double* dout)
{
  for( int j = nbiter-1 ; j >= 0 ; j--)
  {
    __enzyme_autodiff( comp<Fun1>, enzyme_dup, x, dx,
                                           enzyme_dup,out,dout);
  }
}

/*
//This adjoint doesn't use the commutativity of + but need an additional forward pass
void fun2adjoint( double* x, double* dx, double * out, double* dout)
{
  for( int j = 0 ; j < nbiter ; j++)
  {
     comp<Fun1>(x, out);
  }
  double fout = 0;
  for( int j = nbiter-1 ; j >= 0 ; j--)
  {
    __enzyme_autodiff( comp<Fun1>, enzyme_dup, x, dx,
                                           enzyme_dupnoneed,&fout,dout);
  }
} */

class Fun2{};template<> void comp<Fun2>(  double*  __restrict__ x,  double*__restrict__  out  )
{
  for( int  j = 0; j < nbiter ; j++)
  {
     comp<Fun1>(x,out);
  }
}

template< typename T>
void demo( )
{
  double x = 3.0;
  {
  double out = 0.0;
  comp<T>( &x, &out );
  cout << "out without enzyme" << endl;
  cout << out << endl;

  }

  {
  double dx = 0.0;
  double gout = 1.0;
  double out = 0.0;

  __enzyme_autodiff( comp<T>, enzyme_dup, &x, &dx,
                                         enzyme_dup,&out,&gout);

  cout << "out with enzyme" << endl;
  cout << out << endl;
  cout << "dx with enzyme" << endl;
  cout << dx << endl;

  }
}

void demoManual()
{
  double x = 3.0;
  {
  double dx = 0.0;
  double gout = 1.0;
  double out = 0.0;

  fun2adjoint( &x, &dx,&out,&gout);

  cout << "out with enzyme" << endl;
  cout << out << endl;
  cout << "dx with enzyme" << endl;
  cout << dx << endl;
  }

}

int main(int argc, char** argv )
{
  cout << "testMemoryAllocator "<< endl;
  cout << "demofun1() " << endl;
  demo<Fun1>();
  cout << "demo manual" << endl;
  demoManual();
  cout << "demofun2() " << endl;
  demo<Fun2>(); // Fails
}

[wsmoses]

I agree with your assessment. In this case it is beneficial to do this. In other cases [e.g. a recomputable function call], this may not be beneficial since the function call may take longer to compute than the overhead incurred by the caching. As such designing what conditions the checkpointing should apply here will be critical [as well as the actual implementation].

[wsmoses]

Update to this, the first level of caching here has been implemented (namely at a function level). Specifically, if buf is in a separate function (like if Fun1 isn't inlined) this will recompute buf in the reverse pass. We haven't yet adding the support for loops (required for fun3), but that should come soon.

[unrealwill]

@wsmoses Thanks

[wsmoses]

This should not allocate anymore on latest main (and feel free to reopen otherwise)