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
|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.
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))
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.
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).
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.
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 https://mlir.llvm.org/docs/PassManagement/#pass-pipeline-registration 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.
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)
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,
The ShapeToShapeLoweringPass is a different one and it lowers num_elements to reduce.
This function is also related to the num_elements operation.
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?
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.
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.
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