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

[frgossen]

Summary

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

[herhut]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

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.

[frgossen]

Address comments

[herhut]

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

[silvas]

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

[frgossen]

Rebase

[herhut]

In D81871#2096035, @silvas wrote:

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]

Address concerns

[frgossen]

Move patterns to dedicates pass

[silvas]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

this sounds like it should be in Conversion/ShapeToStandard

[frgossen]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

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]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

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.

[rriddle]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

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.

32 ↗

(On Diff #271371)

(here = whatever lowering would actually use this)

[herhut]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

28 ↗

(On Diff #271371)

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.

32 ↗

(On Diff #271371)

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,

37 ↗

(On Diff #271371)

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

[frgossen]

Rebase

[frgossen]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

28 ↗

(On Diff #271371)

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.
@pifon2a

37 ↗

(On Diff #271371)

This function is also related to the num_elements operation.
@pifon2a

[herhut]

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

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

28 ↗

(On Diff #271371)

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.

37 ↗

(On Diff #271371)

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

[silvas]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

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]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

32 ↗

(On Diff #271371)

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]

Rename pass

[jpienaar]

mlir/include/mlir/Dialect/Shape/Transforms/Passes.h

35 ↗

(On Diff #272506)

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.

32 ↗

(On Diff #271371)

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.

mlir/lib/Dialect/Shape/Transforms/CanonicalizeShapeToStandard.cpp

1 ↗

(On Diff #272506)

This no longer matches renamed file

[frgossen]

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)
See https://reviews.llvm.org/D82644

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_GetExtentOp : Shape_Op<"get_extent", [NoSideEffect]> {
   ];
 
   let extraClassDeclaration = [{
     /// Get the `dim` value as integer if it is constant.
     Optional<int64_t> getConstantDim();
   }];
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
   let summary = "Converts a standard index to a shape size";
   let description = [{
     Converts a standard index to a `shape.size`. This operation and its
     inverse, `size_to_index`, facilitate index conversion between the standard
     and the shape dialect.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

def DimOp : Std_Op<"dim", [NoSideEffect]> {
   let results = (outs Index:$result);
 
   let assemblyFormat = [{
     attr-dict $memrefOrTensor `,` $index `:` type($memrefOrTensor)
   }];
 
   let builders = [
     OpBuilder<"OpBuilder &builder, OperationState &result, "
-              "Value memrefOrTensor, int64_t index">
+              "Value memrefOrTensor, int64_t index">,
+    OpBuilder<"OpBuilder &builder, OperationState &result, "
+              "Value memrefOrTensor, Value index">,
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
   }];
 
   let hasFolder = 1;

mlir/lib/Dialect/Shape/IR/Shape.cpp

 //===- Shape.cpp - MLIR Shape Operations ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Shape/IR/Shape.h"
 
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Traits.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Support/raw_ostream.h"
 
OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   Builder builder(getContext());
   return builder.getIndexTensorAttr(extents);
 }
 
 //===----------------------------------------------------------------------===//
 // GetExtentOp
 //===----------------------------------------------------------------------===//
 
+void GetExtentOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<GetExtentShapeOfCanonicalization>(context);
+}
+
 Optional<int64_t> GetExtentOp::getConstantDim() {
   if (auto constSizeOp = dim().getDefiningOp<ConstSizeOp>()) {
     return constSizeOp.value().getLimitedValue();
   }
   return llvm::None;
 }
 
 OpFoldResult GetExtentOp::fold(ArrayRef<Attribute> operands) {

mlir/lib/Dialect/Shape/IR/ShapeCanonicalization.td

 include "mlir/Dialect/Shape/IR/ShapeOps.td"
+include "mlir/Dialect/StandardOps/IR/Ops.td"
 
 def EqualBinaryOperands : Constraint<CPred<"$0 == $1">>;
 
 def AllInputShapesEq : Constraint<CPred< [{
   llvm::all_of($0, [&](mlir::Value val) {
     return $0[0] == val;
   })
 }]>>;
 
 // Canonicalization patterns.
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $lhs, $rhs),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(EqualBinaryOperands $lhs, $rhs)]>;
 
 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
+
+def GetExtentShapeOfCanonicalization : Pat<
+  (Shape_GetExtentOp (Shape_ShapeOfOp $arg), $idx),
+  (Shape_IndexToSizeOp (DimOp $arg, (Shape_SizeToIndexOp $idx)))>;

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 
 void DimOp::build(OpBuilder &builder, OperationState &result,
                   Value memrefOrTensor, int64_t index) {
   auto loc = result.location;
   Value indexValue = builder.create<ConstantIndexOp>(loc, index);
+  build(builder, result, memrefOrTensor, indexValue);
+}
+
+void DimOp::build(OpBuilder &builder, OperationState &result,
+                  Value memrefOrTensor, Value index) {
   auto indexTy = builder.getIndexType();
-  build(builder, result, indexTy, memrefOrTensor, indexValue);
+  build(builder, result, indexTy, memrefOrTensor, index);
 }
 
 Optional<int64_t> DimOp::getConstantIndex() {
   if (auto constantOp = index().getDefiningOp<ConstantOp>())
     return constantOp.getValue().cast<IntegerAttr>().getInt();
   return {};
 }
 

mlir/test/Dialect/Shape/canonicalize.mlir

 func @f(%arg0 : !shape.shape) {
   // CHECK-NEXT: shape.const_witness true
   // CHECK-NEXT: consume.witness
   // CHECK-NEXT: return
   %0 = shape.cstr_broadcastable %arg0, %arg0
   "consume.witness"(%0) : (!shape.witness) -> ()
   return
 }
+
+// -----
+
+// Canonicalize `get_extent` not to rely on `shape_of`.
+// CHECK-LABEL: @get_extent
+// CHECK-SAME: (%[[ARG:.*]]: tensor<2x?xf32>, %[[IDX:.*]]: !shape.size) -> !shape.size
+func @get_extent(%arg : tensor<2x?xf32>, %idx : !shape.size) -> !shape.size {
+  // CHECK-NEXT: %[[IDX_AS_INDEX:.*]] = shape.size_to_index %[[IDX]]
+  // CHECK-NEXT: %[[RESULT_AS_INDEX:.*]] = dim %[[ARG]], %[[IDX_AS_INDEX]] : tensor<2x?xf32>
+  // CHECK-NEXT: %[[RESULT:.*]] = shape.index_to_size %[[RESULT_AS_INDEX]]
+  // CHECK-NEXT: return %[[RESULT]] : !shape.size
+  %shape = shape.shape_of %arg : tensor<2x?xf32>
+  %result = shape.get_extent %shape, %idx
+  return %result : !shape.size
+}