[mlir] Enhance InferShapedTypeOpInterface and move LinalgOps to use them.

[mravishankar]

Summary

For most Linalg operations, the shape of the output can be computed
from the shape.values of the inputs. This change enhance the
InferShapedTypeOpInterface to allow moving all Linalg Structured
operations as well as some other Linalg ops to use this
interface. Adding a canonicalization to dim ops that are getting the
dim of the result of an op that implements this interface allows
removing many patterns across these ops that were implementing the
same functionality.

[nicolasvasilache]

Nice!

@stellaraccident this is getting towards something pluggable from the DSL.

[stellaraccident]

Nice! So then if i plug the shape maps through named ops, I can also implement the interface, right?

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

165

Since this is generic, do you also want to early exit if there are no results?

[mravishankar]

In D97532#2589666, @stellaraccident wrote:

Nice! So then if i plug the shape maps through named ops, I can also implement the interface, right?

@nicolasvasilache and @stellaraccident . I guess I kind of cheated here. Right now this assumes that all the Linalg Structured ops use the default of mapping shapes -> loops and loops -> shapes, so you can resolve the output shapes based on input shapes using the affine maps. The correct way to this would be to have a separate Interface method on the LinalgOp that provides a default implementation. Each operation for which the default affine map based shape computation is not valid will then have to overide this default implementation of the linalg ops. I think its a change that is possible to do right now. Will make that change. Then the python front end could do nothing and named ops will pick up the standard method. Or they can override the default interface on the LinalgOp to provide their own shape transfer function.
Let me know if you see any gotchas here.

