[mlir][linalg] Derive padded subtensor shape from linalg op outputs

[springerm]

Summary

Add a canonicalization pattern that rewrites memref::DimOps accordingly.

Depends On D103226

[nicolasvasilache]

Let's add a proper folding hook and use createAndFold please, this has multiple benefits over all clients using internal op-specific knowledge.

[springerm]

rebase + rewrite as canonicalization pattern

[springerm]

In D103132#2781348, @nicolasvasilache wrote:

Let's add a proper folding hook and use createAndFold please, this has multiple benefits over all clients using internal op-specific knowledge.

The folder would have to be implemented as part of DimOp::fold in MemRefOps.cpp. It would have to check if the DimOp operand is a LinalgOp and if so, replace it with a new DimOp on the LinalgOp's respective output. This means putting linalg-specific code into the memref dialect.

A canonicalization pattern may be more suitable.

[nicolasvasilache]

A folder would be better because createOrFold cleans a lot of things up by construction without relying on calling a canonicalization pass in weird places.
In this particular this is fine for now because coarse-grained canonicalization is fine.

Still, I wonder if this is the justification we've been wanting to get the TiedOperandInterface in place and refactor using this: @silvas @mehdi_amini for an extra set of eyes.

[springerm]

minor changes

[springerm]

Submitting this commit as is, but something like TiedOperandInterface would in fact be useful here. Will revisit this if we end up adding sth like TiedOperandInterface.

[silvas]

DimOp should already have a canonicalizer that will use InferShapedTypeOpInterface to rewrite dim(linalgop(x)) to some_expr_of(dim(x)). Do you have a test case where that doesn't happen?

One difference is that I think the current way that this canonicalization plays out is that it calculates the output shape in terms of the ins. Whereas what you are proposing here calculates it in terms of the outs (which is trivial). This came up in the past and @mravishankar argued for the "shape in terms of ins" approach, whereas I argued for the "shape in terms of outs" approach. I'll let you experts discuss this matter.

I'm generally against TiedOperandInterface (at least the way that IREE's is defined) because it has very vague semantics (it conflates a number of things -- the definition literally says something like "depending on the context it is interpreted as either A, B, or C" which doesn't have the precision needed to guide transformations safely). If the information needed is that the shape of an output is the same as the shape of an input, then InferShapedTypeOpInterface is already sufficient to express that.

[springerm]

You are right, the existing canonicalizations + CSE are good enough. This commit is not needed, but I didn't realize that until after I wrote the pattern ;)

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

let extraClassDeclaration = [{
     using RegionBuilderFunType =
       llvm::function_ref<void(ImplicitLocOpBuilder &b, Block &, ValueRange)>;
     RegionBuilderFunType getRegionBuilder(StringRef name) {
       return namedStructuredOpRegionBuilders.lookup(name);
     }
     private:
       llvm::StringMap<RegionBuilderFunType> namedStructuredOpRegionBuilders;
   }];
+  let hasCanonicalizer = 1;
 }
 
 // Whether a type is a RangeType.
 def LinalgIsRangeTypePred : CPred<"$_self.isa<RangeType>()">;
 def Range : DialectType<Linalg_Dialect, LinalgIsRangeTypePred, "range">;
 
 #endif // LINALG_BASE

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

OpBuilder getBuilder() {
     builder.setInsertionPointToEnd(&block);
     return builder;
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
+// LinalgOp
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Rewrite a memref::DimOp on a LinalgOp result to a memref::DimOp on the
+/// LinalgOp's respective output. E.g.:
+/// ```
+/// %0 = linalg.matmul ins(%a, %b : tensor<?x?xf32>, tensor<?x?xf32>)
+///                    outs(%c : tensor<?x?xf32>) -> tensor<?x?xf32>
+/// %r = memref.dim %0, %c0 : tensor<?x?xf32>
+/// ```
+/// is rewritten to:
+/// ```
+/// %0 = linalg.matmul ins(%a, %b : tensor<?x?xf32>, tensor<?x?xf32>)
+///                    outs(%c : tensor<?x?xf32>) -> tensor<?x?xf32>
+/// %r = memref.dim %c, %c0 : tensor<?x?xf32>
+/// ```
+struct DimOfLinalgResultCanonicalization : OpRewritePattern<memref::DimOp> {
+  using OpRewritePattern<memref::DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::DimOp dimOp,
+                                PatternRewriter &rewriter) const override {
+    auto linalgOp = dimOp.memrefOrTensor().getDefiningOp<LinalgOp>();
+    if (!linalgOp)
+      return failure();
+
+    // TODO(gysit): Use LinalgOp interface function that returns an OpOperand&
+    // corresponding to the given OpResult.
+    auto idx = dimOp.memrefOrTensor().cast<OpResult>().getResultNumber();
+    rewriter.replaceOpWithNewOp<memref::DimOp>(dimOp, linalgOp.getOutput(idx),
+                                               dimOp.index());
+    return success();
+  }
+};
+} // namespace
+
+void LinalgDialect::getCanonicalizationPatterns(
+    RewritePatternSet &results) const {
+  results.add<DimOfLinalgResultCanonicalization>(this->getContext());
+}
+
+//===----------------------------------------------------------------------===//
 // CopyOp
 //===----------------------------------------------------------------------===//
 void CopyOp::regionBuilder(ImplicitLocOpBuilder &b, Block &block,
                            ValueRange captures) {
   assert(block.getNumArguments() == 2 && "CopyOp regionBuilder expects 2 args");
   b.create<linalg::YieldOp>(block.getArgument(0));
 }
 

mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp

auto resultTensorTypes =
       ValueRange(newOperands).take_back(opToPad.getNumOutputs()).getTypes();
   ValueRange otherOperands = opToPad.getAssumedNonShapedOperands();
   newOperands.append(otherOperands.begin(), otherOperands.end());
   linalg::LinalgOp paddedOp =
       opToPad.clone(rewriter, loc, resultTensorTypes, newOperands);
 
   // Recover the subtensor out of the new static results. This keeps the
   // original linalg op around because it uses the dims of the original results.
-  // This later folds away.
+  // This is canonicalized away by DimOfLinalgResultCanonicalization.
   SmallVector<Value> paddedSubviewResults;
   paddedSubviewResults.reserve(opToPad->getNumResults());
   SetVector<Operation *> newUsersOfOpToPad;
   for (auto it : llvm::zip(opToPad->getResults(), paddedOp->getResults())) {
     auto rank = std::get<0>(it).getType().cast<RankedTensorType>().getRank();
     SmallVector<OpFoldResult> offsets(rank, rewriter.getIndexAttr(0));
     auto sizes = llvm::to_vector<4>(llvm::map_range(
         llvm::seq<unsigned>(0, rank), [&](unsigned d) -> OpFoldResult {