[mlir][Linalg] NFC - Factor out Linalg functionality for shape and loop bounds computation

[nicolasvasilache]

Summary

This revision refactors code used in various Linalg transformations and makes it a first class citizen to the LinalgStructureOpInterface. This is in preparation to allowing more advanced Linalg behavior but is otherwise NFC.

[nicolasvasilache]

Simplify hooks.

[bondhugula]

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

776 ↗

(On Diff #306842)

Typo

798 ↗

(On Diff #306842)

Rephrase "used to the".

873 ↗

(On Diff #306842)

It's not really "operand dimensions" but "operand dimension sizes".

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

13–14 ↗

(On Diff #306842)

Includes not in sorted order. You also don't need many of the includes below as a result.

18 ↗

(On Diff #306842)

You don't need this include.

19 ↗

(On Diff #306842)

You don't need this either.

20 ↗

(On Diff #306842)

Not needed either.

25 ↗

(On Diff #306842)

Not needed.

26 ↗

(On Diff #306842)

Not needed. Prune these includes overall.

38 ↗

(On Diff #306842)

Doc comment missing.

70 ↗

(On Diff #306842)

Use i = 0, e = ... form to avoid repeated evaluation.

76 ↗

(On Diff #306842)

if (!attr) and early return. Will reduce indent / more readable for the entire block below.

133–134 ↗

(On Diff #306842)

Use .isa<AffineConstantExpr>() / isa<AffineSymbolExpr>() instead of getKind() comparison.

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

335–336 ↗

(On Diff #306842)

Avoid auto here.

[nicolasvasilache]

Address

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

13–14 ↗

(On Diff #306842)

Seems commonplace to isolate the matching include when such a unique include exists.
See e.g. PDL.cpp

133–134 ↗

(On Diff #306842)

This is moved code that will be deleted in the next revision.
But sure, thanks.

[nicolasvasilache]

Rebase

[nicolasvasilache]

Rebase

[nicolasvasilache]

Fix includes

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let extraClassDeclaration = [{
     void setMemRef(Value value) { setOperand(0, value); }
     MemRefType getMemRefType() {
       return getMemRef().getType().cast<MemRefType>();
     }
 
     operand_range getIndices() { return {operand_begin() + 1, operand_end()}; }
   }];
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 
   let assemblyFormat = "$memref `[` $indices `]` attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // LogOp
 //===----------------------------------------------------------------------===//

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

 
 OpFoldResult LoadOp::fold(ArrayRef<Attribute> cstOperands) {
   /// load(memrefcast) -> load
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
+namespace {
+/// Fold a load on a tensor_to_memref operation into an extract_element on the
+/// corresponding tensor.
+struct LoadOfTensorToMemref : public OpRewritePattern<LoadOp> {
+  using OpRewritePattern<LoadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LoadOp load,
+                                PatternRewriter &rewriter) const override {
+    auto tensorToMemref = load.memref().getDefiningOp<TensorToMemrefOp>();
+    if (!tensorToMemref)
+      return failure();
+
+    rewriter.replaceOpWithNewOp<ExtractElementOp>(load, tensorToMemref.tensor(),
+                                                  load.indices());
+    return success();
+  }
+};
+} // end anonymous namespace.
+
+void LoadOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                         MLIRContext *context) {
+  results.insert<LoadOfTensorToMemref>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // MemRefCastOp
 //===----------------------------------------------------------------------===//
 
 Value MemRefCastOp::getViewSource() { return source(); }
 
 bool MemRefCastOp::areCastCompatible(Type a, Type b) {
   auto aT = a.dyn_cast<MemRefType>();

mlir/test/Dialect/Standard/canonicalize.mlir

 //       CHECK:   return %[[D]] : index
 func @dim_of_tensor_load(%arg0: memref<?xf32>) -> index {
   %c0 = constant 0 : index
   %0 = tensor_load %arg0 : memref<?xf32>
   %1 = dim %0, %c0 : tensor<?xf32>
   return %1 : index
 }
 
+// Test case: Folding of load(tensor_to_memref(%v, %idxs))
+//            -> extract_element(%v, %idx)
+// CHECK-LABEL: func @load_from_tensor_to_memref(
+//  CHECK-SAME:     %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index
+//  CHECK-SAME:     %[[TENSOR:[0-9a-z]+]]: tensor<?x?xf32>
+//       CHECK:   %[[RES:.*]] = extract_element %[[TENSOR]][%[[IDX0]], %[[IDX1]]]
+//   CHECK-NOT:   load
+//       CHECK:   return %[[RES]] : f32
+func @load_from_tensor_to_memref(%arg0: index, %arg1: index, %arg2: tensor<?x?xf32>) -> f32 {
+  %0 = tensor_to_memref %arg2 : memref<?x?xf32>
+  %1 = load %0[%arg0, %arg1] : memref<?x?xf32>
+  return %1 : f32
+}
+
 // Test case: Folding of dim(dynamic_tensor_from_elements %idx) -> %idx
 // CHECK-LABEL: func @dim_of_dynamic_tensor_from_elements(
 //  CHECK-SAME:     %[[IDX0:[0-9a-z]+]]: index, %[[IDX1:[0-9a-z]+]]: index
 //   CHECK-NOT:   dim
 //       CHECK:   return %[[IDX1]] : index
 func @dim_of_dynamic_tensor_from_elements(%arg0: index, %arg1: index) -> index {
   %c3 = constant 3 : index
   %0 = dynamic_tensor_from_elements %arg0, %arg1 {