Split InferShapedTypeOpInterface to create ReifyRankedShapeTypeShapeOpInterface.

[mravishankar]

Summary

The reifyReturnTypeShapesPerResultDim method supports shape
inference for rsults that are ranked types. These are used lower in
the codegeneration stack than its counter part reifyReturnTypeShapes
which also supports unranked types, and is more suited for use higher
up the compilation stack. To have separation of concerns, this method
is split into its own interface.
See discussion : https://llvm.discourse.group/t/better-layering-for-infershapedtypeopinterface/3823

[jpienaar]

Looks like good cleanup thanks. I just focused interface side as I don't know the rest as well.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

150

Not to bike shed, but is the additonal shape needed?

ShapedType is the c++ class, and just ReifyRankedShapedTypeOpInterface seems descriptive and unambiguous

155

Vector too correct?

165

Indent more to avoid the document generator adding weird spaces :)

170

The last sentence is probably not needed given interface name.

177

OOC Should we perhaps add a typedef for this?

[mravishankar]

Addressing comments

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

150

Sounds good. I just copy pasted from Discourse. I like this better.

155

Sure, but vectors so far are statically shaped. So the interface not really needed there.

165

Ugh! Tabs. Fixed.

[mravishankar]

Rebase

[mravishankar]

Fix build warning and update bazel file

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

 #define MLIR_DIALECT_LINALG_IR_LINALGINTERFACES_H_
 
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 
 namespace mlir {
 namespace linalg {
 class LinalgOp;
 
 /// OpOperand vector that implicitly converts to a Value vector.
 struct OpOperandVector : public SmallVector<OpOperand *> {

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

let extraClassDeclaration = [{
 
     /// Compute the static loop sizes necessary to vectorize the computation.
     /// This is done by applying `getShapesToLoopsMap` to
     /// `createFlatListOfOperandStaticDims`.
     SmallVector<int64_t, 4> computeStaticLoopSizes();
 
     /// Returns the value that expresses the shape of the output in terms of
     /// shape of the input operands where possible
-    LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
-        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes);
+    LogicalResult reifyResultShapes(OpBuilder &b,
+        ReifiedRankedShapedTypeDims &reifiedReturnShapes);
 
     //========================================================================//
     // 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.td

class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
   // 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,
-     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-         ["reifyReturnTypeShapesPerResultDim"]>]> {
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
   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 = [
     OpBuilder<(ins "ArrayRef<OpFoldResult>":$sizes, "Type":$elementType,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let hasCanonicalizer = 1;
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments,
-     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-         ["reifyReturnTypeShapesPerResultDim"]>,
-     NoSideEffect]> {
+    [AttrSizedOperandSegments, NoSideEffect,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
   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.
class Linalg_ReshapeLikeOp<string mnemonic, list<OpTrait> traits = []> :
   let parser = [{ return ::parseReshapeLikeOp(parser, result); }];
 }
 
 def IndexListArrayAttr :
   TypedArrayAttrBase<I64ArrayAttr, "Array of 64-bit integer array attributes">;
 
 class Linalg_TensorReshapeOp<string mnemonic> : Linalg_ReshapeLikeOp<
     mnemonic,
-    [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-      ["reifyReturnTypeShapesPerResultDim"]>]>,
+    [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]>,
     Arguments<(ins AnyTensor:$src,
                    IndexListArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {
   let extraClassDeclaration = commonExtraClassDeclaration # [{
     RankedTensorType getSrcType() {
       return src().getType().cast<RankedTensorType>();
     }
     RankedTensorType getResultType() {

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.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, InferShapedTypeOpInterface])> {
+       LinalgStructuredInterface, ReifyRankedShapedTypeOpInterface])> {
   code structuredOpsBaseDecls = [{
     // Return whether the op accesses the iteration indices.
     bool hasIndexSemantics() {
       Operation *op = this->getOperation();
       if(op->getNumRegions() == 0 || op->getRegion(0).empty())
         return false;
       return !op->getRegion(0).front().getOps<IndexOp>().empty();
     }
 
-    LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
-        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
-      return cast<LinalgOp>(getOperation()).reifyReturnTypeShapesPerResultDim(b,
+    LogicalResult reifyResultShapes(OpBuilder &b,
+        ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
+      return cast<LinalgOp>(getOperation()).reifyResultShapes(b,
           reifiedReturnShapes);
     }
   }];
 }
 
 class LinalgStructured_Op<string mnemonic, list<OpTrait> props>
   : LinalgStructuredBase_Op<mnemonic,
        !listconcat(props, [

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h

 // Patterns
 //===----------------------------------------------------------------------===//
 
 /// Appends patterns for folding memref.subview ops into consumer load/store ops
 /// into `patterns`.
 void populateFoldSubViewOpPatterns(RewritePatternSet &patterns);
 
 /// Appends patterns that resolve `memref.dim` operations with values that are
+/// defined by operations that implement the
+/// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
+/// operands.
+void populateResolveRankedShapeTypeResultDimsPatterns(
+    RewritePatternSet &patterns);
+
+/// Appends patterns that resolve `memref.dim` operations with values that are
 /// defined by operations that implement the `InferShapedTypeOpInterface`, in
 /// terms of shapes of its input operands.
 void populateResolveShapedTypeResultDimsPatterns(RewritePatternSet &patterns);
 
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an operation pass to fold memref.subview ops into consumer
 /// load/store ops into `patterns`.
 std::unique_ptr<Pass> createFoldSubViewOpsPass();
 
 /// Creates an operation pass to resolve `memref.dim` operations with values
 /// that are defined by operations that implement the
-/// `InferShapedTypeOpInterface`, in terms of shapes of its input operands.
+/// `ReifyRankedShapeTypeShapeOpInterface`, in terms of shapes of its input
+/// operands.
+std::unique_ptr<Pass> createResolveRankedShapeTypeResultDimsPass();
+
+/// Creates an operation pass to resolve `memref.dim` operations with values
+/// that are defined by operations that implement the
+/// `InferShapedTypeOpInterface` or the `ReifyRankedShapeTypeShapeOpInterface`,
+/// in terms of shapes of its input operands.
 std::unique_ptr<Pass> createResolveShapedTypeResultDimsPass();
 
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 
 } // namespace memref
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES_H

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.td

let description = [{
     from/to the original memref.
   }];
   let constructor = "mlir::memref::createFoldSubViewOpsPass()";
   let dependentDialects = [
       "AffineDialect", "memref::MemRefDialect", "vector::VectorDialect"
   ];
 }
 
+def ResolveRankedShapeTypeResultDims :
+    Pass<"resolve-ranked-shaped-type-result-dims"> {
+  let summary = "Resolve memref.dim of result values of ranked shape type";
+  let description = [{
+    The pass resolves memref.dim of result of operations that
+    implement the `ReifyRankedShapedTypeOpInterface` in terms of
+    shapes of its operands.
+  }];
+  let constructor =
+      "mlir::memref::createResolveRankedShapeTypeResultDimsPass()";
+  let dependentDialects = [
+    "memref::MemRefDialect", "tensor::TensorDialect"
+  ];
+}
+
 def ResolveShapedTypeResultDims : Pass<"resolve-shaped-type-result-dims"> {
   let summary = "Resolve memref.dim of result values";
   let description = [{
     The pass resolves memref.dim of result of operations that
-    implement the `InferShapedTypeOpInterface` in terms of shapes of
-    its operands.
+    implement the `InferShapedTypeOpInterface` or
+    `ReifyRankedShapedTypeOpInterface` in terms of shapes of its
+    operands.
   }];
   let constructor = "mlir::memref::createResolveShapedTypeResultDimsPass()";
   let dependentDialects = [
     "memref::MemRefDialect", "tensor::TensorDialect"
   ];
 }
 
 #endif // MLIR_DIALECT_MEMREF_TRANSFORMS_PASSES
 

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

 
 //===----------------------------------------------------------------------===//
 // InsertSliceOp
 //===----------------------------------------------------------------------===//
 
 def Tensor_InsertSliceOp : BaseOpWithOffsetSizesAndStrides<
     Tensor_Dialect, "insert_slice",
     [NoSideEffect, AttrSizedOperandSegments, OffsetSizeAndStrideOpInterface,
-     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-         ["reifyReturnTypeShapesPerResultDim"]>,
+     DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>,
      TypesMatchWith<"expected result type to match dest type",
                     "dest", "result", "$_self">]> {
   let summary = "insert_slice operation";
   let description = [{
     The "insert_slice" operation insert a tensor `source` into another
     tensor `dest` as specified by the operation's offsets, sizes and strides
     arguments.
 

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
+using ReifiedRankedShapedTypeDims = SmallVector<SmallVector<Value>>;
+
 /// ShapedTypeComponents that represents the components of a ShapedType.
 /// The components consist of
 ///  - A ranked or unranked shape with the dimension specification match those
 ///    of ShapeType's getShape() (e.g., dynamic dimension represented using
 ///    ShapedType::kDynamicSize)
 ///  - A element type, may be unset (nullptr)
 ///  - A attribute, may be unset (nullptr)
 /// Used by ShapedType type inferences.

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

StaticInterfaceMethod<
                       $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. Only one of this method or
-      `reifyReturnTypeShapesPerResultDim` needs to be overriden by the
-      operation. This interface is supposed to be workable during dialect
+      result shape. This interface is supposed to be workable during dialect
       conversion (e.g. convert from tensor world to buffer world),
       where `getOperand` may be invalid. For example, some ops (e.g.
       dynamic_reshape(input, target_shape)) may depend on their operands
       to calculate the result shape. When the `matchAndRewrite ` method
       of a conversion pattern is called, the operands of the op to convert
       may have been converted into other types, which makes it invalid to
       call the `getOperand` method of such op directly inside the
       conversion pattern.  To solve this problem, this interface follows
       the design of the conversion pattern, that is, accepting passed in
       operands to avoid calling `getOperand` directly inside the interface
       implementation.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
           "::mlir::ValueRange":$operands,
           "::mlir::SmallVectorImpl<::mlir::Value> &":$reifiedReturnShapes),
       /*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. 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. Consequently, this
-      interface can only be overridden if the return types are ranked.
-
-      If both this method and `reifyReturnTypeShapes` are overridden
-      by the operation, `reifyReturnTypeShapes` takes precedence. This
-      method is intended to be used when the shape of each result, dim
-      pair can be computed independently. Using this method avoids
-      adding additional instructions to aggregate individual dimension
-      of a result shape into an single `Value` (and consequently
-      avoids the need to extract the value from the shape on the
-      client side).
-      }],
-      /*retTy=*/"::mlir::LogicalResult",
-      /*methodName=*/"reifyReturnTypeShapesPerResultDim",
-      /*args=*/(ins "::mlir::OpBuilder&":$builder,
-          "::mlir::SmallVectorImpl<::mlir::SmallVector<::mlir::Value>>&"
-          :$reifiedReturnShapes),
-      /*methodBody=*/[{}],
-      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
     >
   ];
 }
 
 // Convenience class grouping together type and shaped type op interfaces for
 // ops that have tensor return types.
 class InferTensorType<list<string> overridenMethods = []> {
   list<OpTrait> traits = [
     // Op implements infer type op interface.
     InferTypeOpInterface,
     // The op will have methods implementing the ShapedType type inference
     // interface.
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
     // The op produces tensors and will use the ShapedType type infer interface
     // along with knowledge that it is producing Tensors to infer the type.
     NativeOpTrait<"InferTensorType">
   ];
 }
 defvar InferTensorTypeWithReify = InferTensorType<[
     "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
+
+def ReifyRankedShapedTypeOpInterface :
+    OpInterface<"ReifyRankedShapedTypeOpInterface"> {
+  let description = [{
+    Interface to compute the shape of the result of an operation when
+    the result is a ranked shape type, i.e. `RankedTensorType` or
+    `MemRefType`.
+  }];
+  let cppNamespace = "::mlir";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/[{
+        Reify the shape of the result of an operation (typically in
+        terms of shape of its operands)
+
+        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. Each of these vectors is expected to be of size equal to
+        rank of the corresponding result. If the shape of a particular
+        result cannot be computed it must be empty.
+      }],
+      /*retTy=*/"LogicalResult",
+      /*methodName=*/"reifyResultShapes",
+      /*args=*/(ins "::mlir::OpBuilder &":$builder,
+        "ReifiedRankedShapedTypeDims &":$reifiedReturnShapes)
+    >
+  ];
+}
+
 #endif // MLIR_INFERTYPEOPINTERFACE

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;
 };
 
-LogicalResult LinalgOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult
+LinalgOp::reifyResultShapes(OpBuilder &b,
+                            ReifiedRankedShapedTypeDims &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)
   //   subMapOfResultShapes = (d0, d1, d2) -> (d0 + d1, d1)
   //   shapesToLoopsMap = (d0, d2, d2, d3, d4, d5) -> (d0, d3, d2)
   //   resultShapesFromInputShapes = subMapOfResultDim.compose(shapesToLoopMap)

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

struct FoldInitTensorWithTensorReshapeOp
     : public OpRewritePattern<TensorReshapeOp> {
   using OpRewritePattern<TensorReshapeOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(TensorReshapeOp reshapeOp,
                                 PatternRewriter &rewriter) const override {
     if (!reshapeOp.src().template getDefiningOp<InitTensorOp>())
       return failure();
     Location loc = reshapeOp.getLoc();
-    SmallVector<SmallVector<Value>, 4> resultShapes;
-    if (failed(reshapeOp.reifyReturnTypeShapesPerResultDim(rewriter,
-                                                           resultShapes)) ||
+    ReifiedRankedShapedTypeDims resultShapes;
+    if (failed(reshapeOp.reifyResultShapes(rewriter, resultShapes)) ||
         !llvm::hasSingleElement(resultShapes))
       return failure();
     Value initTensor = rewriter.create<InitTensorOp>(
         loc, getAsOpFoldResult(resultShapes[0]),
         reshapeOp.getResultType().getElementType());
     if (initTensor.getType() != reshapeOp.getResultType()) {
       rewriter.replaceOpWithNewOp<tensor::CastOp>(
           reshapeOp, reshapeOp.getResultType(), initTensor);

 void InitTensorOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
   results.add<FoldInitTensorWithDimOp, FoldInitTensorWithExtractSliceOp,
               FoldInitTensorWithTensorReshapeOp<TensorExpandShapeOp>,
               FoldInitTensorWithTensorReshapeOp<TensorCollapseShapeOp>,
               ReplaceStaticShapeDims>(context);
 }
 
-LogicalResult InitTensorOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &builder,
-    SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult InitTensorOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   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>(getLoc(), getStaticSize(dim));
       }));
   reifiedReturnShapes.emplace_back(std::move(shapes));
   return success();
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);
 }
 
-LogicalResult PadTensorOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult PadTensorOp::reifyResultShapes(
+    OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   Location loc = getLoc();
   auto lowPad = getMixedLowPad();
   auto highPad = getMixedHighPad();
   SmallVector<Value> shapes;
   for (auto dim : llvm::seq<int64_t>(0, getSourceType().getRank())) {
     // Shape along each dimension is source dim + low pad + high pad.
     SmallVector<Value> mapOperands;
     mapOperands.push_back(b.createOrFold<tensor::DimOp>(loc, source(), dim));
void TensorCollapseShapeOp::getCanonicalizationPatterns(
   results
       .add<CollapseReshapeOps<TensorCollapseShapeOp>,
            CollapseMixedReshapeOps<TensorCollapseShapeOp, TensorExpandShapeOp>,
            FoldFillWithTensorReshape<TensorCollapseShapeOp>,
            FoldInitTensorWithTensorReshapeOp<TensorCollapseShapeOp>,
            FoldReshapeWithConstant<TensorCollapseShapeOp>>(context);
 }
 
-LogicalResult TensorExpandShapeOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult TensorExpandShapeOp::reifyResultShapes(
+    OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   auto resultShape =
       getAsValues(b, getLoc(),
                   getReshapeOutputShapeFromInputShape(
                       b, getLoc(), src(), getResultType().getShape(),
                       getReassociationMaps()));
   reifiedReturnShapes.emplace_back(std::move(resultShape));
   return success();
 }
 
-LogicalResult TensorCollapseShapeOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &b, SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult TensorCollapseShapeOp::reifyResultShapes(
+    OpBuilder &b, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   auto resultShape =
       getAsValues(b, getLoc(),
                   getReshapeOutputShapeFromInputShape(
                       b, getLoc(), src(), getResultType().getShape(),
                       getReassociationMaps()));
   reifiedReturnShapes.emplace_back(std::move(resultShape));
   return success();
 }

mlir/lib/Dialect/MemRef/Transforms/ResolveShapedTypeResultDims.cpp

-//===- ResolveShapedTypeResultDims.cpp - Resolve memref.dim ops of result values
-//-------===//
+//===- ResolveShapedTypeResultDims.cpp - Resolve dim ops of result values -===//
 //
 // 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 pass resolves `memref.dim` operations of result values in terms of
 // shapes of their operands using the `InferShapedTypeOpInterface`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
-/// Helper method to get the `Value` that is the shape of the `resultIdx`-th
-/// result at dimension `dimIndex` from the `ShapedTypeOpInterface`.
-/// TODO(ravishankarm): This is better put as a interface utility method
-/// somewhere, but that would imply the interface will depend on the `tensor`
-/// dialect. Ideally maybe a utility method in the `tensor` dialect.
-static Value getResultDimFromShapeInterface(OpBuilder &builder, OpResult result,
-                                            int64_t dimIndex) {
-  unsigned resultNumber = result.getResultNumber();
-  auto shapedTypeOp = dyn_cast<InferShapedTypeOpInterface>(result.getOwner());
-  Location loc = result.getOwner()->getLoc();
+namespace {
+/// Fold dim of an operation that implements the InferShapedTypeOpInterface
+template <typename OpTy>
+struct DimOfShapedTypeOpInterface : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy dimOp,
+                                PatternRewriter &rewriter) const override {
+    OpResult dimValue = dimOp.source().template dyn_cast<OpResult>();
+    if (!dimValue)
+      return failure();
+    auto shapedTypeOp =
+        dyn_cast<InferShapedTypeOpInterface>(dimValue.getOwner());
     if (!shapedTypeOp)
-    return nullptr;
+      return failure();
+
+    Optional<int64_t> dimIndex = dimOp.getConstantIndex();
+    if (!dimIndex)
+      return failure();
 
-  // The interface exposes two methods, one that returns the shape of all the
-  // results as `Value` and other that returns the shape as a list of
-  // `SmallVector<Value>`. The former takes precedence over the latter. So first
-  // check if the op implements the first interface method or the second, and
-  // get the value to use appropriately.
     SmallVector<Value> reifiedResultShapes;
-  if (succeeded(shapedTypeOp.reifyReturnTypeShapes(
-          builder, result.getOwner()->getOperands(), reifiedResultShapes))) {
-    if (reifiedResultShapes.size() <= resultNumber)
-      return nullptr;
-    Value resultShape = reifiedResultShapes[resultNumber];
+    if (failed(shapedTypeOp.reifyReturnTypeShapes(
+            rewriter, shapedTypeOp->getOperands(), reifiedResultShapes)))
+      return failure();
+
+    if (reifiedResultShapes.size() != shapedTypeOp->getNumResults())
+      return failure();
+
+    Value resultShape = reifiedResultShapes[dimValue.getResultNumber()];
     auto resultShapeType = resultShape.getType().dyn_cast<RankedTensorType>();
     if (!resultShapeType || !resultShapeType.getElementType().isa<IndexType>())
-      return nullptr;
-    return builder.create<tensor::ExtractOp>(
-        loc, resultShape, builder.createOrFold<ConstantIndexOp>(loc, dimIndex));
-  }
-
-  SmallVector<SmallVector<Value>> reifiedResultShapesPerDim;
-  if (failed(shapedTypeOp.reifyReturnTypeShapesPerResultDim(
-          builder, reifiedResultShapesPerDim)))
-    return nullptr;
-  if (reifiedResultShapesPerDim.size() <= resultNumber ||
-      reifiedResultShapesPerDim[resultNumber].size() !=
-          static_cast<size_t>(result.getType().cast<ShapedType>().getRank()))
-    return nullptr;
-  OpFoldResult valueOrAttr = reifiedResultShapesPerDim[resultNumber][dimIndex];
-  if (auto attr = valueOrAttr.dyn_cast<Attribute>())
-    return builder.createOrFold<ConstantIndexOp>(
-        loc, attr.cast<IntegerAttr>().getInt());
-  return valueOrAttr.get<Value>();
+      return failure();
+
+    Location loc = dimOp->getLoc();
+    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(
+        dimOp, resultShape,
+        rewriter.createOrFold<ConstantIndexOp>(loc, *dimIndex));
+    return success();
   }
+};
 
-namespace {
 /// Fold dim of an operation that implements the InferShapedTypeOpInterface
 template <typename OpTy>
-struct DimOfShapedTypeOpInterface : public OpRewritePattern<OpTy> {
+struct DimOfReifyRankedShapedTypeOpInterface : public OpRewritePattern<OpTy> {
   using OpRewritePattern<OpTy>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(OpTy dimOp,
                                 PatternRewriter &rewriter) const override {
     OpResult dimValue = dimOp.source().template dyn_cast<OpResult>();
     if (!dimValue)
       return failure();
-    auto shapedTypeOp =
-        dyn_cast<InferShapedTypeOpInterface>(dimValue.getOwner());
-    if (!shapedTypeOp)
+    auto rankedShapeTypeOp =
+        dyn_cast<ReifyRankedShapedTypeOpInterface>(dimValue.getOwner());
+    if (!rankedShapeTypeOp)
       return failure();
 
     Optional<int64_t> dimIndex = dimOp.getConstantIndex();
     if (!dimIndex)
       return failure();
-    Value replacement =
-        getResultDimFromShapeInterface(rewriter, dimValue, *dimIndex);
-    if (!replacement)
+
+    SmallVector<SmallVector<Value>> reifiedResultShapes;
+    if (failed(
+            rankedShapeTypeOp.reifyResultShapes(rewriter, reifiedResultShapes)))
+      return failure();
+
+    if (reifiedResultShapes.size() != rankedShapeTypeOp->getNumResults())
       return failure();
-    rewriter.replaceOp(dimOp, replacement);
+
+    unsigned resultNumber = dimValue.getResultNumber();
+    auto sourceType = dimValue.getType().dyn_cast<RankedTensorType>();
+    if (reifiedResultShapes[resultNumber].size() !=
+        static_cast<size_t>(sourceType.getRank()))
+      return failure();
+
+    rewriter.replaceOp(dimOp, reifiedResultShapes[resultNumber][*dimIndex]);
     return success();
   }
 };
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Pass registration
 //===----------------------------------------------------------------------===//
 
 namespace {
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/MemRef/Transforms/Passes.h.inc"
 
+struct ResolveRankedShapeTypeResultDimsPass final
+    : public ResolveRankedShapeTypeResultDimsBase<
+          ResolveRankedShapeTypeResultDimsPass> {
+  void runOnOperation() override;
+};
+
 struct ResolveShapedTypeResultDimsPass final
     : public ResolveShapedTypeResultDimsBase<ResolveShapedTypeResultDimsPass> {
   void runOnOperation() override;
 };
+
 } // namespace
 
+void memref::populateResolveRankedShapeTypeResultDimsPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<DimOfReifyRankedShapedTypeOpInterface<memref::DimOp>,
+               DimOfReifyRankedShapedTypeOpInterface<tensor::DimOp>>(
+      patterns.getContext());
+}
+
 void memref::populateResolveShapedTypeResultDimsPatterns(
     RewritePatternSet &patterns) {
   // TODO: Move tensor::DimOp pattern to the Tensor dialect.
   patterns.add<DimOfShapedTypeOpInterface<memref::DimOp>,
                DimOfShapedTypeOpInterface<tensor::DimOp>>(
       patterns.getContext());
 }
 
+void ResolveRankedShapeTypeResultDimsPass::runOnOperation() {
+  RewritePatternSet patterns(&getContext());
+  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
+  if (failed(applyPatternsAndFoldGreedily(getOperation()->getRegions(),
+                                          std::move(patterns))))
+    return signalPassFailure();
+}
+
 void ResolveShapedTypeResultDimsPass::runOnOperation() {
   RewritePatternSet patterns(&getContext());
+  memref::populateResolveRankedShapeTypeResultDimsPatterns(patterns);
   memref::populateResolveShapedTypeResultDimsPatterns(patterns);
   if (failed(applyPatternsAndFoldGreedily(getOperation()->getRegions(),
                                           std::move(patterns))))
     return signalPassFailure();
 }
 
 std::unique_ptr<Pass> memref::createResolveShapedTypeResultDimsPass() {
   return std::make_unique<ResolveShapedTypeResultDimsPass>();
 }
+
+std::unique_ptr<Pass> memref::createResolveRankedShapeTypeResultDimsPass() {
+  return std::make_unique<ResolveRankedShapeTypeResultDimsPass>();
+}

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

 OpFoldResult InsertSliceOp::fold(ArrayRef<Attribute>) {
   if (getSourceType().hasStaticShape() && getType().hasStaticShape() &&
       getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   return OpFoldResult();
 }
 
-LogicalResult InsertSliceOp::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &builder,
-    SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+LogicalResult InsertSliceOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
   reifiedReturnShapes.resize(1, SmallVector<Value>(getType().getRank()));
   for (auto dim : llvm::seq<int64_t>(0, getType().getRank())) {
     reifiedReturnShapes[0][dim] =
         builder.createOrFold<tensor::DimOp>(getLoc(), dest(), dim);
   }
   return success();
 }
 

mlir/test/Interfaces/InferShapedTypeOpInterface/resolve-shaped-type-result-dims.mlir

 //  CHECK-SAME:   %[[ARG_1:[a-z0-9]*]]: tensor<?x5xf32>)
 //   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
 //   CHECK-DAG:   %[[C2:.+]] = constant 2 : index
 //   CHECK-DAG:   %[[C3:.+]] = constant 3 : index
 //   CHECK-DAG:   %[[C5:.+]] = constant 5 : index
 //   CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG_1]], %[[C0]]
 //   CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG_0]], %[[C2]]
 //       CHECK:   return %[[D0]], %[[C5]], %[[C2]], %[[C3]], %[[D1]]
-
-// -----
-
-func @result_shape_and_per_dim(%arg0 : tensor<2x3x?xf32>, %arg1 : tensor<?x5xf32>)
-    -> (index, index, index, index, index) {
-  %c0 = constant 0 : index
-  %c1 = constant 1 : index
-  %c2 = constant 2 : index
-  %0:2 = "test.op_with_result_shape_and_per_dim_interface"(%arg0, %arg1)
-      : (tensor<2x3x?xf32>, tensor<?x5xf32>) -> (tensor<?x5xf32>, tensor<2x3x?xf32>)
-  %1 = tensor.dim %0#0, %c0 : tensor<?x5xf32>
-  %2 = tensor.dim %0#0, %c1 : tensor<?x5xf32>
-  %3 = tensor.dim %0#1, %c0 : tensor<2x3x?xf32>
-  %4 = tensor.dim %0#1, %c1 : tensor<2x3x?xf32>
-  %5 = tensor.dim %0#1, %c2 : tensor<2x3x?xf32>
-  return %1, %2, %3, %4, %5 : index, index, index, index, index
-}
-// CHECK-LABEL: func @result_shape_and_per_dim(
-//  CHECK-SAME:   %[[ARG_0:[a-z0-9]*]]: tensor<2x3x?xf32>
-//  CHECK-SAME:   %[[ARG_1:[a-z0-9]*]]: tensor<?x5xf32>)
-//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
-//   CHECK-DAG:   %[[C2:.+]] = constant 2 : index
-//   CHECK-DAG:   %[[C3:.+]] = constant 3 : index
-//   CHECK-DAG:   %[[C5:.+]] = constant 5 : index
-//   CHECK-DAG:   %[[D0:.+]] = tensor.dim %[[ARG_1]], %[[C0]]
-//   CHECK-DAG:   %[[S0:.+]] = tensor.from_elements %[[D0]], %[[C5]]
-//   CHECK-DAG:   %[[D0_OUT:.+]] = tensor.extract %[[S0]][%[[C0]]]
-//   CHECK-DAG:   %[[D1:.+]] = tensor.dim %[[ARG_0]], %[[C2]]
-//   CHECK-DAG:   %[[S1:.+]] = tensor.from_elements %[[C2]], %[[C3]], %[[D1]]
-//   CHECK-DAG:   %[[D1_OUT:.+]] = tensor.extract %[[S1]][%[[C2]]]
-//       CHECK:   return %[[D0_OUT]], %[[C5]], %[[C2]], %[[C3]], %[[D1_OUT]]

mlir/test/lib/Dialect/Test/TestDialect.cpp

auto currShape = llvm::to_vector<4>(llvm::map_range(
           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
         }));
     shapes.push_back(builder.create<tensor::FromElementsOp>(
         getLoc(), builder.getIndexType(), currShape));
   }
   return success();
 }
 
-LogicalResult
-OpWithResultShapePerDimInterfaceOp ::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &builder,
-    llvm::SmallVectorImpl<llvm::SmallVector<Value>> &shapes) {
-  Location loc = getLoc();
-  shapes.reserve(getNumOperands());
-  for (Value operand : llvm::reverse(getOperands())) {
-    auto currShape = llvm::to_vector<4>(llvm::map_range(
-        llvm::seq<int64_t>(
-            0, operand.getType().cast<RankedTensorType>().getRank()),
-        [&](int64_t dim) -> Value {
-          return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
-        }));
-    shapes.emplace_back(std::move(currShape));
-  }
-  return success();
-}
-
-LogicalResult OpWithResultShapeAndPerDimInterfaceOp::reifyReturnTypeShapes(
-    OpBuilder &builder, ValueRange operands,
-    llvm::SmallVectorImpl<Value> &shapes) {
-  Location loc = getLoc();
-  shapes.reserve(operands.size());
-  for (Value operand : llvm::reverse(operands)) {
-    auto currShape = llvm::to_vector<4>(llvm::map_range(
-        llvm::seq<int64_t>(
-            0, operand.getType().cast<RankedTensorType>().getRank()),
-        [&](int64_t dim) -> Value {
-          return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
-        }));
-    shapes.push_back(builder.create<tensor::FromElementsOp>(
-        getLoc(), builder.getIndexType(), currShape));
-  }
-  return success();
-}
-
-LogicalResult
-OpWithResultShapeAndPerDimInterfaceOp ::reifyReturnTypeShapesPerResultDim(
-    OpBuilder &builder,
-    llvm::SmallVectorImpl<llvm::SmallVector<Value>> &shapes) {
+LogicalResult OpWithResultShapePerDimInterfaceOp::reifyResultShapes(
+    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
   Location loc = getLoc();
   shapes.reserve(getNumOperands());
   for (Value operand : llvm::reverse(getOperands())) {
     auto currShape = llvm::to_vector<4>(llvm::map_range(
         llvm::seq<int64_t>(
             0, operand.getType().cast<RankedTensorType>().getRank()),
         [&](int64_t dim) -> Value {
           return builder.createOrFold<tensor::DimOp>(loc, operand, dim);

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

def OpWithResultShapeInterfaceOp : TEST_Op<"op_with_result_shape_interface",
       [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
           ["reifyReturnTypeShapes"]>]> {
   let arguments = (ins AnyRankedTensor:$operand1, AnyRankedTensor:$operand2);
   let results = (outs AnyRankedTensor:$result1, AnyRankedTensor:$result2);
 }
 
 def OpWithResultShapePerDimInterfaceOp :
     TEST_Op<"op_with_result_shape_per_dim_interface",
-        [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-             ["reifyReturnTypeShapesPerResultDim"]>]> {
-  let arguments = (ins AnyRankedTensor:$operand1, AnyRankedTensor:$operand2);
-  let results = (outs AnyRankedTensor:$result1, AnyRankedTensor:$result2);
-}
-
-def OpWithResultShapeAndPerDimInterfaceOp :
-    TEST_Op<"op_with_result_shape_and_per_dim_interface",
-        [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-             ["reifyReturnTypeShapes", "reifyReturnTypeShapesPerResultDim"]>]> {
+        [DeclareOpInterfaceMethods<ReifyRankedShapedTypeOpInterface>]> {
   let arguments = (ins AnyRankedTensor:$operand1, AnyRankedTensor:$operand2);
   let results = (outs AnyRankedTensor:$result1, AnyRankedTensor:$result2);
 }
 
 def IsNotScalar : Constraint<CPred<"$0.getType().getRank() != 0">>;
 
 def UpdateAttr : Pat<(I32ElementsAttrOp $attr),
                      (I32ElementsAttrOp ConstantAttr<I32ElementsAttr, "0">),

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     name = "LinalgInterfaces",
     srcs = ["lib/Dialect/Linalg/IR/LinalgInterfaces.cpp"],
     hdrs = ["include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h"],
     includes = ["include"],
     deps = [
         ":Affine",
         ":DialectUtils",
         ":IR",
+        ":InferTypeOpInterface",
         ":LinalgInterfacesIncGen",
         ":LinalgStructuredOpsIncGen",
         ":MemRefDialect",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )