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

[frgossen]

Summary

When the origin of a shape is an extent tensor the operation get_extent can be
lowered directly to extract_element.
This choice circumvents the necessity to materialize the shape in memory.

Depends On D82644

[herhut]

Nice!

mlir/lib/Conversion/ShapeToStandard/ShapeToStandardPatterns.td

 
 // 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)>;
+def GetExtentFromExtentTensorConversion : Pattern<
+    (Shape_GetExtentOp (Shape_FromExtentTensorOp $extents), $idx),
+    [
+      (Shape_SizeToIndexOp:$std_idx $idx),
+      (ExtractElementOp:$std_result $extents, (NativeCodeCall<"ValueRange({$0})"> $std_idx)),
+      (Shape_IndexToSizeOp $std_result)
+    ],
+    [],
+    (addBenefit 10)>;
+

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

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
 }
 
+// -----
+
+// Express `get_extent` as `std.extract_element` when it relies directly on the
+// outcome of a `from_extent_tensor` operation.
+// CHECK-LABEL: @get_extent_from_extent_tensor
+// CHECK-SAME: (%[[EXTENTS:.*]]: tensor<?xindex>, %[[IDX:.*]]: index) -> index
+func @get_extent_from_extent_tensor(%extents : tensor<?xindex>,
+                                    %idx : !shape.size) -> !shape.size {
+  // CHECK: %[[RESULT:.*]] = extract_element %[[EXTENTS]][%[[IDX]]] : tensor<?xindex>
+  // CHECK: return %[[RESULT]] : index
+  %shape = shape.from_extent_tensor %extents : tensor<?xindex>
+  %result = shape.get_extent %shape, %idx
+  return %result : !shape.size
+}
+