[mlir][Linalg] Add canonicalization of linalg op -> dim op.

[mravishankar]

Summary

Add canonicalization to replace use of the result of a linalg
operation on tensors in a dim operation, to use one of the operands of
the linalg operations instead. This allows the linalg op itself to be
deleted when all its non-dim uses are removed (say through tiling, etc.)

[nicolasvasilache]

Also @pifon2a

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

1734

prevent

1748

This revision introduces OpInterface support for getLoopsToShapesMap and getShapesToLoopsMap and dropped special-case logic.

Let's please use those instead of reintroducing special-case logic, this will be more future-proof.

This should prob take the form of a new interface method Value inferDimFromInputs(unsigned resultIdx, unsigned dim) which:

1. computes the index of result(resultIdx)[dim] in the flattened list of shapes (i.e. sum the ranks for all input operands and results < resultIdx + dim). let's call this "IDX"
2. extract the map consisting of the single IDX result from getLoopsToShapesMap
3. compose 2. with getShapesToLoopsMap that is the expression of the output dim you want in terms of the expression of the input dims.
4. call applyMapToValues with 3. and the flattened list of dims.

e.g. assume some hypothetical op that resembles matmul with a twist: O(i + j, j) += A(i, k) * B(k, j):

• getLoopsToShapesMap is (i, j, k) -> (i, k, k, j, i + j, j)
• getShapesToLoopsMap is (d0, d1, d2, d3, d4, d5) -> (d0, d3, d1)

Assume you wanted the expression of dim res, 0 (i.e. i + j, a.k.a dim A, 0 + dim B, 1).

Step 1. IDX = 4
Step 2. extract IDX from getLoopsToShapesMap, i.e. (i, j, k) -> (i + j)
Step 3. compute ([ (i, j, k) -> (0, 0, 0, 0, i + j, 0) ]).compose( [(d0, d1, d2, d3, d4, d5) -> (d0, d3, d1)] ) returns [d0, d1, d2, d3, d4, d5) -> (d0 + d1)]
Step 4. call applyMapToValues(b, loc, [d0, d1, d2, d3, d4, d5) -> (d0 + d1)], op.createFlatListOfOperandDims())

In step 4. you will prob need to adjust to only take the input operands.
By default we also capture the output because of output buffers.

Does this make sense ?

[hanchung]

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

1782–1783

nit: use auto because there is a dyn_cast.

https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

[mravishankar]

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

1748

I think it makes sense. But there is a corner case here that might make the pattern rewriter go into a loop, and it doesnt happen today. Take an op where the loop shape actually depends on the shape of the output.

%cst = constant 0.0 : f32
%1 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],
                                 iterator_types = ["parallel", "parallel"]}
         ins() {
         ^bb0():
             linalg.yield %cst : f32
        } -> tensor<?x?xf32>

