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

[Gil Rapaport]

aniragil wrote:

> Ioana is OOO for a few days, so let me quickly reply to keep to maintain the momentum.
> 
> > * This actually seems more like a generic func -> func optimization/canonicalization than an emitc-specific lowering.
> 
> That's a good point. In fact, the initial implementation tackled this at the Tensor level: #194404

That PR seems open and under review, why is was it abandoned?

> 
> However, if we are to consider both function definitions and call-sites, then we need some global (module level) rewrite, e.g. a conversion pass and that makes `FuncToEmitC` a perfect candidate! A regular rewrite pattern doesn't seem like the right tool and we couldn't find a suitable alternative (or any prior art).

Wouldn't a module pass doing that transformation suffice?

> 
> Also, there is a question of how general should this be? Ultimately, we are tackling the mismatch between Tensor and C semantics (hence many C-specific assumptions, e.g. that functions cannot return arrays) and for that reason it feels like EmitC conversion is the right level. As in, why would the Tensor dialect be concerned about functions returning tensors? (*)

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.

> 
> Note that this PR makes quite a few assumptions, e.g. "the function has exactly one result" + "the result is a (...) memref with exactly one element". That's sufficient for our needs today and it also allowed Ioana to keep this PR relatively simple.

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

> For more generic cases, we probably should invest into your proposal to switch to memref descriptors when lowering to EmitC.

Right, the memref-to-array lowering causes all sorts of issues and limitations. The memref-descriptor proposal should hopefully handle this completely and uniformly (as in the LLVM dialect), preferably before these issue cause this kind of workarounds which would be hard to deprecate later.

> Btw, there used to be a "de-tensorize" pass, but that was removed in #177579. We could re-visit that at some point, but that transformation focused on Linalg Ops rather than e.g. functions.

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.

> 
> > * What happens if the returned memref is stored to or deallocated by the caller?
> 
> That will be a problem if there are users, yes. This patch bails out in such cases. If we generalise to also look at users, then we would:
> 
>     * erase `memref.dealloc`,
> 
>     * rewrite `memref.store` as `emit.assign`?
> 
> 
> That should work, right?

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?

> 
> > * Why not lower `memref<0>`'s to a pointer?
> 
> IMHO it's a convenient convention to treat `memref<1xi64>` + `memref<f64>` as plain scalars, but admittedly it's just a convention. We could treat those as pointers, but would that gain us anything?

So in the motivation paragraph it states "EmitC cannot return arrays by value" as the problem. For zero-ranked memrefs pointers seem to be a straightforward and robust way to lower such functions.
I agree doing sort of a "Mem2Reg" is preferable when provably possible (which isn't always simple), but it sounds more like a canonicalization than a lowering.

> 
> To be perfectly honest, I feel that we are entering uncharted waters here and "we" (EmitC) may end-up having to iterate a bit. I am under the impression that MemRef descriptors will make lowering to EmitC more natural. In the meantime, we try to adapt to every case that we hit. But "descriptors" will likely take a while and perhaps first (independantly of this PR) we should jump on a call to discuss the plans in more detail. From MLGO perspective, I don't see any other requirements just now.

Agreed. I started laying the groundwork for memref descriptors as structs with with #198198 (another PR soon to follow), then I hope I can upload a patch and an RFC we can use as a basis for more detailed planning

> 
> (*) Just to be clear - I am trying to say to it's not 100% clear to me what the right place for this transformation is. Right now I am leaning towards EmitC conversions, but if you feel differently, lets continue discussing.

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?

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