[jpienaar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

109

Does this result correspond to operand of op as defined for the op or as part of "flattened"?

115

dim isn't documented. Is this saying that if you have a (tensor<10x20xf32>) you can specify 0, 1 to get constant 10?

It would seem you get exactly this by reify & get_element_at & canonicalizing. How many of these custom helpers are generated by creating all these methods vs composition? What is the impact on codesize? These seem very pro cases where the dimension computation is very separate and so this is cheap, I think that is a special case.

118

By your definition, the below is only valid if there is a single result.

157

Why this restriction? This is returning an array of results in particular because each of them is for a different value.

168

Again this is making the assumption that computing each dim indivdualy is cheaper than the whole, so just compute all and make the base case the per index and dim one.

[benvanik]

Thanks for doing this Mahesh! This will really help unblock me.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1570–1571

this should probably be dyn_cast_or_null, in the case where dim is used on a block/function argument (who have null owners)

[silvas]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

122

Remove the special case for one result. It's not hard to just implement the more general case (just add assert(resultIndex == 0) at the top.

[mravishankar]

Rebase.

[mravishankar]

After an offline conversation with Jacques, I realize I had misread the existing interface. I havent changed anything in this patch, but a follow up patch (https://reviews.llvm.org/D97887) goes back to the previous interface and each op that uses this interface needs to use the reifyReturnTypeShapes to compute the shape of all results. PTAL and let me know which works better. I prefer the one in the patch (https://reviews.llvm.org/D97887) cause it is less ambiguous.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

109

This was meant to be the result index if the op returns multiple indices.

115

See followup patch that addresses this concern to some extent.

W.R.T codesize, AFAIK, the methods are just default implementations that are on the interface (and every op that needs to override it). So it shouldnt impact codesize, though there is some vtable jumping that will happen on invoking this method.

118

Yeah, the documentation needs to be changed. Will take a pass at it based on comments with follow up patch (this method might not even be needed)

122

Its useful to have special case for resultIndex == 0. See the canonicalization on DimOp. I dont need to call different interface methods depending on number of results of the operation.

Overall though, I found this interface is maybe a bit clunky. The follow up patch goes back to the old interface and each op is expected to return the shape of all its results. That works well in practice with canonicalizations removing all unnecessary code.

157

I misread the interface as I mentioned to you. See follow up patch.

165

Sure, I will add that, but please do look at the follow up patch. This issue will not exists on that interface.

168

See followup patch that avoids this.

[mravishankar]

Retiring this. its folded with the dependent patch.

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.td

Operation::operand_range getAssumedNonShapedOperands() {
         assert((t.isSignlessIntOrIndexOrFloat() || t.template isa<VectorType>())
                &&"expected scalar or vector type");
       }
       return res;
     }
 
     /// 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);
+    Value reifyReturnTypeShapeForResult(OpBuilder &b, unsigned resultIdx,
+        int64_t dim);
 
     //========================================================================//
     // Helper functions to mutate the `operand_segment_sizes` attribute.
     // These are useful when cloning and changing operand types.
     //========================================================================//
     void setNumInputs(unsigned num) { setOperandSegmentAt(0, num); }
     void setNumOutputBuffers(unsigned num) { setOperandSegmentAt(1, num); }
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.h

 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
 #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"
 
 #include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
 namespace linalg {

 SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
                                   ArrayRef<AffineExpr> b);
 
 /// Return the dims that are `iteratorTypeName` loops in the LinalgOp `op`.
 /// Assumes `op` is a LinalgOp.
 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
 } // namespace linalg
 } // namespace mlir
 
 namespace mlir {
 namespace linalg {

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_OPS
 #define LINALG_OPS
 
 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"
 
 // Base class for Linalg dialect ops that do not correspond to library calls.
 class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Linalg_Dialect, mnemonic, traits> {
   // For every linalg op, there needs to be a:
   //   * void print(OpAsmPrinter &p, ${C++ class of Op} op)
   //   * LogicalResult verify(${C++ class of Op} op)
   //   * ParseResult parse${C++ class of Op}(OpAsmParser &parser,
   //                                         OperationState &result)
   // functions.
   let printer = [{ return ::print(p, *this); }];
   let verifier = [{ return ::verify(*this); }];
   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,
+         ["reifyReturnTypeShape"]>]> {
   let summary = "operation to define a tensor of particular value";
 
   let description = [{
     `linalg.init_tensor` is an operation that materializes a tensor of
     a given shape. The shape could be dynamic or static.
   }];
 
   let arguments =
let builders = [
     OpBuilderDAG<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let hasCanonicalizer = 1;
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments, NoSideEffect]> {
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShape"]>,
+     NoSideEffect]> {
   let summary = "tensor pad operation";
   let description = [{
     `linalg.pad_tensor` is an operation that pads the `source` tensor
     with given `low` and `high` padding config.
 
     The PadTensor operation supports the following arguments:
 
     * source: the "base" tensor on which to pad.
code commonExtraClassDeclaration = [{
     SmallVector<ReassociationExprs, 4> getReassociationExprs() {
       return
         llvm::to_vector<4>(llvm::map_range(reassociation(),
           [](Attribute a) {
             return llvm::to_vector<2>(
               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)
   }];
 }
 
 def Linalg_ReshapeOp : Linalg_ReshapeLikeOp<"reshape",
     [DeclareOpInterfaceMethods<ViewLikeOpInterface>]>,
def Linalg_ReshapeOp : Linalg_ReshapeLikeOp<"reshape",
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     MemRefType getSrcType() { return src().getType().cast<MemRefType>(); }
     MemRefType getResultType() { return result().getType().cast<MemRefType>(); }
   }];
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
-def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
+def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<
+    "tensor_reshape",
+    [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+      ["reifyReturnTypeShape", "reifyReturnTypeShapes"]>]>,
     Arguments<(ins AnyTensor:$src,
                    AffineMapArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
   let summary = "linalg.tensor_reshape produces a new reshaped tensor.";
   let description = [{
     The `linalg.reshape` op produces a new tensor whose sizes are a
     reassociation of the original `src`.
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

 //===----------------------------------------------------------------------===//
 
 #ifndef LINALG_STRUCTURED_OPS
 #define LINALG_STRUCTURED_OPS
 
 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.
 // Linalg ops that correspond to library calls operate on ShapedType as their
 // first operands. These may be optionally followed by non-view operands
 // 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
     // equal to numLoops.
     unsigned getNumPayloadInductionVariables() {
       return isa<IndexedGenericOp>(this->getOperation()) ? getNumLoops() : 0;
     }
+
+    Value reifyReturnTypeShapeForResult(OpBuilder &b, unsigned resultIdx,
+        int64_t dim) {
+      return cast<LinalgOp>(getOperation()).
+          reifyReturnTypeShapeForResult(b, resultIdx, dim);
+    }
   }];
 }
 
 class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
   : LinalgStructuredBase_Op<mnemonic,
        !listconcat(props, [
          DeclareOpInterfaceMethods<MemoryEffectsOpInterface>])> {
   code structuredOpsDecls = structuredOpsBaseDecls # [{

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_
 #define MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_
 
 #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"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 /// ShapedTypeComponents that represents the components of a ShapedType.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

StaticInterfaceMethod<
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"inferReturnTypeComponents",
       /*args=*/(ins "::mlir::MLIRContext*":$context,
                     "::mlir::Optional<::mlir::Location>":$location,
                     "::mlir::ValueRange":$operands,
                     "::mlir::DictionaryAttr":$attributes,
                     "::mlir::RegionRange":$regions,
                     "::mlir::SmallVectorImpl<::mlir::ShapedTypeComponents>&":
-                      $inferredReturnShapes)
+                      $inferredReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return failure(); }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{Reify the shape computation for the operation.
+
+      Insert operations using the given OpBuilder that computes the result
+      shape along a particular dimension of result at `resultIndex`.
+      }],
+      /*retTy=*/"::mlir::Value",
+      /*methodName=*/"reifyReturnTypeShapeForResult",
+      /*args=*/(ins "::mlir::OpBuilder&":$builder,
+                    "unsigned":$resultIndex,
+                    "int64_t":$dim),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        if (resultIndex == 0) return $_op.reifyReturnTypeShape(builder, dim);
+        return nullptr; }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{Reify the shape computation for the operation.
+
+      Insert operations using the given OpBuilder that computes the
+      result shape along a particular dimension. Valid only when the
+      operation has a single result value.
+      }],
+      /*retTy=*/"::mlir::Value",
+      /*methodName=*/"reifyReturnTypeShape",
+      /*args=*/(ins "::mlir::OpBuilder&":$builder,
+                    "int64_t":$dim),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return nullptr; }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{Reify the shape computation for the operation.
+
+      Insert operations using the given OpBuilder that computes the result
+      shape for a result at `resultIndex`.
+      }],
+      /*retTy=*/"::mlir::LogicalResult",
+      /*methodName=*/"reifyReturnTypeShapesForResult",
+      /*args=*/(ins "::mlir::OpBuilder&":$builder,
+                    "unsigned":$resultIndex,
+                    "::mlir::SmallVectorImpl<::mlir::Value>&":$reifiedReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{
+        if (resultIndex == 0)
+          return $_op.reifyReturnTypeShapes(builder, reifiedReturnShapes);
+        return ::mlir::failure();
+      }]
     >,
     InterfaceMethod<
       /*desc=*/[{Reify the shape computation for the operation.
 
       Insert operations using the given OpBuilder that computes the result
-      shape.
+      shape. Valid only when the operation has a single result value.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
                     "::mlir::SmallVectorImpl<::mlir::Value>&":$reifiedReturnShapes),
       /*methodBody=*/[{}],
-      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
+      /*defaultImplementation=*/[{
+        ShapedType type =
+          $_op.getOperation()->getResult(0).getType().template dyn_cast<ShapedType>();
+        if (!type) return failure();
+        for (auto dim : llvm::seq<int64_t>(0, type.getRank())) {
+          Value shape = $_op.reifyReturnTypeShape(builder, dim);
+          if (!shape) return failure();
+          reifiedReturnShapes.push_back(shape);
+        }
+        return success();
+      }]
     >,
   ];
 }
 
 #endif // MLIR_INFERTYPEOPINTERFACE

mlir/lib/Dialect/Linalg/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalg
   LinalgInterfaces.cpp
   LinalgOps.cpp
   LinalgTypes.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgInterfacesIncGen
   MLIRLinalgOpsIncGen
   MLIRLinalgStructuredOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffine
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
   MLIRStandard
   MLIRTensor
   )

