diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td
@@ -1087,18 +1087,20 @@
     >,
     InterfaceMethod<
       /*desc=*/[{
-        Return the position in the results of the affine map computed
-        by getLoopsToShapesMap() that represents the shape of the
-        result value at a dimension.
+        Return the range of position in the result of the affine map
+        computed by getLoopsToShapesMap() which correspond to the
+        AffineExprs used to access the outputs of the operation.
       }],
-      /*retTy=*/"Optional<unsigned>",
-      /*methodName=*/"getResultValueDimPositionInLoopsToShapeMap",
-      /*args=*/(ins "unsigned":$resultIdx, "unsigned":$dim),
+      /*retTy=*/"std::pair<unsigned, unsigned>",
+      /*methodName=*/"getResultsPositionInLoopsToShapeMap",
+      /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
-        if (resultIdx >= getNumOutputs()) return {};
-        return getOperandDimPositionInLoopsToShapeMap(
-            getNumInputs() + resultIdx, dim);
+        return
+          {*getOperandDimPositionInLoopsToShapeMap(getNumInputs(), 0),
+           (*getOperandDimPositionInLoopsToShapeMap
+                 (getNumInputs() + getNumOutputs() - 1,
+                  getOutputShapedType(getNumOutputs()-1).getRank() - 1)) + 1};
       }]
     >,
 
@@ -1188,8 +1190,8 @@
 
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
-    Optional<Value> inferResultDimFromInputShapes
-      (OpBuilder &b, Location loc, unsigned resultIdx, unsigned im);
+    LogicalResult reifyReturnTypeShapes(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes);
 
     //========================================================================//
     // Helper functions to mutate the `operand_segment_sizes` attribute.
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h
@@ -22,6 +22,7 @@
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/CopyOpInterface.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "mlir/Support/LLVM.h"
@@ -107,13 +108,6 @@
 void getDimsOfType(Operation *op, StringRef iteratorTypeName,
                    SmallVectorImpl<AffineExpr> &res);
 
-/// For reshape operation, compute the shape of the output based on the result
-/// type and shape of the input.
-SmallVector<Value, 4>
-getReshapeOutputShapeFromInputShape(OpBuilder &b, Location loc, Value src,
-                                    ArrayRef<int64_t> dstStaticShape,
-                                    ArrayRef<AffineMap> reassociation);
-
 namespace detail {
 LogicalResult verifyStructuredOpInterface(Operation *op);
 } // namespace detail
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td
@@ -15,6 +15,7 @@
 
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ViewLikeInterface.td"
@@ -33,7 +34,10 @@
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
-def Linalg_InitTensorOp : Linalg_Op<"init_tensor", [NoSideEffect]> {
+def Linalg_InitTensorOp : Linalg_Op<"init_tensor",
+    [NoSideEffect,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapes"]>]> {
   let summary = "operation to define a tensor of particular value";
 
   let description = [{
@@ -126,7 +130,10 @@
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments, NoSideEffect]> {
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapes"]>,
+     NoSideEffect]> {
   let summary = "tensor pad operation";
   let description = [{
     `linalg.pad_tensor` is an operation that pads the `source` tensor
@@ -348,11 +355,6 @@
               a.cast<AffineMapAttr>().getValue().getResults());
           }));
     }
-    SmallVector<Value, 4> getOutputShape(OpBuilder &b, Location loc) {
-      return getReshapeOutputShapeFromInputShape(
-          b, loc, src(), getResultType().getShape(),
-          getReassociationMaps());
-    }
   }];
   let assemblyFormat = [{
     $src $reassociation attr-dict `:` type($src) `into` type(results)
@@ -417,7 +419,10 @@
   let hasCanonicalizer = 1;
 }
 
-def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
+def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<
+    "tensor_reshape",
+    [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+      ["reifyReturnTypeShapes"]>]>,
     Arguments<(ins AnyTensor:$src,
                    AffineMapArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
--- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td
@@ -17,6 +17,7 @@
 include "mlir/Dialect/Linalg/IR/LinalgBase.td"
 include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 // Base Tablegen class for Linalg ops.
@@ -25,7 +26,7 @@
 // depending on the specific Linalg op.
 class LinalgStructuredBase_Op<string mnemonic, list<OpTrait> props>
   : Op<Linalg_Dialect, mnemonic, !listconcat(props, [
-       LinalgStructuredInterface])> {
+       LinalgStructuredInterface, InferShapedTypeOpInterface])> {
   code structuredOpsBaseDecls = [{
     // Return the number of induction variables in the basic block. This should
     // always be 0 for index-free linalg ops. For IndexedGeneric, this must be
@@ -33,6 +34,12 @@
     unsigned getNumPayloadInductionVariables() {
       return isa<IndexedGenericOp>(this->getOperation()) ? getNumLoops() : 0;
     }
+
+    LogicalResult reifyReturnTypeShapes(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+      return cast<LinalgOp>(getOperation()).reifyReturnTypeShapes(b,
+          reifiedReturnShapes);
+    }
   }];
 }
 
diff --git a/mlir/include/mlir/Interfaces/InferTypeOpInterface.h b/mlir/include/mlir/Interfaces/InferTypeOpInterface.h
--- a/mlir/include/mlir/Interfaces/InferTypeOpInterface.h
+++ b/mlir/include/mlir/Interfaces/InferTypeOpInterface.h
@@ -16,6 +16,7 @@
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
diff --git a/mlir/include/mlir/Interfaces/InferTypeOpInterface.td b/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
--- a/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
+++ b/mlir/include/mlir/Interfaces/InferTypeOpInterface.td
@@ -97,21 +97,28 @@
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::mlir::SmallVectorImpl<::mlir::ShapedTypeComponents>&":
-                      $inferredReturnShapes)
+                      $inferredReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.
 
-      Insert operations using the given OpBuilder that computes the result
-      shape.
+      Insert operations using the given OpBuilder that computes the
+      result shape. The `reifiedReturnShapes` is expected to be
+      populated with as many vectors as the number of results of the
+      op (empty if the shape of a result value cannot be computed). If
+      the returned shape for a result is not empty, its size must
+      match the rank of the shaped type returned.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
-                    "::mlir::SmallVectorImpl<::mlir::Value>&":$reifiedReturnShapes),
+          "::mlir::SmallVectorImpl<SmallVector<::mlir::Value>>&"
+          :$reifiedReturnShapes),
       /*methodBody=*/[{}],
       /*defaultImplementation=*/[{ return ::mlir::failure(); }]
-    >,
+    >
   ];
 }
 
diff --git a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
--- a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt
@@ -14,6 +14,7 @@
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRDialectUtils
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
--- a/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp
@@ -233,57 +233,57 @@
   llvm::SmallSet<unsigned, 4> positions;
 };
 
-Optional<Value> LinalgOp::inferResultDimFromInputShapes(OpBuilder &b,
-                                                        Location loc,
-                                                        unsigned resultIdx,
-                                                        unsigned dim) {
+LogicalResult LinalgOp::reifyReturnTypeShapes(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
   // An example that helps understand the logic below.
   // Consider the following expression O(i+j, j) += A(i,k) * B(k, j)
   // We want to express the shape of dim 0 of O in terms of shape of the inputs.
   // This is achieved as follows.
   //   loopsToShapesMap = (d0, d1, d2) -> (d0, d2, d2, d1, d0 + d1, d1)
-  //   subMapOfResultDim = (d0, d1, d2) -> (d0 + d1)
+  //   subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)
   //   shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
-  //   resultFromFromInputDim = subMapOfResultDim.compose(shapesToLoopMap)
-  //     = (d0, d1, d2, d3, d4, d5) -> (d0 + d1)
+  //   resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)
+  //     = (d0, d1, d2, d3, d4, d5) -> (d0 + d1, d1)
   AffineMap loopsToShapesMap = getLoopsToShapesMap();
 
   // Find the position in the above map that represents the shape of the
   // result:dim being inferred.
-  Optional<unsigned> resultDimSubMapPos =
-      getResultValueDimPositionInLoopsToShapeMap(resultIdx, dim);
-  if (!resultDimSubMapPos)
-    return {};
+  auto resultShapesSubMapPos = getResultsPositionInLoopsToShapeMap();
 
   /// From loopsToShapesMap extract the submap that represents the shape of the
-  /// (resultIdx, dim) needed
-  AffineMap loopToResultDimShapeMap =
-      loopsToShapesMap.getSubMap(*resultDimSubMapPos);
-  AffineMap operandShapesToResultDimMap =
-      loopToResultDimShapeMap.compose(getShapesToLoopsMap());
+  /// (resultIdx, dim) needed.
+  SmallVector<unsigned, 4> resultPosRange =
+      llvm::to_vector<4>(llvm::seq<unsigned>(resultShapesSubMapPos.first,
+                                             resultShapesSubMapPos.second));
+  AffineMap loopToResultsShapeMap = loopsToShapesMap.getSubMap(resultPosRange);
+  AffineMap resultShapesFromInputShapesMap =
+      loopToResultsShapeMap.compose(getShapesToLoopsMap());
 
   // Check that the result dim map does not contain the positions corresponding
   // to the outputs.
   llvm::SmallSet<unsigned, 4> outputDims;
-  unsigned outputDimPosStart =
-      getResultValueDimPositionInLoopsToShapeMap(0, 0).getValue();
-  unsigned outputDimPosEnd =
-      getResultValueDimPositionInLoopsToShapeMap(getNumOutputs() - 1,
-                                                 getOutputOpOperands()
-                                                         .back()
-                                                         .get()
-                                                         .getType()
-                                                         .cast<ShapedType>()
-                                                         .getRank() -
-                                                     1)
-          .getValue();
-  llvm::for_each(llvm::seq<unsigned>(outputDimPosStart, outputDimPosEnd),
+  llvm::for_each(resultPosRange,
                  [&outputDims](unsigned dim) { outputDims.insert(dim); });
   HasAffineDimExprVisitor checkDimExpr(outputDims);
-  if (checkDimExpr.visit(operandShapesToResultDimMap.getResult(0)))
-    return llvm::None;
-  return applyMapToValues(b, loc, operandShapesToResultDimMap,
-                          createFlatListOfOperandDims(b, loc))[0];
+  Location loc = getOperation()->getLoc();
+  auto allResultDimValues =
+      applyMapToValues(b, loc, resultShapesFromInputShapesMap,
+                       createFlatListOfOperandDims(b, loc));
+  unsigned pos = 0;
+  ArrayRef<AffineExpr> shapeExprs = resultShapesFromInputShapesMap.getResults();
+  for (auto resultIdx : llvm::seq<unsigned>(0, getNumOutputs())) {
+    ShapedType resultType = getOutputShapedType(resultIdx);
+    SmallVector<Value> shapes;
+    for (unsigned dim : llvm::seq<unsigned>(0, resultType.getRank())) {
+      if (checkDimExpr.visit(shapeExprs[pos]))
+        shapes.push_back(b.create<DimOp>(loc, getOutput(resultIdx), dim));
+      else
+        shapes.push_back(allResultDimValues[pos]);
+      pos++;
+    }
+    reifiedReturnShapes.emplace_back(std::move(shapes));
+  }
+  return success();
 }
 
 LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -20,6 +20,7 @@
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
 #include "llvm/ADT/DenseMap.h"
@@ -770,33 +771,6 @@
     return success();
   }
 };
-
-/// Canonicalize a `linalg.init_tensor` -> `dim` pattern by replacing the `dim`
-/// with
-/// - A constant value if the size is static along the dimension.
-/// - The dynamic value that defines the size of the result of
-///   `linalg.init_tensor` op.
-struct ReplaceDimOfInitTensorOp : public OpRewritePattern<DimOp> {
-  using OpRewritePattern<DimOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    auto initTensorOp = dimOp.memrefOrTensor().getDefiningOp<InitTensorOp>();
-    if (!initTensorOp)
-      return failure();
-    auto dimIndex = dimOp.index().getDefiningOp<ConstantIndexOp>();
-    if (!dimIndex)
-      return failure();
-    int64_t index = dimIndex.getValue();
-    if (!initTensorOp.isDynamicSize(index)) {
-      rewriter.replaceOpWithNewOp<ConstantIndexOp>(
-          dimOp, initTensorOp.getStaticSize(index));
-    } else {
-      rewriter.replaceOp(dimOp, initTensorOp.getDynamicSize(index));
-    }
-    return success();
-  }
-};
 } // namespace
 
 namespace {
@@ -830,10 +804,12 @@
     if (!reshapeOp.src().getDefiningOp<InitTensorOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
-    SmallVector<Value, 4> resultShapeValues =
-        reshapeOp.getOutputShape(rewriter, loc);
+    SmallVector<SmallVector<Value>, 4> resultShapes;
+    if (failed(reshapeOp.reifyReturnTypeShapes(rewriter, resultShapes)) ||
+        !llvm::hasSingleElement(resultShapes))
+      return failure();
     Value initTensor = rewriter.create<InitTensorOp>(
-        loc, resultShapeValues, reshapeOp.getResultType().getElementType());
+        loc, resultShapes[0], reshapeOp.getResultType().getElementType());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(
         reshapeOp, reshapeOp.getResultType(), initTensor);
     return success();
@@ -843,9 +819,23 @@
 
 void InitTensorOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results
-      .insert<FoldInitTensorWithSubTensorOp, FoldInitTensorWithTensorReshapeOp,
-              ReplaceDimOfInitTensorOp, ReplaceStaticShapeDims>(context);
+  results.insert<FoldInitTensorWithSubTensorOp,
+                 FoldInitTensorWithTensorReshapeOp, ReplaceStaticShapeDims>(
+      context);
+}
+
+LogicalResult InitTensorOp::reifyReturnTypeShapes(
+    OpBuilder &builder,
+    SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  Location loc = getLoc();
+  auto shapes = llvm::to_vector<4>(llvm::map_range(
+      llvm::seq<int64_t>(0, getType().getRank()), [&](int64_t dim) -> Value {
+        if (isDynamicSize(dim))
+          return getDynamicSize(dim);
+        return builder.create<ConstantIndexOp>(loc, getStaticSize(dim));
+      }));
+  reifiedReturnShapes.emplace_back(std::move(shapes));
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
@@ -997,6 +987,33 @@
                                         builder);
 }
 
