TiledArray code generation

[bimalgaudel]

Consider the following sum of tensor networks:

$$(A(a_1 x_1) B(x_1 x_2)) C(x_2 b_1) + (A(a_1 y_1) B(y_1 y_2)) D(y_2 b_1) + (C(a_1 x_1) D(x_1 x_2)) E(x_2 b_1)$$

This PR attempts to support TiledArray code generation from it that looks like the following or its equivalent.

// declare the total evaluation result
TArray R;

// first and second term share a common intermediate (A*B)
// evaluate it as part of the first term evaluation

TArray I1;
I1("a1,x2") = A("a1,x1") * B("x1,x2");
R("a1,a2")  = I1("a1,x2") * C("x2,b1");
// done evaluating first term

// reuse the intermediate from first term evaluation
R("a1,a2") += I1("a1,y2") * D("y2,b1");
// done evaluating second term

// third term shares no intermediates
// no need to split up the evaluation like for the first term
R("a1,b1") += (C("a1,x1") * D("x1,x2")) * E("x2,b1");
// done evaluating third term

[bimalgaudel]

@Krzmbrzl I think export framework will be useful here?

I tried a couple approaches to support this in one sitting, but the splitting up of single term evaluation into intermediary steps only when strictly necessary turned out to require non-trivial amount of coding. It is doable for sure, but better made use of the export framework.

@evaleev

[Krzmbrzl]

the splitting up of single term evaluation into intermediary steps only when strictly necessary turned out to require non-trivial amount of coding.

What exactly do you mean by that?

[bimalgaudel]

the splitting up of single term evaluation into intermediary steps only when strictly necessary turned out to require non-trivial amount of coding.

What exactly do you mean by that?

Look at the generated code for the first term and the last term in the example above.
I1 is separately declared and evaluated for the first term.

[Krzmbrzl]

So you want to have binary contractions only when this allows intermediate reuse and otherwise want to use the unfactorized equations? Why?

At the moment the export framework is based on evaluation trees and hence is restricted to binarized contractions as input. Technically, you could probably re-assemble the original expression in the backend and only export that but that seems like a fair amount of work to get right in general.

[bimalgaudel]

So you want to have binary contractions only when this allows intermediate reuse and otherwise want to use the unfactorized equations? Why?

Not exactly what I meant. In TiledArray binary evaluations are happening behind the scenes and evaluation order can be enforced with parenthesization. Consider the following evaluation code for $(A(a_1) B(b_1)) C(c_1)$ using TiledArray.

TArray R; // the final result is in R

TArray I1;
I1("a1,b1") = A("a1") * B("b1");
R("a1,b1,c1") = I1("a1,b1") * C("c1");

However, TiledArray also supports the following:

TArray R;
R("a1,b1,c1") = (A("a1") * B("b1")) * C("c1");

No need to create the temporary I1 and one less thing to worry about life-time/resource management. TArray is a distributed object, and TiledArray knows how to manage its life-time and resources properly.

PS I added the missing parenthesis in the last statement of the generated example code..

[Krzmbrzl]

I see. Does that mean there is an actual disadvantage in creating the intermediates explicitly even in cases where it's not necessary? I mean, other than code readability

But I guess the bottom line is that it can likely be done right now by adding some logic into the backend generator but it kinda goes against the design philosophy of the export framework (backend generators should contain only a minimal amount of logic in order to be easily written).

I would propose focusing on getting the export framework with its current abilities and restrictions merged first and then think of how to generalize it to support non-binary contractions as well. I think this would also be nice for generating code that is interpreted by a specialized tensor compiler.

[evaleev]

I agree, minimal binarized TA code emission should be the initial goal. Dispatch to more general backends capable of evaluating whole TNs is longer term.