mlir/lib/Dialect/Linalg/IR/LinalgInterfaces.cpp

struct HasAffineDimExprVisitor
   bool visitConstantExpr(AffineConstantExpr constExpr) { return false; }
 
   bool visitSymbolExpr(AffineSymbolExpr symbolExpr) { return false; }
 
 private:
   llvm::SmallSet<unsigned, 4> positions;
 };
 
-Optional<Value> LinalgOp::inferResultDimFromInputShapes(OpBuilder &b,
-                                                        Location loc,
-                                                        unsigned resultIdx,
-                                                        unsigned dim) {
+Value LinalgOp::reifyReturnTypeShapeForResult(OpBuilder &b, unsigned resultIdx,
+                                              int64_t dim) {
   // 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)
   //   shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
   //   resultFromFromInputDim = subMapOfResultDim.compose(shapesToLoopMap)
unsigned outputDimPosEnd =
                                                          .getType()
                                                          .cast<ShapedType>()
                                                          .getRank() -
                                                      1)
           .getValue();
   llvm::for_each(llvm::seq<unsigned>(outputDimPosStart, outputDimPosEnd),
                  [&outputDims](unsigned dim) { outputDims.insert(dim); });
   HasAffineDimExprVisitor checkDimExpr(outputDims);
+  Location loc = getOperation()->getLoc();
   if (checkDimExpr.visit(operandShapesToResultDimMap.getResult(0)))
-    return llvm::None;
+    return b.createOrFold<DimOp>(loc, getOutput(resultIdx), dim);
   return applyMapToValues(b, loc, operandShapesToResultDimMap,
                           createFlatListOfOperandDims(b, loc))[0];
 }
 
 LogicalResult mlir::linalg::detail::verifyStructuredOpInterface(Operation *op) {
   LinalgOp linalgOp = cast<LinalgOp>(op);
   // Expect at least one shaped operand.
   // This means an op that constructs a tensor out of indices cannot be a

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
if (newType == op.getType())
       return failure();
     auto newOp =
         rewriter.create<InitTensorOp>(op.getLoc(), newType, dynamicSizes,
                                       rewriter.getI64ArrayAttr(staticSizes));
     rewriter.replaceOpWithNewOp<tensor::CastOp>(op, op.getType(), newOp);
     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 {
 /// Since `init_tensor` operation creates a tensor needed only for its shape, a
 /// subtensor of this is also needed only for its shape. The result can be
 /// replaced by a new init_tensor operation of the same size as the subtensor
 /// op.
 struct FoldInitTensorWithSubTensorOp : public OpRewritePattern<SubTensorOp> {
struct FoldInitTensorWithTensorReshapeOp
     : public OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.src().getDefiningOp<InitTensorOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
-    SmallVector<Value, 4> resultShapeValues =
-        reshapeOp.getOutputShape(rewriter, loc);
+    SmallVector<Value, 4> resultShapeValues;
+    if (failed(reshapeOp.reifyReturnTypeShapesForResult(rewriter, 0,
+                                                        resultShapeValues)))
+      return failure();
     Value initTensor = rewriter.create<InitTensorOp>(
         loc, resultShapeValues, reshapeOp.getResultType().getElementType());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(
         reshapeOp, reshapeOp.getResultType(), initTensor);
     return success();
   }
 };
 } // namespace
 
 void InitTensorOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
