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

[frgossen]

Summary

Eliminate the subsequent applications of index_to_size and size_to_index.

[silvas]

Thanks!

[frgossen]

Remove trailing white line

[frgossen]

White line

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

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   }];
 
   let arguments = (ins Shape_SizeType:$arg);
   let results = (outs Index:$result);
 
   let assemblyFormat = "$arg attr-dict";
 
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
 }
 
 def Shape_YieldOp : Shape_Op<"yield",
     [HasParent<"ReduceOp">,
      NoSideEffect,
      ReturnLike,
      Terminator]> {
   let summary = "Returns the value to parent op";

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

 OpFoldResult SizeToIndexOp::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 SizeToIndexOp::getCanonicalizationPatterns(
+    OwningRewritePatternList &patterns, MLIRContext *context) {
+  patterns.insert<IndexToSizeToIndexCanonicalization>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(YieldOp op) {
   auto *parentOp = op.getParentOp();
   auto results = parentOp->getResults();
   auto operands = op.getOperands();

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

 include "mlir/Dialect/Shape/IR/ShapeOps.td"
 
 def EqualBinaryOperands : Constraint<CPred<"$0 == $1">>;
 
 def AllInputShapesEq : Constraint<CPred< [{
   llvm::all_of($0, [&](mlir::Value val) {
     return $0[0] == val;
   })
 }]>>;
 
 // Canonicalization patterns.
+
 def CstrBroadcastableEqOps : Pat<(Shape_CstrBroadcastableOp:$op $lhs, $rhs),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(EqualBinaryOperands $lhs, $rhs)]>;
 
 def CstrEqEqOps : Pat<(Shape_CstrEqOp:$op $shapes),
   (Shape_ConstWitnessOp ConstBoolAttrTrue),
   [(AllInputShapesEq $shapes)]>;
+
+def IndexToSizeToIndexCanonicalization : Pat<
+  (Shape_SizeToIndexOp (Shape_IndexToSizeOp $arg)),
+  (replaceWithValue $arg)>;
+

mlir/test/Dialect/Shape/canonicalize.mlir

 func @dont_canonicalize_rank(%arg : tensor<*xf32>) -> !shape.size {
 // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<*xf32>
 // CHECK-DAG: %[[SIZE:.*]] = shape.rank %[[SHAPE]]
 // CHECK-DAG: return %[[SIZE]] : !shape.size
 %shape = shape.shape_of %arg : tensor<*xf32>
 %rank = shape.rank %shape
 return %rank : !shape.size
 }
+
+// Canonicalize redundant conversion from `index` to `size` and back.
+// CHECK-LABEL: @index_to_size_to_index
+// 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
+}
+