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

[frgossen]

Summary

Simplify pattern of successive shape.get_extent and shape.from_extent_tensor
operations and express them with the standard operation extract_element.
This transformation later circumvents the necessity to materialize shapes in
memory.

Depends On D81871

[pifon2a]

mlir/lib/Dialect/Shape/Transforms/SimplifyShapePatterns.td

4

is this really needed?

[frgossen]

mlir/lib/Dialect/Shape/Transforms/SimplifyShapePatterns.td

4

Unfortunately, yes.
I tried to get away with out this and it seems that the declarative rewrite patterns can deal with variadic arguments but not with ValueRange.
I'd be happy if someone knows a nice solution.

[frgossen]

Inline NativeCodeCall

[frgossen]

Rebase

[frgossen]

https://reviews.llvm.org/D82645

mlir/lib/Dialect/Shape/Transforms/SimplifyShapePatterns.td

 include "mlir/Dialect/Shape/IR/ShapeOps.td"
 include "mlir/Dialect/StandardOps/IR/Ops.td"
 
+def createSingletonValueRange : NativeCodeCall<"ValueRange({$0})">;
+
+// Rewrite patterns.
+
 def GetExtentShapeOfConversion : Pat<
     (Shape_GetExtentOp (Shape_ShapeOfOp $arg), $idx),
     (Shape_IndexToSizeOp (DimOp $arg, (Shape_SizeToIndexOp $idx)))>;
+
+def GetExtentFromExtentTensorConversion : Pattern<
+    (Shape_GetExtentOp (Shape_FromExtentTensorOp $extents), $idx),
+    [
+      (Shape_SizeToIndexOp:$std_idx $idx),
+      (ExtractElementOp:$std_result $extents, (createSingletonValueRange $std_idx)),
+      (Shape_IndexToSizeOp $std_result)
+    ]>;

mlir/test/Dialect/Shape/simplify-shape.mlir

 // RUN: mlir-opt --simplify-shape --split-input-file %s | FileCheck %s
 
 // Express `get_extent` as `std.dim` when it relies directly on the outcome of a
 // `shape_of` operation.
 // CHECK-LABEL: @get_extent
 // CHECK-SAME: (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: !shape.size) -> !shape.size
 func @get_extent(%arg : tensor<2x3xf32>, %idx : !shape.size) -> !shape.size {
   // CHECK-DAG: %[[STD_IDX:.*]] = shape.size_to_index %[[IDX]]
   // CHECK-DAG: %[[STD_RESULT:.*]] = dim %[[ARG]], %[[STD_IDX]] : tensor<2x3xf32>
   // CHECK-DAG: %[[RESULT:.*]] = shape.index_to_size %[[STD_RESULT]]
   // CHECK-DAG: return %[[RESULT]] : !shape.size
   %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
+// CHECK-SAME: (%[[EXTENTS:.*]]: tensor<?xindex>, %[[IDX:.*]]: !shape.size) -> !shape.size
+func @get_extent(%extents : tensor<?xindex>, %idx : !shape.size) -> !shape.size {
+  // CHECK-DAG: %[[STD_IDX:.*]] = shape.size_to_index %[[IDX]]
+  // CHECK-DAG: %[[STD_RESULT:.*]] = extract_element %[[EXTENTS]][%[[STD_IDX]]] : tensor<?xindex>
+  // CHECK-DAG: %[[RESULT:.*]] = shape.index_to_size %[[STD_RESULT]]
+  // CHECK-DAG: return %[[RESULT]] : !shape.size
+  %shape = shape.from_extent_tensor %extents: tensor<?xindex>
+  %result = shape.get_extent %shape, %idx
+  return %result : !shape.size
+}