+LogicalResult PadTensorOp::reifyReturnTypeShapes(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  Location loc = getLoc();
+  auto getAsValues = [&](ArrayRef<OpFoldResult> foldResults) {
+    return llvm::to_vector<4>(
+        llvm::map_range(foldResults, [&](OpFoldResult foldResult) -> Value {
+          if (Attribute attr = foldResult.dyn_cast<Attribute>())
+            return b.create<ConstantIndexOp>(loc,
+                                             attr.cast<IntegerAttr>().getInt());
+          return foldResult.get<Value>();
+        }));
+  };
+  auto lowPad = getAsValues(getMixedLowPad());
+  auto highPad = getAsValues(getMixedHighPad());
+  auto shape = llvm::to_vector<4>(llvm::map_range(
+      llvm::seq<int64_t>(0, getSourceType().getRank()),
+      [&](int64_t dim) -> Value {
+        Value sourceDim = b.create<DimOp>(loc, source(), dim);
+        AffineExpr expr = b.getAffineDimExpr(0) + b.getAffineSymbolExpr(0) +
+                          b.getAffineSymbolExpr(1);
+        return applyMapToValues(b, loc, AffineMap::get(1, 2, expr),
+                                {sourceDim, lowPad[dim], highPad[dim]})[0];
+      }));
+  reifiedReturnShapes.emplace_back(std::move(shape));
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
@@ -1384,9 +1401,10 @@
       }));
 }
 
