[Mlir-commits] [mlir] [EmitC] Scalarize single-element memref returns (PR #198604)

[Andrzej Warzyński]

banach-space wrote:

Thanks for the feedback, Gil! 

> That PR seems open and under review, why is was it abandoned?
 
TBH, we got busy trying different approaches. We also wanted to keep it open for reference - in case you suggested that as an option :) 

If you are happy with the path suggested here, we can close that one. In the meantime, I suggest minimising the noise.  

> Wouldn't a module pass doing that transformation suffice?

Yes, @ioghiban will sent a separate PR with a module pass shortly. (*)

(*) IMHO a separate PR is "cleaner" than completely rewriting an existing PR, e.g. this one. I hope that that's OK?

> Tensors, being pure SSA values, make such transformations much simpler than memrefs with their load-store/aliasing semantics. Since at some point each tensor should either be bufferized into a memref or casted to a scalar, it seems easier to avoid bufferization than later proving they can be removed.

These are very good points for implementing this at the Tensor level, yes. 

In fact, we should actually either ask the authors of those models or patch the TFLite importer to avoid those 1-element tensors to begin with - that's my long-term goal 😅 

In the meantime, tweaking this at EmitC/MemRef level feels like the most straightforward path for us. That's modulo all the issues that you've mentioned and which we are trying to avoid by bailing out conservatively.

> The patch indeed aims to be handle only specific cases, but not sure I'd call it simple :)

Fair :)

> That actually supports the direction of this canonicalization being done earlier. IIUC from the Discourse discussion the pass wasn't "wrong" in any way, just not owned/developed by anyone.

We couldn't even find any users of that code :( 

It does feel like a potential Tensor canonicalization. Even once we gain memref-descriptors in EmitC, fixing this at the Tensor level would be "neater".

My personal preference would be to start with the approach taken here and then possibly generalise to Tensor. Basically, gain experience by applying this within EmitC and then use that learning to generalise. I know that this makes EmitC a guinea pig, but it allows us to approach this incrementally  and discuss between ourselves (i.e. EmitC) before expanding.

> Hmmm ... so the code (and commit message) actually state the transformation will only be done on private and unused functions. But aren't such functions dead and can just be deleted?

Good point. But then how do we "label" code that we know won't be used in any other MLIR module (so MLIR transformations are safe), but can be used from C++ (i.e. outside of MLIR)? 

That's roughly the situation that we have here and feels like a gap between bridging MLIR and C 🤔 Should we look into creating another attribute?

> For zero-ranked memrefs pointers seem to be a straightforward and robust way to lower such functions.

You might be right (I don't want to be definitive without trying), but then we will be special-casing for `memref<f32>`. That's fine with me, but some people dislike that. I also prefer to treat everything consistently and in general view `f32`, `memref<f32>` and `memref<1xf32>` as identical types (i.e. scalars). 

> So perhaps a temporary solution could be the "MLGO-specific passes" mechanism you and @mtrofin suggested: Instead of converting zero-rank memrefs through the common lowering passes while we're unsure if and where such a transformation belongs, could this transformation be implemented as a module pass working on `func.func` and run by MLGO before lowering to emitc?

Sounds like a plan, thank you for the suggestion! Like I mentioned earlier, Ioana will be sending a PR soonish.

https://github.com/llvm/llvm-project/pull/198604