transformations: (stencil-tensorize-z-dim) Support loops

[n-io]

Add support for enclosing unknown ups such as scf.for

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 89.95%. Comparing base (37702db) to head (8bd0fa7).
Report is 20 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3118      +/-   ##
==========================================
+ Coverage   89.90%   89.95%   +0.04%     
==========================================
  Files         421      425       +4     
  Lines       53306    53568     +262     
  Branches     8263     8297      +34     
==========================================
+ Hits        47927    48186     +259     
- Misses       4041     4043       +2     
- Partials     1338     1339       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dk949]

Could you please add a description as to what this PR aims to achieve?

[dk949]

I don't fully understand what this change does. Is it to relax the restrictions on what can be tensorized?

[n-io]

I don't fully understand what this change does. Is it to relax the restrictions on what can be tensorized?

The type conversion pass helps dealing with unknown ops. Notably, including scf.for. Running the type conversion pass first has the effect of the old checks "input is tensorised and output is not tensorised" no longer working. So the ExternalLoadOp conversion was not needed at all, and some other rewrite passes needed their checks (esp for output) revisited.

[dk949]

I don't fully understand what this change does. Is it to relax the restrictions on what can be tensorized?

The type conversion pass helps dealing with unknown ops. Notably, including scf.for. Running the type conversion pass first has the effect of the old checks "input is tensorised and output is not tensorised" no longer working. So the ExternalLoadOp conversion was not needed at all, and some other rewrite passes needed their checks (esp for output) revisited.

Makes sense, thank you