Today, the loop computation fails to get the right loop shape for such ops (the method createLoopRanges I think will throw an error cause createFlatListOfOperandDims does not append the shape of the results that makes affineApplyMapToValues(b, loc, getShapesToLoopsMap(), createFlatListOfOperandDims() fail due to mismatch in number of values passed and numdims of the map. (The next CL in the stack fixes that issue).
So here we need to guard against that. I like the approach you described above, but this canonicalization should fail in such cases otherwise it will put the rewriter in a loop. It doesnt happen today cause of what I described in the previous paragraph)

[mravishankar]

Address comments.

[hanchung]

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

89

IIRC, Value can be nullptr, so we don't need optional here?

114–118

Let's initialize resultDimSubMapPos with dim, so we don't need the last statement ...+=dim?

[mravishankar]

Rebase

[mravishankar]

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

89

Thats true, but the signature of Optional makes it explicit that this can return llvm::None. Its easier to follow that. We can dcoument that it returns nullptr on failure, but that might just be easier to miss. This makes it more explicit. WDYT?

114–118

I feel like that would make the logic more complicated. This is easier to follow IMO. Will leave it as is if thats OK.

[hanchung]

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

89

Yes, you are right on that. It is indeed more explicit. Personally, I don't have a strong feeling on this because I know Value could be nullptr and it could ends up with an assertion at some point. It already has Optional semantic to me.

Another point is that Optional introduces a bit more memory usage (ie sizeof(bool) and padding may be added). If Value itself carries this information and people commonly use it, I'd prefer returning nullptr on failure.

However, I went through some usage, and found that some people use Optional<Value> and some people use Value directly. There is no convention here. So I'll leave the decision to you.

114–118

Sure, both are okay to me. I just feel that the comment states what it's doing and it looks simpler to me. But both looks good to me.

[mravishankar]

Adapting this to work with current setup of linalg operations. The
specifics might change based on
https://llvm.discourse.group/t/linalg-and-shapes/2421.

[mravishankar]

Rebase

[mravishankar]

Rebase

[nicolasvasilache]

Going through my unreviewed stack, note that with the revamp of linalg on tensors @pifon2a sent a diff today that implements the functionality with out tensors,

[mravishankar]

Rebase

[mravishankar]

In D93076#2480253, @nicolasvasilache wrote:

Going through my unreviewed stack, note that with the revamp of linalg on tensors @pifon2a sent a diff today that implements the functionality with out tensors,

Ok, Ill leave it upto you if this needs to be pushed in. I rebased it on top of the linalg on tensors changes. It still describes the shape in terms of "ins". THat might still be worth it. PTAL and let me know.

[nicolasvasilache]

ok, my apologies, this looks great, I'll do a proper review in a bit but consider it landed :)

[nicolasvasilache]

Thanks!

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

127

Please give this a good name and add to the LinalgStructuredOpInterface.

128

inferred

132

Just iterate of getInputShapes ?
If it's not in the LinalgStructuredOpsInterface just add it too, plz

[mravishankar]

Address comments

[mravishankar]

Rebase

[mravishankar]

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

132

Not needed anymore after interface method is added.

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

if (auto d = result.dyn_cast<AffineDimExpr>()) {
       if (res[d.getPosition()].offset)
         continue;
       res[d.getPosition()] = Range{zeroVal, viewSizes[idx], oneVal};
     }
   }
   return res;
 }
 
 /// Forward declarations.
 template <typename NamedStructuredOpType>
 static void buildNamedStructuredOpRegionAndAttributes(
     OpBuilder &opBuilder, OperationState &result, TypeRange inputTypes,
     TypeRange outputBufferTypes, TypeRange initTensorTypes,
     TypeRange resultTypes);
 
 static ParseResult
 parseCommonStructuredOpParts(OpAsmParser &parser, OperationState &result,
                              SmallVectorImpl<Type> &inputTypes,
                              SmallVectorImpl<Type> &outputBufferTypes,
                              SmallVectorImpl<Type> &initTensorTypes);
 
 template <typename NamedStructuredOpType>
 static ParseResult
 parseNamedStructuredOpRegion(OpAsmParser &parser, Region &region,
                              TypeRange inputTypes, TypeRange outputBufferTypes,
                              TypeRange initTensorTypes, TypeRange resultTypes);
 static ParseResult
 parseNamedStructuredOpResults(OpAsmParser &parser,
                               SmallVectorImpl<Type> &resultTypes);
 
 template <typename NamedStructuredOpType>
 static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
                                           OperationState &result);
 
 template <typename NamedStructuredOpType>
 static void printCommonStructuredOpParts(OpAsmPrinter &p,
                                          NamedStructuredOpType op);
 
 static void printNamedStructuredOpResults(OpAsmPrinter &p,
                                           TypeRange resultTypes);
 
 template <typename NamedStructuredOpType>
 static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op);
 
 template <typename NamedStructuredOpType>
 static LogicalResult verifyNamedStructuredOp(NamedStructuredOpType op);
 
 /// This is a common class used for patterns of the form
 /// ```
 ///    someop(memrefcast) -> someop
 /// ```
 /// It folds the source of the memref_cast into the root operation directly.
 static LogicalResult foldMemRefCast(Operation *op) {
   bool folded = false;
   for (OpOperand &operand : op->getOpOperands()) {
     auto castOp = operand.get().getDefiningOp<MemRefCastOp>();
     if (castOp && canFoldIntoConsumerOp(castOp)) {
       operand.set(castOp.getOperand());
       folded = true;
     }
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 prevert 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();
+    LinalgOp linalgOp = dimValue.getDefiningOp<LinalgOp>();
+    if (!linalgOp)
+      return failure();
+
+    // Get the indexing map of the current result value which is the operand of
+    // the dim op.
+    auto indexingMaps = linalgOp.getIndexingMaps();
+    unsigned resultIndex = dimValue.cast<OpResult>().getResultNumber();
+    AffineMap resultIndexMap =
+        indexingMaps[resultIndex + linalgOp.getNumInputs()];
+    AffineDimExpr resultShapeDimExpr =
+        resultIndexMap.getResult(*dimIndex).dyn_cast<AffineDimExpr>();
+    if (!resultShapeDimExpr)
+      return failure();
+    unsigned resultDimPosition = resultShapeDimExpr.getPosition();
+    // If there is another indexing map from the input that uses the same
+    // AffineDimExpr as the output, then we can use that instead. For example,
+    // in `linalg.matmul`
+    //
+    // LHS indexing map : affine_map<(d0, d1, d2) -> (d0, d2)>
+    // Output indexing map : affine_map<(d0, d1, d2) -> (d0, d1)>
+    //
+    // so
+    //
+    //  dim %{{matmul-output}}, %c0
+    //
+    // is same as
+    //
+    //  dim %{{matmul-LHS}}, %c0
+    //
+    // since
+    //   {{LHS indexing map}}.getResult(0) == {{RHS indexing map}}.getResult(0)
+    for (unsigned inputIndex :
+         llvm::seq<unsigned>(0, linalgOp.getNumInputs())) {
+      unsigned operandIndex =
+          linalgOp.getOperandIndexForInputIndex(inputIndex).getValue();
+      for (auto inputMapResultExpr :
+           enumerate(indexingMaps[operandIndex].getResults())) {
+        AffineDimExpr inputMapResultDimExpr =
+            inputMapResultExpr.value().dyn_cast<AffineDimExpr>();
+        if (!inputMapResultDimExpr)
+          continue;
+        unsigned inputMapResultDimPosition =
+            inputMapResultDimExpr.getPosition();
+        if (inputMapResultDimPosition != resultDimPosition)
+          continue;
+        rewriter.replaceOpWithNewOp<DimOp>(dimOp, linalgOp.getInput(inputIndex),
+                                           inputMapResultExpr.index());
+        return success();
+      }
+    }
+    return failure();
+  }
+};
+
 } // 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<EraseDeadLinalgOp>();                                       \
     results.insert<FoldTensorCastOp>();                                        \
     results.insert<DeduplicateInputs>();                                       \
