message: [MLIR][MHLO] Replace shape reification in shape simplification
  Shape reification was introduced as a fix after the integration of
  Instead, replace `extract(shape_of)` with
  `dim`.
diff: @@ -150,42 +150,22 @@ struct BroadcastRemoveSubsumedOperandsPattern
  }
  };
  -#define GEN_PASS_CLASSES
  -#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
  +struct ExtractFromExtentTensorCanonicalizationPattern
  + : public OpRewritePattern<tensor::ExtractOp> {
  + using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
  -// Resolve shape of result of an operation in terms of shape of its operands
  -// using the `InferShapedTypeOpInterface`. Most Linalg operations implement this
  -// interface.
  -struct ShapeOfOpWithReifyPerDim : public OpRewritePattern<shape::ShapeOfOp> {
  - using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
  -
  - LogicalResult matchAndRewrite(shape::ShapeOfOp op,
  + LogicalResult matchAndRewrite(tensor::ExtractOp op,
  PatternRewriter &rewriter) const override {
  - if (!op.arg().getType().isa<ShapedType>()) return failure();
  - auto arg = op.arg().dyn_cast<OpResult>();
  - if (!arg) return failure();
  - auto defining_op = dyn_cast<InferShapedTypeOpInterface>(arg.getOwner());
  - if (!defining_op) return failure();
  -
  - SmallVector<SmallVector<Value>> reified_result_shape_per_dim;
  - if (failed(defining_op.reifyReturnTypeShapesPerResultDim(
  - rewriter, reified_result_shape_per_dim)))
  - return failure();
  - Value reified_result_shape = rewriter.create<tensor::FromElementsOp>(
  - op.getLoc(), reified_result_shape_per_dim[arg.getResultNumber()]);
  - rewriter.replaceOp(op, reified_result_shape);
  + auto shape_of_op = op.tensor().getDefiningOp<ShapeOfOp>();
  + if (!shape_of_op) return failure();
  + Value index = op.indices().front();
  + rewriter.replaceOpWithNewOp<memref::DimOp>(op, shape_of_op.arg(), index);
  return success();
  }
  };
  -// Simplify the shape expressions involving results of the operation. This makes
  -// the operations that are still not DCE-ed because of their use in shape
  -// expression dead.
  -LogicalResult simplifyShapeExprs(FuncOp func, MLIRContext *context) {
  - RewritePatternSet patterns(context);
  - patterns.insert<ShapeOfOpWithReifyPerDim>(context);
  - return applyPatternsAndFoldGreedily(func, std::move(patterns));
  -}
  +#define GEN_PASS_CLASSES
  +#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
  struct ShapeSimplification
  : public ShapeSimplificationBase<ShapeSimplification> {
  @@ -206,13 +186,12 @@ struct ShapeSimplification
  op->getCanonicalizationPatterns(patterns, context);
  }
  - patterns.insert<BroadcastRemoveSubsumedOperandsPattern>(context);
  + patterns.insert<BroadcastRemoveSubsumedOperandsPattern,
  + ExtractFromExtentTensorCanonicalizationPattern>(context);
  auto func = getFunction();
  if (failed(applyPatternsAndFoldGreedily(func, std::move(patterns))))
  return signalPassFailure();
  -
  - if (failed(simplifyShapeExprs(func, context))) return signalPassFailure();
  }
  };
  