Disallow async nestings that violate read after write dependencies

[abadams]

Fixes #7867
Fixes #5175

Discovered while trying to improve the async schedule in the async_gpu performance test so that it's less flaky.

[zvookin]

Does this have any connection to #5202 ? There is also #5195 . It would be great to either get those moved forward or closed. @vksnk .

[abadams]

Looks strongly related to #5202, which I had forgotten about. Not sure about #5195

[vksnk]

#5202 looks related for sure, but #5195 is a sepatate issue.

I was planning to ressurect the PRs above (finally have some spare cycles), but I don't think this PR should be blocked on these. I can update accordingly once this is merged in.

[abadams]

I think this is the same as #5202

I like this solution slightly better, because it happens after tightening producer consumer nodes so it seems less possible for it to capture false dependencies. Open to other opinions though. I should at least incorporate the test from #5202 , and this PR has the error test but not the correctness test.

[abadams]

This also fixes #5201 (makes it error out), but that revealed my error message is wrong. I think the constraint is that if you have a producer consumer pair, where the consumer is async, then the store_at of the producer can never be between the compute_at and store_at levels of the consumer, and the compute_at of the producer can only be between the compute_at and store_at levels of the consumer if the producer is also async.

[abadams]

Correction:

If you have a producer consumer pair, where the consumer is async, and the compute_at of the producer is in between the store_at and compute_at of the consumer, then the producer must be both async and store_at outside the store_at level of the consumer.

[abadams]

I cherry-picked the tests from #5201, and one of them fails, so clearly something is still wrong :(

[vksnk]

The only thing I remember is that this bug was pretty tricky and required some graph level reasoning, but the rest is pretty foggy for now:( I'll need to reread the PR to remember it again.

I also remember there was another test I found which was failing, it's not in PR (all tests in the PR should be passing), but I sort of vaguely remember what was happening there. I'm going to try to reproduce it again.

[abadams]

I think the failing test is actually caused by the new sliding window behavior. There are too many release calls to the storage folding semaphores because the consumer has acquired extra loop iterations.

[vksnk]

I think I found a failing test, it's hanging in my branch for #5202:

    {
        Func producer1, producer2, consumer;
        Var x, y;

        producer1(x, y) = x + y;
        producer2(x, y) = producer1(x, y);
        consumer(x, y) = producer1(x, y - 1) + producer2(x, y + 1);

        consumer.compute_root();

        producer1.store_at(consumer, Var::outermost()).compute_at(consumer, y).async();
        producer2.store_root().compute_at(consumer, y).async();
        consumer.bound(x, 0, 16).bound(y, 0, 16);

        Buffer<int> out = consumer.realize({16, 16});

        out.for_each_element([&](int x, int y) {
            int correct = 2 * (x + y);
            if (out(x, y) != correct) {
                printf("out(%d, %d) = %d instead of %d\n",
                       x, y, out(x, y), correct);
                exit(-1);
            }
        });
    }

[vksnk]

Ok, I just tried it in your branch and it does produce an error message instead of hanging.

[abadams]

Dillon doesn't remember the reasoning behind that line, so I just removed it. This will need testing inside Google to see if it breaks anything.

[vksnk]

Dillon doesn't remember the reasoning behind that line, so I just removed it. This will need testing inside Google to see if it breaks anything.

I can do testing inside Google, if it's ready.

[abadams]

Should be ready. I just did a merge with main.

[steven-johnson]

I'll pull this into Google to test

[steven-johnson]

AFAICT, the async() directive is not used anywhere inside Google

[abadams]

I was under the impression @vksnk was using it

[vksnk]

I thought we wanted to test it in Google, because no one could remember what the removed lines in storage_folding were for.

I am only planning to use async for something, but not quite yet.

[steven-johnson]

There were some usages of it for a long-defunct project, but that code was deleted a while back.

[abadams]

Ah, right. This is a change to storage folding, and it needs to be tested for that reason.

[steven-johnson]

I don't see any google3 failures.

[abadams]

Just needs a review then.

[vksnk]

"to the producer" -> "in the producer"?