[spirv-opt] Reassociation Pass

[alister-chowdhury]

Hey,

Currently folding rules act on a instruction-by-instruction basis, which works for a lot of cases, but struggles to deal with temporaries in the middle of an expression.

e.g:

10 + a + 10       => a + 20
10 + b + a + 10   => no fold

(3 * a) + (-3 *a)     => 0
(3 * a) + 1 + (-3 *a) => no fold

(3 * a) + (3 * b)     => 3 * (a + b)
(3 * a) + c + (3 * b) => no fold

Because each rule would need to account for every possibility and combination, it's not very practical to
make rules for these sorts of situations.

Locally I've been working on a reassociation pass (for FP arithmetic) and was wondering if it is something that you would want (since it's not simply adding extra folding rules).

The general premise is to reassociate independant chains of: OpFAdd, OpFSub, OpFNegate, OpFMul, OpFDiv, OpVectorTimesScalar.

Constants are moved closer together e.g:

  (A + 5 + B - C - 7) => (5 - 7 + A + B - C) => (-2 + A + B - C)
  (A * 5 * B * C * 7) => (5 * 7 * A * B * C) => (35 * A * B * C)

Variables far apart can cancel each other out e.g:

 (A + B + C - A - B)   => C
 (A * B * C / (A * B)) => C

Variables can be factored e.g:

  (A * x * x + B * x + C) => x * (A * x + B) + C
  A * x + B * x + C * x   => x * (A + B + C)

With a few extra folding rules that operate on the chains directly.

Many thanks,
Alister

[s-perron]

We could do something like that. I would require a more through review. The difficulties with reassociation is that if you are trying to make local changes to get to a particular normal form, but different folding rules might have different ways in which they work best.

Here are some other ideas on how to get what you want. See what works best.

Can you add individual folding rules to do the reassociation moving constant "to the left".

(a + 10) -> (10 + a)

and

(10 + a ) + b -> 10 + (a + b)

You can also sort the variables based on their id the same way. This would not be the most efficient sort, but how long are most chains?

We do not currently use the value number table in folding rules, but you could.

(B * A) * C / (A * B)) => C

You could do this without reassociation by checking the value number of the expressions

(expr1 * C) / (expr2) => C if the value numbers for expr1 and expr2 are the same.

Try to think how smaller folding rules could work together as well:

((Ax) + (Bx)) + C*x => (A+B)x + Cx => ((A+B)+C) * x.

[alister-chowdhury]

Can you add individual folding rules to do the reassociation moving constant "to the left".

(a + 10) -> (10 + a)

After #6371, most commutative expressions should already get de-duplicated and from what I remember all the folding rules
handled both sides.

(10 + a ) + b -> 10 + (a + b)

I think this would be possible, where you hoist expressions downstream.
The main downside would be that you constantly need to create temporary instructions (bloating the result_id) and having to check if any other instruction was dependent each time.

You can also sort the variables based on their id the same way. This would not be the most efficient sort, but how long are most chains?

Most chains are typically short, but if they're smaller than say 4 instructions, they should probably be ignored
for reassociation anyway.

But they can absolutely get very large, take something like:

float LocalBias;
float GlobalBias;
StructuredBuffer<float2>   Inputs;
RWStructuredBuffer<float>  Output;

float AccumBias(float x , float y)
{
    return ((x * LocalBias) + y) / GlobalBias;
}

[numthreads(64, 1, 1)]
void test_0(uint index : SV_DispatchThreadID)
{
    float Integrated = 0.0;
    [unroll]
    for(uint i=0; i < 8; ++i)
    {
        float2 V = Inputs[index * 8 + i];
        Integrated += AccumBias(V.x, V.y);
    }
    Output[index] = Integrated;
}

As a human, looking at this in detail, I know this can be reduced to:
(sum(x)*LocalBias + sum(y)) / GlobalBias

But to encapsulate with folding rules alone would be extremely difficult.

Before

After

(B * A) * C / (A * B)) => C

This specific case is already handled in the folding rules.

But something like:
((A * B) * ((C * D) * E) / (B * D) = (A * C * E)

Would not be and in order to really do that, I think you'd need to evaluate the entire dependency chain and might as well just do it with a dedicated pass.

---

A large part of the motivation of this comes from dealing with material graphs authored by artists who (understandably) don't typically think how engine provided nodes will gel together.

Another example of the sort of random things you see in materials:

StructuredBuffer<float2>    in_data;
RWStructuredBuffer<float2>  data_out;

float CustomFunction0(float x)
{
    float _0 = (x * 2.000000);
    float _1 = _0 + 5.000000;
    float _2 = (x * 5.000000);
    float _3 = _1 + _2;
    return _3;
}

[numthreads(64, 1, 1)]
void MainCS(uint index : SV_DispatchThreadID)
{
    float2 uv0 = in_data[index * 2 + 0];
    float2 uv1 = in_data[index * 2 + 1];
    float u = CustomFunction0(uv0.x) + CustomFunction0(uv1.y);
    float v = CustomFunction0(uv1.x) + CustomFunction0(uv0.y);
    data_out[index] = float2(u, v);
}

Before

After