-  results
-      .insert<FoldInitTensorWithSubTensorOp, FoldInitTensorWithTensorReshapeOp,
-              ReplaceDimOfInitTensorOp, ReplaceStaticShapeDims>(context);
+  results.insert<FoldInitTensorWithSubTensorOp,
+                 FoldInitTensorWithTensorReshapeOp, ReplaceStaticShapeDims>(
+      context);
+}
+
+Value InitTensorOp::reifyReturnTypeShape(OpBuilder &builder, int64_t dim) {
+  if (isDynamicSize(dim))
+    return getDynamicSize(dim);
+  return builder.create<ConstantIndexOp>(getLoc(), getStaticSize(dim));
 }
 
 //===----------------------------------------------------------------------===//
 // PadTensorOp
 //===----------------------------------------------------------------------===//
 
 /// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
 static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
for (int i = 0; i < rank; ++i) {
     auto highValue = builder.createOrFold<SubIOp>(loc, resultDimSize, dimOp);
     high.push_back(highValue);
     low.push_back(builder.createOrFold<ConstantIndexOp>(loc, 0));
   }
   return PadTensorOp::createPadScalarOp(type, source, pad, low, high, loc,
                                         builder);
 }
 
+Value PadTensorOp::reifyReturnTypeShape(OpBuilder &b, int64_t dim) {
+  Location loc = getLoc();
+  auto getAsValue = [&](OpFoldResult valueOrAttr) -> Value {
+    if (Attribute attr = valueOrAttr.dyn_cast<Attribute>())
+      return b.create<ConstantIndexOp>(loc, attr.cast<IntegerAttr>().getInt());
+    return valueOrAttr.get<Value>();
+  };
+  auto lowPad = getAsValue(getMixedLowPad()[dim]);
+  auto highPad = getAsValue(getMixedHighPad()[dim]);
+  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, highPad})[0];
+}
+
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//
 
 /// Collapse reassociation maps that are used in pair of reshape ops where one
 /// is a producer and other is the consumer. Only valid to use this method when
 /// both the producer and consumer are collapsing dimensions or both are
 /// expanding dimensions.
