How to update a vector and reuse it for the same calculation (c++)?

[CuddlyCat]

I want to solve a differential equation

df(i)/dt = A(i,j,k) * f(j) * f(k)

with A a (sparse,sparse,sparse) 3-fold tensor and f a dense vector with TACO. I am using a numerical scheme that takes initial values for f, a set time step size dt and the number of time steps you want to make. Essentially, you can calculate the progression of f according to the differential equation by calculating for each time step fnext(i) = fnow(i)+ param() A(i,j,k) * fnow(j) * fnow(k) with fnow the vector at the current time t and fnext the vector at the time t+dt. This operation itself works well with Taco, after initializing the values of fnow I define the calculation of fnext and with fnext.compile(); fnext.assemble(); fnext.compute(); one calculation works.

My issue is, that I want to compute the same calculation multiple times, updating the values of fnow after each calculation. Ideally, something like this should work:

fnext(i) = fnow(i)+ param() A(i,j,k) * fnow(j) * fnow(k)
fnext.compile();
for(int timestep=1; timestep<=totalTimesteps; timestep++){
   fnext.assemble();
   fnext.compute();
   fnow=fnext; //placeholder for some kind of updating...
}

I have tried the different ways to update fnow that I could gather from the online documentation, but whatever I do, it looks like the updating is never successful. Instead the calculation is executed totalTimesteps times always with fnow set to the initial values. Notice, that I am repeatedly using exactly the same tensor A(i,j,k) and the compilation and assembly takes quite some computational time. I want to use Taco to optimize computational time, and am confused that the updating step seems to be quite complicated. I can usually solve this with some pointer switching, but here I am stumped.

How can I update the values in fnow efficiently so that the computation of fnext uses the new values?

I have written c++ code and compile and execute it with the Taco-library. Ideally, Taco could generate a function calculateTimestep(&fnext,&fnow) for a given A(i,j,k), that performs the calculation fnext(i) = fnow(i)+ param() A(i,j,k) * fnow(j) * fnow(k) and that I can call repeatedly with different reference arguments fnext and fnow. Is there a way to do that?

[stephenchouca]

I think you have the right general idea, but the problem is that the computation is defined outside of the loop, which is why it is always using the initial fnow. If you instead define the computation within the loop, then each execution of the computation will use the updated fnow as intended; the example below shows how that would work:

int main() {
  Tensor<double> a({3}, Dense);
  Tensor<double> c({3}, Dense);
  c.insert({0}, 1.0);
  c.insert({1}, 2.0);
  c.insert({2}, 3.0);
  for (int t = 0; t < 10; ++t) {
    Tensor<double> b({3}, Dense);
    IndexVar i;
    b(i) = a(i) + c(i);
    // The next three (commented-out) lines are invoked implicitly when `a = b` is executed:
    // b.compile();
    // b.assemble();
    // b.compute();
    a = b;
  }
  std::cout << a << std::endl;  // `a` should equal [10, 20, 30]
  return 0;
}

Note that even though compile is invoked in every iteration, since the kernel being compiled is the same, the TACO library will actually cache and reuse the compiled code under the hood. This means that only the first iteration of the loop will actually incur additional overhead to generate and compile code.

[skullbaselab]

Very cool.