+    results.insert<ReplaceDimOfLinalgOpResult>(context);                       \
   }                                                                            \
                                                                                \
   LogicalResult XXX::fold(ArrayRef<Attribute>,                                 \
                           SmallVectorImpl<OpFoldResult> &) {                   \
     return foldMemRefCast(*this);                                              \
   }
 
 CANONICALIZERS_AND_FOLDERS(ConvOp)

mlir/test/Dialect/Linalg/canonicalize.mlir

func @linalg_effects(%a : tensor<?x?xf32>, %b : memref<?x?xf32>, %c : tensor<?x?xf32>) {
   %t = linalg.matmul ins(%a, %b : tensor<?x?xf32>, memref<?x?xf32>)
                     init(%c : tensor<?x?xf32>) -> tensor<?x?xf32>
 
   // CHECK-NOT:   %{{.*}} = linalg.matmul
   linalg.matmul ins(%a, %c : tensor<?x?xf32>, tensor<?x?xf32>)
                outs(%b : memref<?x?xf32>)
   return
 }
+
+// -----
+
+func @remove_dim_result_uses
+  (%arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
+   %arg2 : tensor<?x?xf32>) -> (index, index, index, index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = linalg.matmul ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+    init(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = dim %0, %c0 : tensor<?x?xf32>
+  %2 = dim %0, %c1 : tensor<?x?xf32>
+  %3 = linalg.generic
+    {indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0)>,
+                      affine_map<(d0, d1, d2) -> (d0, d2)>,
+          affine_map<(d0, d1, d2) -> (d0, d2)>],
+     iterator_types = ["parallel", "reduction", "parallel"]}
+    ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+    init(%arg2 : tensor<?x?xf32>) {
+    ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32):
+      %4 = mulf %arg3, %arg4 : f32
+      %5 = addf %4, %arg5 : f32
+      linalg.yield %5 : f32
+    } -> tensor<?x?xf32>
+  %6 = dim %3, %c0 : tensor<?x?xf32>
+  %7 = dim %3, %c1 : tensor<?x?xf32>
+  return %1, %2, %6, %7 : index, index, index, index
+}
+// CHECK-LABEL: func @remove_dim_result_uses
+//  CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//  CHECK-SAME:   %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>
+//   CHECK-DAG:   %[[C0:.+]] = constant 0 : index
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//   CHECK-DAG:   %[[T0:.+]] = dim %[[ARG0]], %[[C0]]
+//   CHECK-DAG:   %[[T1:.+]] = dim %[[ARG1]], %[[C1]]
+//   CHECK-DAG:   %[[T2:.+]] = dim %[[ARG0]], %[[C1]]
+//   CHECK-DAG:   %[[T3:.+]] = dim %[[ARG1]], %[[C1]]
+//       CHECK:   return %[[T0]], %[[T1]], %[[T2]], %[[T3]]
 No newline at end of file