Thanks, Sergei! LGTM. @nicolasvasilache, @aartbik, could we please have final approval from anybody from your side? Thanks!

Thanks, Aart! LGTM. Just minor comments.

Thanks for working on this, Sergei! Some comments.

Thanks, Alex!

Thanks!

Thanks for working on this, Aart! I think the progressive lowering approach that you are taking here is very on point! I’m not working with AMX but it would be great if you could add me as a reviewer to the related code reviews. It would be very educational for me since this approach is also applicable to similar internal problems we have.

Thanks for addressing this issue! Would this work if we compiled in release mode and with asserts enabled?

In D97304#2620165, @ezhulenev wrote:

They are actually used but only inside assert, is it a release build when it is a warning?

Sorry, Alex. Let me look into that.

Let's address this in a separate patch since the vector dialect doesn't currently have the pass infrastructure, so it will be quite a bit of somewhat unrelated code.

It looks like elementType and isI32 methods in PolynomialApproximation.cpp are not used. Would it be possible to remove them to avoid the warnings?

Thanks for fixing this bug, Vinayaka! The direction looks good to me. Adding some comments.

Thanks! I'll commit this tomorrow if no more comments.

Addressed feedback, thanks!
I'll commit this tomorrow if no more comments.

Addressed feedback. Thanks!

Addressed feedback. Thanks!

Agreed on the blanket. But how about the simple folding away of unmasked (or all-true/false mask). In the long run, it makes a lot more sense to have one place with that logic (folding/canonicalizaiton) and have all other rewriting rule that expect unmasked transfers simply work on the l/s instead?

For now, I would say that hiding the patterns behind a populate function (but not a canonicalization) would be the simplest way forward.

More ops going through vector loads/stores! Great! Thanks, Aart. LGTM.

In the sense that e.g. unmasked one becomes regular loads, yes. But AFAIK, this is very much in the spirit of MLIR, to canonicalize everything to the simplest form possible, and have patterns in the rewriting rules that expect only the simplest form possible to avoid cluttering the patterns with implicit canonicalization.

Don't we want to make some of these folders/canonicalization patterns? That way, everyone doing rewriting benefits automatically, and does not need to explicitly pull in these new specific rewriting patterns?

Thanks! I will commit it tomorrow if no more comments.

Addressed feedback.

Just to clarify, "prune" is ignored for post order walks?

Addressed feedback.

• Addressed Sergei's feedback.
• Minor fix to 'iter_args' tests.
• Minor changes to registration APIs.

Thanks for working on this, Sergei! Some initial comments!

Restoring previous approved version. I'll commit it if no more comments.
I'll create a separate review for the prune changes.

Can you separate out the prune functionality into a different commit?

Kind ping :)

This new diff introduces important changes to the previous version.
Hopefully, @rriddle/@frgossen/@bondhugula can take another look.

Let me revisit a few things in this patch and come back.

Thanks, Aart! I just replied to some comments. Nothing else from my side!
Maybe it would be good that somebody else from your side took a look, as well.

(sorry, replace "first dimension" with "last dimension" in my prev. comment :))

Hey Aart!

LGTM! Just minor comments. I think Uday planned to commit it. Otherwise, I could commit it tomorrow together with one of mine.

Addressing feedback.

Does this patch handle iter_args in an inner loop? For example:

• Passing order as an extra parameter.
• Fixing typo.

Addressing feedback:

• Use enum to encode walk order.
• Set region, block and operation walks to post-order by default.
• Adjust walks in Liveness and NumberOfExecutions to align with the new default traversal.

Please also see: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/tensorflow/utils/visitor_util.h

It looks great! Thank you so much!

I think the add between affine.load and affine.for was handled by hasNonAffineUsersOnThePath.

Thanks for addressing the comments!

I would like to know what @bondhugula and @andydavis1 think but I think we might want to add the proper
%0 -> %i1 dependence to the MDG,

If you could improve the lit test so that loads and stores are also checked, that should be all. Thanks!

https://reviews.llvm.org/D97030 is somehow related. I think the common solution is to represent these missing dependences in the MDG. We may have to coordinate the fixes.

Hey tungld,

Thank you so much for the fix! Some comments below.

It works after the last update. Thanks!
LGTM. Some minor comments. Please, wait for @ftynse's final approval.

Maybe there's something off on my side but I can't see the diff between the last version of the patch (Diff 3) and Base. When I set up that view in "Revision Contents->History", it shows the Diff against a previous version of the patch, which doesn't include all the changes. Something seems to be off.

If this is a commitment that you are working on starting to connect the pieces next then fair enough :)