thank you! For our MLIR use case this is exactly what we need.

Rebase

Fix bugs

There are two places where in MLIR you can put an attribute to coroutine functions:

Build is failing: https://buildkite.com/llvm-project/premerge-checks/builds/70123#b97ba1d4-10ba-4287-899a-0bcf5260aefe

Reset number of unrallable loops if decide not to uroll

Cleanup

Fix LLVM style violations

Address comments

Revert accidental changes

It would be great to have the cost model to be a part of core MLIR, defined with op interface, but for it to be usable for all the potential targets it will require a lot of discussions (e.g. how cost model can capture GPU memory coalescing, should it be the concern of the cost model at all, or it should just track compute?). Let's start an RFC to figure out how it should look like. In the meantime we can hide the current crude implementation behind the AsyncBlockSizeComputationFunction API, and keep implementation inside the Tensorflow compiler. We can tune it for the compilation pipeline we have right now, without need to worry about the general case.

I have a failing test for this patch: Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

Basically add a pass something like X86vectorMathApproximationsPass?

But avx and avx512 rsqrt have different accuracy

Maybe we can have generic x86vector.rsqrt operation (without avx or avx512 prefix), and then it can be unrolled and lowered to ISA specific intrinsics by the x86dialect passes? /cc @aartbik

In D112736#3094908, @silvas wrote:

Can we move this RsqrtOp lowering to the x86vector dialect? It seems really out of place in the math dialect, which should remain target agnostic. (also, since this approximation is newton-raphson based instead of polynomial based, it seems out of place in PolynomialApproximation.cpp as well)

In D112736#3094843, @silvas wrote:

This doesn't seem to be the right pass to do target-specific lowerings. Can you add math.rsqrt and move the vector unrolling to a generic pass in the vector dialect? Then this pass wouldn't depend on target-details.

Address comments

Included this fix in: https://reviews.llvm.org/D112736

There is one failed test

Can you please sync to head and re-upload, can't apply the patch to the head

We should migrate it to min/max after https://github.com/llvm/llvm-project/commit/9fb57c8c1dd87df36daf7b6f7dee3b7423475afc