static SmallVector<Value, 4> getExpandedOutputShapeFromInputShape(
   return llvm::to_vector<4>(llvm::map_range(
       llvm::seq<int64_t>(0, dstStaticShape.size()), [&](int64_t dim) {
         return getExpandedOutputDimFromInputShape(builder, loc, dim, src,
                                                   dstStaticShape, reassociation,
                                                   expandedDimToCollapsedDim);
       }));
 }
 
-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(
                    builder, loc, src, dstStaticShape, reassocation)
              : getCollapsedOutputShapeFromInputShape(
                    builder, loc, src, dstStaticShape, reassocation);
 }
 
LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
     if (!attr || !attr.isSplat())
       return failure();
     DenseElementsAttr newAttr = DenseElementsAttr::getFromRawBuffer(
         reshapeOp.getResultType(), attr.getRawData(), true);
     rewriter.replaceOpWithNewOp<ConstantOp>(reshapeOp, newAttr);
     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);
+}
+
+Value TensorReshapeOp::reifyReturnTypeShape(OpBuilder &b, int64_t dim) {
+  return getReshapeOutputDimFromInputShape(b, getLoc(), dim, src(),
+                                           getResultType().getShape(),
+                                           getReassociationMaps());
+}
+
+LogicalResult TensorReshapeOp::reifyReturnTypeShapes(
+    OpBuilder &b, SmallVectorImpl<Value> &reifiedReturnShapes) {
+  auto resultShape = getReshapeOutputShapeFromInputShape(
+      b, getLoc(), src(), getResultType().getShape(), getReassociationMaps());
+  reifiedReturnShapes.append(resultShape.begin(), resultShape.end());
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // YieldOp
 //===----------------------------------------------------------------------===//
 
 static void print(OpAsmPrinter &p, linalg::YieldOp op) {
   p << op.getOperationName();
for (auto t : containers)
       argTypes.push_back(getElementTypeOrSelf(t));
 
   // RAII.
   OpBuilder::InsertionGuard guard(opBuilder);
   Block *body = opBuilder.createBlock(&region, /*insertPt=*/{}, argTypes);
   unsigned actual = body->getNumArguments();
   unsigned expected = NamedStructuredOpType::getNumRegionArgs();
   if (expected != actual) {
-    if (errorHandler) errorHandler(expected, actual);
+    if (errorHandler)
+      errorHandler(expected, actual);
     return;
   }
 
   opBuilder.setInsertionPointToStart(body);
   mlir::edsc::ScopedContext scope(opBuilder, opBuilder.getUnknownLoc());
   NamedStructuredOpType::regionBuilder(*body, captures);
 
   // indexing_maps is an auto-generated method.
LogicalResult matchAndRewrite(Operation *op,
     // Clone op.
     Operation *newOp =
         linalgOp.clone(rewriter, op->getLoc(), newResultTypes, newOperands);
     rewriter.replaceOp(op, newOp->getResults());
 
     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 {
 // Deduplicate redundant args of a linalg op.
 // An arg is redundant if it has the same Value and indexing map as another.
 struct DeduplicateInputs : public RewritePattern {
   DeduplicateInputs(PatternBenefit benefit = 1)
       : RewritePattern(benefit, MatchAnyOpTypeTag()) {}

 };
 } // namespace
 
 #define CANONICALIZERS_AND_FOLDERS(XXX)                                        \
   void XXX::getCanonicalizationPatterns(OwningRewritePatternList &results,     \
                                         MLIRContext *context) {                \
     results.insert<DeduplicateInputs, EraseDeadLinalgOp, FoldTensorCastOp,     \
                    RemoveIdentityLinalgOps>();                                 \
-    results.insert<ReplaceDimOfLinalgOpResult>(context);                       \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
                           SmallVectorImpl<OpFoldResult> &) {                   \
     return foldMemRefCast(*this);                                              \
   }
 
 CANONICALIZERS_AND_FOLDERS(ConvOp)

mlir/lib/Dialect/StandardOps/CMakeLists.txt

add_mlir_dialect_library(MLIRStandard
   DEPENDS
   MLIRStandardOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRCallInterfaces
   MLIRCastInterfaces
   MLIRControlFlowInterfaces
   MLIREDSC
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRTensor
   MLIRVectorInterfaces
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Transforms)

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #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"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 // Pull in all enum type definitions and utility function declarations.
LogicalResult matchAndRewrite(DimOp dimOp,
     auto castOp = dimOp.memrefOrTensor().getDefiningOp<CastOpTy>();
     if (!castOp)
       return failure();
     Value newSource = castOp.getOperand();
     rewriter.replaceOpWithNewOp<DimOp>(dimOp, newSource, dimOp.index());
     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();
+
+    Value dimReplacement = shapedTypeOp.reifyReturnTypeShapeForResult(
+        rewriter, dimValue.getResultNumber(), *dimIndex);
+    if (!dimReplacement)
+      return failure();
+
+    rewriter.replaceOp(dimOp, dimReplacement);
+    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);
 }
 
 // ---------------------------------------------------------------------------
 // DivFOp
 // ---------------------------------------------------------------------------
 
 OpFoldResult DivFOp::fold(ArrayRef<Attribute> operands) {
   return constFoldBinaryOp<FloatAttr>(

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 //===- InferTypeOpInterface.cpp - Infer Type Interfaces ---------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains the definitions of the infer op interfaces defined in
 // `InferTypeOpInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 
-#include "mlir/IR/BuiltinTypes.h"
-
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
 LogicalResult mlir::detail::inferReturnTensorTypes(
     function_ref<LogicalResult(

mlir/test/lib/Dialect/Test/TestOps.td

def I32ElementsAttrOp : TEST_Op<"i32ElementsAttr"> {
   let arguments = (ins I32ElementsAttr:$attr);
 }
 
 def IndexElementsAttrOp : TEST_Op<"indexElementsAttr"> {
   let arguments = (ins IndexElementsAttr:$attr);
 }
 
 def OpWithInferTypeInterfaceOp : TEST_Op<"op_with_infer_type_if", [
-    DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+    DeclareOpInterfaceMethods<InferTypeOpInterface,
+        ["inferReturnTypeComponents"]>]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
 def InferTensorType : NativeOpTrait<"InferTensorType">;
 def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_type_if",
   [
      // Op implements infer type op interface.
      InferTypeOpInterface,
      // The op will have methods implementing the ShapedType type infer interface.
-     DeclareOpInterfaceMethods<InferShapedTypeOpInterface>,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+        ["inferReturnTypeComponents"]>,
      // The op produces tensors and will use the ShapedType type infer interface
      // along with knowledge that it is producing Tensors to infer shape.
      InferTensorType
    ]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 
   let extraClassDeclaration = [{