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

A new `InterfaceMethod` is added to `InferShapedTypeOpInterface` that
allows an operation to return the `Value`s for each dim of its
results. It is intended for the case where the `Value` returned for
each dim is computed using the operands and operation attributes. This
interface method is for cases where the result dim of an operation can
be computed independently, and it avoids the need to aggregate all
dims of a result into a single shape value. This also implies that
this is not suitable for cases where the result type is unranked (for
which the existing interface methods is to be used).

Also added is a canonicalization pattern that uses this interface and
resolves the shapes of the output in terms of the shapes of the
inputs. Moving Linalg ops to use this interface, so that many
canonicalization patterns implemented for individual linalg ops to
achieve the same result can be removed in favor of the added
canonicalization pattern.

Differential Revision: https://reviews.llvm.org/D97887

Author: MaheshRavishankar

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

@@ -1087,18 +1087,20 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
     >,
     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};
       }]
     >,
     InterfaceMethod<
@@ -1226,8 +1228,8 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> {
 
     /// 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 reifyReturnTypeShapesPerResultDim(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes);
 
     //========================================================================//
     // Helper functions to mutate the `operand_segment_sizes` attribute.

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 @@ SmallVector<AffineExpr, 4> concat(ArrayRef<AffineExpr> a,
 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

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 @@ class Linalg_Op<string mnemonic, list<OpTrait> traits = []> :
   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,
+         ["reifyReturnTypeShapesPerResultDim"]>]> {
   let summary = "operation to define a tensor of particular value";
 
   let description = [{
@@ -126,7 +130,10 @@ def Linalg_InitTensorOp : Linalg_Op<"init_tensor", [NoSideEffect]> {
 }
 
 def Linalg_PadTensorOp : Linalg_Op<"pad_tensor",
-    [AttrSizedOperandSegments, NoSideEffect]> {
+    [AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+         ["reifyReturnTypeShapesPerResultDim"]>,
+     NoSideEffect]> {
   let summary = "tensor pad operation";
   let description = [{
     `linalg.pad_tensor` is an operation that pads the `source` tensor
@@ -348,11 +355,6 @@ class Linalg_ReshapeLikeOp<string mnemonic, list<OpTrait> traits = []> :
               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 @@ def Linalg_ReshapeOp : Linalg_ReshapeLikeOp<"reshape",
   let hasCanonicalizer = 1;
 }
 
-def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<"tensor_reshape">,
+def Linalg_TensorReshapeOp : Linalg_ReshapeLikeOp<
+    "tensor_reshape",
+    [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+      ["reifyReturnTypeShapesPerResultDim"]>]>,
     Arguments<(ins AnyTensor:$src,
                    AffineMapArrayAttr:$reassociation)>,
     Results<(outs AnyTensor:$result)> {

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,14 +26,20 @@ include "mlir/Interfaces/SideEffectInterfaces.td"
 // 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;
     }
+
+    LogicalResult reifyReturnTypeShapesPerResultDim(OpBuilder &b,
+        SmallVectorImpl<SmallVector<Value>> &reifiedReturnShapes) {
+      return cast<LinalgOp>(getOperation()).reifyReturnTypeShapesPerResultDim(b,
+          reifiedReturnShapes);
+    }
   }];
 }
 

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

@@ -97,21 +97,53 @@ def InferShapedTypeOpInterface : OpInterface<"InferShapedTypeOpInterface"> {
                     "::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. Only one of this method or
+      `reifyReturnTypeShapesPerResultDim` needs to be overriden by the
+      operation.
       }],
       /*retTy=*/"::mlir::LogicalResult",
       /*methodName=*/"reifyReturnTypeShapes",
       /*args=*/(ins "::mlir::OpBuilder&":$builder,
-                    "::mlir::SmallVectorImpl<::mlir::Value>&":$reifiedReturnShapes),
+          "::mlir::SmallVectorImpl<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<SmallVector<::mlir::Value>>&"
+          :$reifiedReturnShapes),
+      /*methodBody=*/[{}],
+      /*defaultImplementation=*/[{ return ::mlir::failure(); }]
+    >
   ];
 }
 
@@ -129,6 +161,7 @@ class InferTensorType<list<string> overridenMethods = []> {
     NativeOpTrait<"InferTensorType">
   ];
 }
-defvar InferTensorTypeWithReify = InferTensorType<["reifyReturnTypeShapes"]>;
+defvar InferTensorTypeWithReify = InferTensorType<[
+    "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
 #endif // MLIR_INFERTYPEOPINTERFACE

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

@@ -14,6 +14,7 @@ add_mlir_dialect_library(MLIRLinalg
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRDialectUtils
+  MLIRInferTypeOpInterface
   MLIRIR
   MLIRParser
   MLIRSideEffectInterfaces

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

@@ -188,7 +188,7 @@ SmallVector<Value, 4> LinalgOp::createFlatListOfOperandDims(OpBuilder &b,
   for (Value v : getShapedOperands()) {
     ShapedType t = v.getType().template cast<ShapedType>();
     for (unsigned i = 0, e = t.getRank(); i < e; ++i)
-      res.push_back(b.create<memref::DimOp>(loc, v, i));
+      res.push_back(b.createOrFold<memref::DimOp>(loc, v, i));
   }
   return res;
 }
@@ -234,57 +234,58 @@ struct HasAffineDimExprVisitor
   llvm::SmallSet<unsigned, 4> positions;
 };
 
-Optional<Value> LinalgOp::inferResultDimFromInputShapes(OpBuilder &b,
-                                                        Location loc,
-                                                        unsigned resultIdx,
-                                                        unsigned dim) {
+LogicalResult LinalgOp::reifyReturnTypeShapesPerResultDim(
+    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.createOrFold<memref::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) {