-SmallVector<Value, 4> mlir::linalg::getReshapeOutputShapeFromInputShape(
-    OpBuilder &builder, Location loc, Value src,
-    ArrayRef<int64_t> dstStaticShape, ArrayRef<AffineMap> reassocation) {
+static SmallVector<Value, 4>
+getReshapeOutputShapeFromInputShape(OpBuilder &builder, Location loc, Value src,
+                                    ArrayRef<int64_t> dstStaticShape,
+                                    ArrayRef<AffineMap> reassocation) {
   return dstStaticShape.size() >
                  static_cast<size_t>(src.getType().cast<ShapedType>().getRank())
              ? getExpandedOutputShapeFromInputShape(
@@ -1635,35 +1653,20 @@
     return success();
   }
 };
-
-/// Canonicalize dim ops that use the output shape with dim of the input.
-struct ReplaceDimOfReshapeOpResult : OpRewritePattern<DimOp> {
-  using OpRewritePattern<DimOp>::OpRewritePattern;
-  LogicalResult matchAndRewrite(DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    Value dimValue = dimOp.memrefOrTensor();
-    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
-    if (!dimIndex)
-      return failure();
-
-    auto reshapeOp = dimValue.getDefiningOp<TensorReshapeOp>();
-    if (!reshapeOp)
-      return failure();
-
-    rewriter.replaceOp(dimOp,
-                       getReshapeOutputDimFromInputShape(
-                           rewriter, dimOp.getLoc(), *dimIndex, reshapeOp.src(),
-                           reshapeOp.getResultType().getShape(),
-                           reshapeOp.getReassociationMaps()));
-    return success();
-  }
-};
 } // namespace
 
 void TensorReshapeOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results.insert<CollapseReshapeOps<TensorReshapeOp>, FoldReshapeWithConstant,
-                 ReplaceDimOfReshapeOpResult>(context);
+  results.insert<CollapseReshapeOps<TensorReshapeOp>, FoldReshapeWithConstant>(
+      context);
+}
+
+LogicalResult TensorReshapeOp::reifyReturnTypeShapes(
+    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+  auto resultShape = getReshapeOutputShapeFromInputShape(
+      b, getLoc(), src(), getResultType().getShape(), getReassociationMaps());
+  reifiedReturnShapes.emplace_back(std::move(resultShape));
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
@@ -2444,49 +2447,6 @@
     return success();
   }
 };
-
-/// Replaces std.dim operations that use the result of a LinalgOp (on tensors)
-/// with std.dim operations that use one of the arguments. For example,
-///
-/// %0 = linalg.matmul ins(%arg0, %arg1, ...)
-/// %1 = dim %0, %c0
-///
-/// with
-///
-/// %1 = dim %arg0, %c0
-///
-/// where possible. With this the result of the `linalg.matmul` is not used in
-/// dim operations. If the value produced is replaced with another value (say by
-/// tiling `linalg.matmul`) will make the `linalg.matmul` truly dead instead of
-/// used in a dim op that would prevent the DCE of this op.
-struct ReplaceDimOfLinalgOpResult : public OpRewritePattern<DimOp> {
-  using OpRewritePattern<DimOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(DimOp dimOp,
-                                PatternRewriter &rewriter) const override {
-    Value dimValue = dimOp.memrefOrTensor();
-    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
-    if (!dimIndex)
-      return failure();
-    auto linalgOp = dimValue.getDefiningOp<LinalgOp>();
-    if (!linalgOp)
-      return failure();
-
-    unsigned resultIndex = dimValue.cast<OpResult>().getResultNumber();
-    Optional<Value> operandDimValue = linalgOp.inferResultDimFromInputShapes(
-        rewriter, dimOp.getLoc(), resultIndex,
-        static_cast<unsigned>(*dimIndex));
-    if (!operandDimValue) {
-      // Its always possible to replace using the corresponding `outs`
-      // parameter.
-      operandDimValue = rewriter.create<DimOp>(
-          dimOp.getLoc(), linalgOp.getOutput(resultIndex), *dimIndex);
-    }
-    rewriter.replaceOp(dimOp, *operandDimValue);
-    return success();
-  }
-};
-
 } // namespace
 
 namespace {
@@ -2649,7 +2609,6 @@
                                         MLIRContext *context) {                \
     results.insert<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp,     \
                    RemoveIdentityLinalgOps>();                                 \
-    results.insert<ReplaceDimOfLinalgOpResult>(context);                       \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
diff --git a/mlir/lib/Dialect/StandardOps/CMakeLists.txt b/mlir/lib/Dialect/StandardOps/CMakeLists.txt
--- a/mlir/lib/Dialect/StandardOps/CMakeLists.txt
+++ b/mlir/lib/Dialect/StandardOps/CMakeLists.txt
@@ -15,6 +15,7 @@
   MLIRCastInterfaces
   MLIRControlFlowInterfaces
   MLIREDSC
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRTensor
diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
--- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
+++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp
@@ -21,6 +21,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Value.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Support/MathExtras.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/STLExtras.h"
@@ -1556,12 +1557,45 @@
     return success();
   }
 };
+
+/// Fold dim of an operation that implements the InferShapedTypeOpInterface
+struct DimOfShapedTypeOpInterface : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dimOp,
+                                PatternRewriter &rewriter) const override {
+    OpResult dimValue = dimOp.memrefOrTensor().dyn_cast<OpResult>();
+    if (!dimValue)
+      return failure();
+    auto shapedTypeOp =
+        dyn_cast<InferShapedTypeOpInterface>(dimValue.getOwner());
+    if (!shapedTypeOp)
+      return failure();
+
+    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
+    if (!dimIndex)
+      return failure();
+
+    SmallVector<SmallVector<Value>> returnShapes;
+    if (failed(shapedTypeOp.reifyReturnTypeShapes(rewriter, returnShapes)))
+      return failure();
+
+    unsigned resultNumber = dimValue.getResultNumber();
+    if (returnShapes.size() <= resultNumber ||
+        returnShapes[resultNumber].size() <= static_cast<size_t>(*dimIndex))
+      return failure();
+
+    rewriter.replaceOp(dimOp, returnShapes[resultNumber][*dimIndex]);
+    return success();
+  }
+};
 } // end anonymous namespace.
 
 void DimOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
   results.insert<DimOfMemRefReshape, DimOfCastOp<TensorToMemrefOp>,
-                 DimOfCastOp<tensor::CastOp>>(context);
+                 DimOfCastOp<tensor::CastOp>, DimOfShapedTypeOpInterface>(
+      context);
 }
 
 // ---------------------------------------------------------------------------
diff --git a/mlir/lib/Interfaces/InferTypeOpInterface.cpp b/mlir/lib/Interfaces/InferTypeOpInterface.cpp
--- a/mlir/lib/Interfaces/InferTypeOpInterface.cpp
+++ b/mlir/lib/Interfaces/InferTypeOpInterface.cpp
@@ -13,8 +13,6 @@
 
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 
-#include "mlir/IR/BuiltinTypes.h"
-
 using namespace mlir;
 
 namespace mlir {
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -683,9 +683,10 @@
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
-    OpBuilder &builder, llvm::SmallVectorImpl<Value> &shapes) {
-  shapes = SmallVector<Value, 1>{
-      builder.createOrFold<DimOp>(getLoc(), getOperand(0), 0)};
+    OpBuilder &builder,
+    llvm::SmallVectorImpl<llvm::SmallVector<Value>> &shapes) {
+  shapes.push_back(SmallVector<Value, 1>{
+      builder.createOrFold<DimOp>(getLoc(), getOperand(0), 0)});
   return success();
 }
 
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -506,13 +506,14 @@
 }
 
 def OpWithInferTypeInterfaceOp : TEST_Op<"op_with_infer_type_if", [
-    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+    DeclareOpInterfaceMethods<InferTypeOpInterface,
+        ["inferReturnTypeComponents"]>]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
 def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_type_if",
-      InferTensorTypeWithReify.traits> {
+      InferTensorType<["inferReturnTypeComponents", "reifyReturnTypeShapes"]>.traits> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
diff --git a/mlir/test/lib/Dialect/Test/TestPatterns.cpp b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
--- a/mlir/test/lib/Dialect/Test/TestPatterns.cpp
+++ b/mlir/test/lib/Dialect/Test/TestPatterns.cpp
@@ -126,10 +126,11 @@
 
   // Use permutations of 2 args as operands.
   auto shapedOp = cast<OpWithShapedTypeInferTypeInterfaceOp>(op);
-  SmallVector<Value, 2> shapes;
-  if (failed(shapedOp.reifyReturnTypeShapes(b, shapes)))
+  SmallVector<SmallVector<Value>, 2> shapes;
+  if (failed(shapedOp.reifyReturnTypeShapes(b, shapes)) ||
+      !llvm::hasSingleElement(shapes))
     return;
-  for (auto it : llvm::enumerate(shapes))
+  for (auto it : llvm::enumerate(shapes[0]))
     op->emitRemark() << "value " << it.index() << ": "
                      << it.value().getDefiningOp();
 }