[MLIR][Shape] Lower `shape.get_extent` to `std.dim` when possible

[frgossen]

Summary

When the shape is derived from a tensor argument the shape extent can be derived
directly from that tensor with std.dim.
This lowering pattern circumvents the necessity to materialize the shape in
memory.

[herhut]

Neat!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

 // represented as indices.
 def IndexToSizeOpConversion : Pat<
     (Shape_IndexToSizeOp $arg),
     (replaceWithValue $arg)>;
 def SizeToIndexOpConversion : Pat<
     (Shape_SizeToIndexOp $arg),
     (replaceWithValue $arg)>;
 
+// Derive shape extent directly from shape origin if possible.
+// This circumvents the necessity to materialize the shape in memory.
+def GetExtentShapeOfConversion : Pat<
+    (Shape_GetExtentOp (Shape_ShapeOfOp $arg), $idx),
+    (Shape_IndexToSizeOp (DimOp $arg, (Shape_SizeToIndexOp $idx))),
+    [],
+    (addBenefit 10)>;

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 // CHECK-SAME: (%[[SHAPE:.*]]: tensor<?xindex>) -> index
 func @rank(%shape : !shape.shape) -> !shape.size {
   // CHECK-DAG: %[[C0:.*]] = constant 0 : index
   // CHECK-DAG: %[[RESULT:.*]] = dim %[[SHAPE]], %[[C0]]
   // CHECK-DAG: return %[[RESULT]] : index
   %rank = shape.rank %shape
   return %rank : !shape.size
 }
+
+// -----
+
+// Express `get_extent` as `std.dim` when it relies directly on the outcome of a
+// `shape_of` operation.
+// CHECK-LABEL: @get_extent_shape_of
+// CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
+func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : !shape.size)
+    -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
+  // CHECK: return %[[RESULT]] : index
+  %shape = shape.shape_of %arg : tensor<2x3xf32>
+  %result = shape.get_extent %shape, %idx
+  return %result : !shape.size
+}
+