Page MenuHomePhabricator

[MLIR][Shape] Simplify `get_extent` when dependent on `shape_of`

Authored by frgossen on Mon, Jun 15, 12:56 PM.



Simplify common pattern of shape.get_extent depending on shape_of.
The successive operations can be replaces with the standard dim operation.
This rewrite later avoids unnecessary shape materialization in memory.

Depends On D81834

Diff Detail

Event Timeline

frgossen created this revision.Mon, Jun 15, 12:56 PM
Herald added a project: Restricted Project. · View Herald TranscriptMon, Jun 15, 12:56 PM
herhut requested changes to this revision.Tue, Jun 16, 12:50 AM
herhut added inline comments.
60 ↗(On Diff #270835)

This can fail if get_extent is not called on a shape_of operation. So you need to use dyn_cast and fail the rewrite in such cases.

Also, this is not really a lowering but rather a canonicalization. This rewrite can happen earlier, as well.

88 ↗(On Diff #270835)

This is dangerous, as there might be uses. If the above pattern is replaced by a canonicalization, then the shape_of operation would be removed because it has no more users.

This revision now requires changes to proceed.Tue, Jun 16, 12:50 AM
frgossen updated this revision to Diff 271022.Tue, Jun 16, 3:44 AM
frgossen marked 2 inline comments as done.

Address comments

herhut accepted this revision.Tue, Jun 16, 4:22 AM

Thanks. Looks good to me, please also wait for @jpienaar to review.

silvas requested changes to this revision.Tue, Jun 16, 9:41 AM
silvas added a subscriber: stellaraccident.

As per discussion today, this pattern is going to "fight" with what we are doing with IREE (and also npcomp). Can we reserve this pattern for the shape to std-with-descriptors lowering?

Does this pattern actually have much value outside lowering? The dim op + casts doesn't seem much more analyzable than get_extent(shape_of(x))

+ @stellaraccident

This revision now requires changes to proceed.Tue, Jun 16, 9:41 AM

As per discussion today, this pattern is going to "fight" with what we are doing with IREE (and also npcomp). Can we reserve this pattern for the shape to std-with-descriptors lowering?

I don't think it fights with it. You will get dim operations in the code anyway that you have to resolve during lowering for IREE. dim(%val) is the canonical way in mlir to get the dimension of a tensor and is understood by most analyses, while get_extend(shape_of(...)) is not.

In any case, I am also ok with this to be a lowering. It just felt more like a canonicalization to me.

frgossen updated this revision to Diff 271299.Wed, Jun 17, 1:52 AM

Address concerns

frgossen updated this revision to Diff 271371.Wed, Jun 17, 7:39 AM

Move patterns to dedicates pass

frgossen retitled this revision from [MLIR][Shape] Lower `get_extent` and `shape_of` to [MLIR][Shape] Simplify `get_extent` when dependent on `shape_of`.Wed, Jun 17, 8:02 AM
frgossen edited the summary of this revision. (Show Details)
silvas requested changes to this revision.Wed, Jun 17, 11:52 AM
silvas added inline comments.

this sounds like it should be in Conversion/ShapeToStandard

This revision now requires changes to proceed.Wed, Jun 17, 11:52 AM
frgossen marked an inline comment as done.Wed, Jun 17, 2:49 PM
frgossen added inline comments.

That is what I tried first but it turns out that these patterns cannot be part of the lowering pass ShapeToStandard (at least I did not find in a nice way).

The lowering pass lowers shape_of before it reaches the get_extent operation (it does not actually lower shape_of yet but soon will).
At the time the root of this CL's pattern (get_extent) is reached there is no shape_of operation left and get_extent(shape_of(arg)) will never be encountered.
The same problem appears with get_extent(from_extent_tensor(arg)).

I thought of the benefits parameter but that allows (if I understand it correctly) to define preceedence only among applicable patterns per root operation.

This is why, I think, we need a separate pass for this partial lowering of the shape dialect.
Similar to the original idea of canonicalisation, this felt like the right place for the pass but I will happily move it elsewhere if there is a nicer place.
I'm also happy to put this back into the lowering pass if someone can explain how to solve my problem :-)

@silvas, @herhut

silvas added inline comments.Thu, Jun 18, 2:48 PM

It is fine for Conversion/ShapeToStandard to expose a multiple passes that together perform the complete conversion process. This pass can be one of the passes there.

A pass pipeline can be exposed which runs all the needed passes in the right order. (see for more info about describing pass pipelines)

Since this pass creates "dim" ops, it does not make sense for it to live inside the shape dialect directory.

rriddle added inline comments.Thu, Jun 18, 2:55 PM

That isn't always the case. If the pattern allows for progressive lowering it can still be in here. I don't see a reason to arbitrarily split the pass pipeline. We do not enforce that *all* patterns have to go directly to the target. We highly encourage adding additional patterns that go through different dialects/operations if a lowering already exists. Not to say that is the case for this pattern, but just a drive by comment.

Also, having a pass that only performs one small lowering seems sketchy to me. Generally we just expose patterns that "real" lowering passes pull in.


(here = whatever lowering would actually use this)

herhut added inline comments.Fri, Jun 19, 1:47 AM

Give this a name that reflects the purpose more. createSimplifyShapeToStandardPass or something? These are in the global mlir namespace, so they need to be somewhat unique.


This should have a name that describes the patterns purpose more.


If we make the use of dim ops a criterion of whether something belongs into one dialect or another, we have to split things way more. One needs standard ops when lowering into SCF for example.

Originally, these patterns (and there will be more over time) were canonicalizations but we agreed that they might not be useful in all scenarios (e.g. when not using descriptors). We could call this SimplifyShapeToDescriptors but then again at this level it is not descriptors, yet.

This pass also does not lower all the way to the standard dialect and these patterns can be run early as a "custom canonicalization" as they still keep shape computations on the shape dialect where they compute new information. Only if the shape computation creates information that is already in the descriptors, we rewrite them here.

So, in conclusion, I would like to keep them in the shape dialect and separate from the ShapeToStandard lowering. We can call them CanonicalizeShapeToStandard maybe? But we should not block this on name bikeshedding too much,

frgossen updated this revision to Diff 272072.Fri, Jun 19, 7:40 AM
frgossen marked 2 inline comments as done.


frgossen added inline comments.Fri, Jun 19, 7:43 AM

The ShapeToShapeLoweringPass is a different one and it lowers num_elements to reduce.
I guess we don't want to merge these two passes because the ShapeToShapeLowering is a pass that we all need - descriptor and non-descriptor world.


This function is also related to the num_elements operation.

herhut accepted this revision.Fri, Jun 19, 8:49 AM

Just names, but we can bikeshed on names some more later.


Wrong line. There are too many comments on this change :(

I wound find createSimplifyShapeToStandardPass better than just createSimplifyShapes is what I wanted to say. That name isn't great either. I though it would be good to somehow convey in the name that this is useful when going to descriptors.


Can you add one for the simplify patterns, as well?

silvas added inline comments.Fri, Jun 19, 12:27 PM

CanonicalizeShapeToStandard sounds fine to me. I still claim that this pass fundamentally is a conversion pass (it fundamentally replaces the "looking at the shape of a tensor" behavior from the one provided by the shape dialect (shape.shape_of) to a different mechanism external to the shape dialect (std.dim)), but I don't care that much.

silvas added inline comments.Fri, Jun 19, 12:48 PM

Also, I'll note that our lowering process here is using mixed dialects. E.g. ShapeToSCF actually introduces extract_element and dim from std and bakes in a certain lowering for !shape.shape to tensors of index which is itself a type from std.

Really all of these are just pieces in "convert !shape.shape to tensor<?xindex>" conversion. Strictly speaking, it would probably make more sense for most of the recent "ShapeToX" work to live in lib/Conversion/ShapeDotShapeToTensorOfIndex, but that's incredibly verbose and breaks from our established "dialect to dialect" way of describing conversions. We should consider it though.

I'll contrast this pass to e.g. the lowering of shape.num_elements to shape.reduce, which is a pure in-dialect conversion that is independently useful no matter what !shape.shape is ultimately being converted to.

frgossen updated this revision to Diff 272506.Mon, Jun 22, 11:39 AM
frgossen marked 10 inline comments as done.

Rename pass

frgossen marked an inline comment as done.Mon, Jun 22, 11:39 AM
silvas accepted this revision.Mon, Jun 22, 11:40 AM
jpienaar added inline comments.Tue, Jun 23, 2:33 PM

We don't have a hard constraint there: we definitely mix dialects during progressive lowering (and already use standard dialect along with others).

This does really look like a pattern that is part of lowering to standard to me too. And I would have expected the pattern to match given the bottom up application of greedy pattern rewriter.


I don't follow canonicalizing between between dialects here. Either it is a canonicalization (and then the dialects don't matter and I don't see why it isn't a canonicalization on the op) or it is part of some legalization.

It doesn't seem safe as a canonicalization though as you have a std.dim in between which doesn't have the same error propagating behavior.


This no longer matches renamed file

frgossen abandoned this revision.Fri, Jun 26, 5:09 AM

It's slightly embarrassing but it seems to work now.
(Don't know what I did wrong last time. I was sure I had tried the benefits parameter)