[mlir][Shape] Add a pattern to turn extract from shape_of into tensor.dim

[bkramer]

Summary

If I remember correctly this wasn't done previously because dim used to
be in the memref dialect.

[jpienaar]

Do we need to limit this to where $arg is a Tensor?

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

19

Add comment ?

53

Is this taking the front because there should be only 1 index as shape of should result in 1D shape?

mlir/test/Dialect/Shape/canonicalize.mlir

1390

Nit: C0 ? it just makes it easier for me to read below (or make %c1 above)

[frgossen]

Thanks!

mlir/test/Dialect/Shape/canonicalize.mlir

1389

%c1?

[bkramer]

Address review comments

[bkramer]

In D111651#3058922, @jpienaar wrote:

Do we need to limit this to where $arg is a Tensor?

tensor.extract only works on tensors, so I think !shape.shape cannot happen.

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

LogicalResult matchAndRewrite(tensor::CastOp op,
     rewriter.replaceOpWithNewOp<ShapeOfOp>(op, ty, shapeOfOp.arg());
     return success();
   }
 };
 } // namespace
 
 void ShapeOfOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
                                             MLIRContext *context) {
-  patterns.add<ShapeOfCastExtentTensor, ShapeOfWithTensor>(context);
+  patterns.add<ShapeOfCastExtentTensor, ShapeOfWithTensor,
+               ExtractFromShapeOfExtentTensor>(context);
 }
 
 LogicalResult mlir::shape::ShapeOfOp::inferReturnTypes(
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<Type> &inferredReturnTypes) {
   if (operands[0].getType().isa<ValueShapeType>())
     inferredReturnTypes.assign({ShapeType::get(context)});

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

 def HasSingleElement : Constraint<CPred< [{
   $0.size() == 1
 }]>>;
 
 def HasStaticShape : Constraint<CPred< [{
   $0.getType().dyn_cast<ShapedType>().hasStaticShape()
 }]>>;
 
+// Helper that takes the first element of a range.
+def TakeFront : NativeCodeCall<"$0.front()">;
+
 // Canonicalization patterns.
 
 def AssumingAllOneOp : Pat<(Shape_AssumingAllOp $args),
                            (replaceWithValue $args),
                            [(HasSingleElement $args)]>;
 
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
def SizeToIndexToSizeCanonicalization : Pat<
   (Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
   (replaceWithValue $arg)>;
 
 // Fold tensor.cast(const_shape) to const_shape. This changes the type of
 // const_shape to the destination type of the cast.
 def TensorCastConstShape : Pat <
   (Tensor_CastOp:$res (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg),
   [(HasStaticShape $res)]>;
+
+// tensor.extract from shape_of -> tensor.dim. We can take the first index
+// because shape_of always returns a 1D tensor.
+def ExtractFromShapeOfExtentTensor : Pat<
+  (Tensor_ExtractOp (Shape_ShapeOfOp $arg), $indices),
+  (Tensor_DimOp $arg, (TakeFront $indices))>;

mlir/test/Dialect/Shape/canonicalize.mlir

func @concretize_broadcast_result_type(%arg0 : tensor<2xindex>,
     %arg1 : tensor<3xindex>) -> tensor<?xindex> {
   // CHECK: %[[CONCR:.*]] = shape.broadcast %[[ARG0]], %[[ARG1]] : tensor<2xindex>, tensor<3xindex> -> tensor<3xindex>
   // CHECK: %[[RES:.*]] = tensor.cast %[[CONCR]] : tensor<3xindex> to tensor<?xindex>
   // CHECK: return %[[RES]]
   %0 = shape.broadcast %arg0, %arg1 : tensor<2xindex>, tensor<3xindex>
       -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_shapeof
+// CHECK-SAME:    %[[ARG0:.*]]: tensor<?x?xf64>
+func @extract_shapeof(%arg0 : tensor<?x?xf64>) -> index {
+ %c1 = constant 1 : index
+// CHECK:        %[[C1:.*]] = constant 1
+ %shape = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>
+// CHECK:        %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C1]]
+ %result = tensor.extract %shape[%c1] : tensor<2xindex>
+// CHECK:        return %[[DIM]]
+ return %result : index
+}