[MLIR][Shape] Canonicalize subsequent `size_to_index` and `index_to_size`

[frgossen]

Summary

Eliminate the subsequent applications of size_to_index and index_to_size.

Depends On D82082

[silvas]

Great! Thanks for these nicely factored patches!

[Kayjukh]

nit: You may want to update the title before landing it

[frgossen]

Remove trailing white line

[frgossen]

White line

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

def Shape_IndexToSizeOp : Shape_Op<"index_to_size", [NoSideEffect]> {
   }];
 
   let arguments = (ins Index:$arg);
   let results = (outs Shape_SizeType:$result);
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_JoinOp : Shape_Op<"join", []> {
   let summary = "Returns the least general shape.size of its operands";
   let description = [{
     An operation that computes the least general shape of input operands.
     This effectively asserts that corresponding static dimensions are equal.
     The behavior is to match each element of the `shape.shape` and propagate the

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

 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;
   return {};
 }
 
+void IndexToSizeOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<SizeToIndexToSizeCanonicalization>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // FromExtentsOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult FromExtentsOp::fold(ArrayRef<Attribute> operands) {
   if (llvm::any_of(operands, [](Attribute a) { return !a; }))
     return nullptr;
   SmallVector<int64_t, 6> extents;

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

 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
 
 def IndexToSizeToIndexCanonicalization : Pat<
   (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
   (replaceWithValue $arg)>;
 
+def SizeToIndexToSizeCanonicalization : Pat<
+  (Shape_IndexToSizeOp (Shape_SizeToIndexOp $arg)),
+  (replaceWithValue $arg)>;
+

mlir/test/Dialect/Shape/canonicalize.mlir

 // CHECK-SAME: (%[[IDX:.*]]: index) -> index
 func @index_to_size_to_index(%index : index) -> index {
   // CHECK: return %[[IDX]] : index
   %size = shape.index_to_size %index
   %result = shape.size_to_index %size
   return %result : index
 }
 
+// -----
+
+// Canonicalize redundant conversion from `size` to `index` and back.
+// CHECK-LABEL: @size_to_index_to_size
+// CHECK-SAME: (%[[SIZE:.*]]: !shape.size) -> !shape.size
+func @size_to_index_to_size(%size : !shape.size) -> !shape.size {
+  // CHECK: return %[[SIZE]] : !shape.size
+  %idx = shape.size_to_index %size
+  %result = shape.index_to_size %idx
+  return %result : !shape.size
+}
+